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

package http

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"

	"github.com/cloudprober/cloudprober/targets/endpoint"
)

const relURLLabel = "relative_url"

// requestBody encapsulates the request body and implements the io.Reader()
// interface.
type requestBody struct {
	b []byte
}

// Read implements the io.Reader interface. Instead of using buffered read,
// it simply copies the bytes to the provided slice in one go (depending on
// the input slice capacity) and returns io.EOF. Buffered reads require
// resetting the buffer before re-use, restricting our ability to use the
// request object concurrently.
func (rb *requestBody) Read(p []byte) (int, error) {
	return copy(p, rb.b), io.EOF
}

// resolveFunc resolves the given host for the IP version.
// This type is mainly used for testing. For all other cases, a nil function
// should be passed to the httpRequestForTarget function.
type resolveFunc func(host string, ipVer int) (net.IP, error)

func hostWithPort(host string, port int) string {
	if port == 0 {
		return host
	}
	return fmt.Sprintf("%s:%d", host, port)
}

// hostHeaderForTarget computes request's Host header for a target.
//  - If host header is set in the probe, it overrides everything else.
//  - If target's fqdn is provided in its labels, use that along with the port.
//  - Finally, use target's name with port.
func hostHeaderForTarget(target endpoint.Endpoint, probeHostHeader string, port int) string {
	if probeHostHeader != "" {
		return probeHostHeader
	}

	if target.Labels["fqdn"] != "" {
		return hostWithPort(target.Labels["fqdn"], port)
	}

	return hostWithPort(target.Name, port)
}

func urlHostForTarget(target endpoint.Endpoint) string {
	if target.Labels["fqdn"] != "" {
		return target.Labels["fqdn"]
	}

	return target.Name
}

func relURLForTarget(target endpoint.Endpoint, probeURL string) string {
	if probeURL != "" {
		return probeURL
	}

	if target.Labels[relURLLabel] != "" {
		return target.Labels[relURLLabel]
	}

	return ""
}

func (p *Probe) httpRequestForTarget(target endpoint.Endpoint, resolveF resolveFunc) *http.Request {
	// Prepare HTTP.Request for Client.Do
	port := int(p.c.GetPort())
	// If port is not configured explicitly, use target's port if available.
	if port == 0 {
		port = target.Port
	}

	urlHost := urlHostForTarget(target)
	ipForLabel := ""

	if p.c.GetResolveFirst() {
		if resolveF == nil {
			resolveF = p.opts.Targets.Resolve
		}

		ip, err := resolveF(target.Name, p.opts.IPVersion)
		if err != nil {
			p.l.Error("target: ", target.Name, ", resolve error: ", err.Error())
			return nil
		}

		ipStr := ip.String()
		urlHost, ipForLabel = ipStr, ipStr
	}

	for _, al := range p.opts.AdditionalLabels {
		al.UpdateForTargetWithIPPort(target, ipForLabel, port)
	}

	// Put square brackets around literal IPv6 hosts. This is the same logic as
	// net.JoinHostPort, but we cannot use net.JoinHostPort as it works only for
	// non default ports.
	if strings.IndexByte(urlHost, ':') >= 0 {
		urlHost = "[" + urlHost + "]"
	}

	url := fmt.Sprintf("%s://%s%s", p.protocol, hostWithPort(urlHost, port), relURLForTarget(target, p.url))

	// Prepare request body
	var body io.Reader
	if len(p.requestBody) > 0 {
		body = &requestBody{p.requestBody}
	}
	req, err := http.NewRequest(p.method, url, body)
	if err != nil {
		p.l.Error("target: ", target.Name, ", error creating HTTP request: ", err.Error())
		return nil
	}

	var probeHostHeader string
	for _, header := range p.c.GetHeaders() {
		if header.GetName() == "Host" {
			probeHostHeader = header.GetValue()
			continue
		}
		req.Header.Set(header.GetName(), header.GetValue())
	}

	// Host header is set by http.NewRequest based on the URL, update it based
	// on various conditions.
	req.Host = hostHeaderForTarget(target, probeHostHeader, port)

	return req
}
<commit_msg>Set Content-Length header on outgoing requests<commit_after>// Copyright 2019-2020 The Cloudprober Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package http

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"

	"github.com/cloudprober/cloudprober/targets/endpoint"
)

const relURLLabel = "relative_url"

// requestBody encapsulates the request body and implements the io.Reader()
// interface.
type requestBody struct {
	b []byte
}

// Read implements the io.Reader interface. Instead of using buffered read,
// it simply copies the bytes to the provided slice in one go (depending on
// the input slice capacity) and returns io.EOF. Buffered reads require
// resetting the buffer before re-use, restricting our ability to use the
// request object concurrently.
func (rb *requestBody) Read(p []byte) (int, error) {
	return copy(p, rb.b), io.EOF
}

// resolveFunc resolves the given host for the IP version.
// This type is mainly used for testing. For all other cases, a nil function
// should be passed to the httpRequestForTarget function.
type resolveFunc func(host string, ipVer int) (net.IP, error)

func hostWithPort(host string, port int) string {
	if port == 0 {
		return host
	}
	return fmt.Sprintf("%s:%d", host, port)
}

// hostHeaderForTarget computes request's Host header for a target.
//  - If host header is set in the probe, it overrides everything else.
//  - If target's fqdn is provided in its labels, use that along with the port.
//  - Finally, use target's name with port.
func hostHeaderForTarget(target endpoint.Endpoint, probeHostHeader string, port int) string {
	if probeHostHeader != "" {
		return probeHostHeader
	}

	if target.Labels["fqdn"] != "" {
		return hostWithPort(target.Labels["fqdn"], port)
	}

	return hostWithPort(target.Name, port)
}

func urlHostForTarget(target endpoint.Endpoint) string {
	if target.Labels["fqdn"] != "" {
		return target.Labels["fqdn"]
	}

	return target.Name
}

func relURLForTarget(target endpoint.Endpoint, probeURL string) string {
	if probeURL != "" {
		return probeURL
	}

	if target.Labels[relURLLabel] != "" {
		return target.Labels[relURLLabel]
	}

	return ""
}

func (p *Probe) httpRequestForTarget(target endpoint.Endpoint, resolveF resolveFunc) *http.Request {
	// Prepare HTTP.Request for Client.Do
	port := int(p.c.GetPort())
	// If port is not configured explicitly, use target's port if available.
	if port == 0 {
		port = target.Port
	}

	urlHost := urlHostForTarget(target)
	ipForLabel := ""

	if p.c.GetResolveFirst() {
		if resolveF == nil {
			resolveF = p.opts.Targets.Resolve
		}

		ip, err := resolveF(target.Name, p.opts.IPVersion)
		if err != nil {
			p.l.Error("target: ", target.Name, ", resolve error: ", err.Error())
			return nil
		}

		ipStr := ip.String()
		urlHost, ipForLabel = ipStr, ipStr
	}

	for _, al := range p.opts.AdditionalLabels {
		al.UpdateForTargetWithIPPort(target, ipForLabel, port)
	}

	// Put square brackets around literal IPv6 hosts. This is the same logic as
	// net.JoinHostPort, but we cannot use net.JoinHostPort as it works only for
	// non default ports.
	if strings.IndexByte(urlHost, ':') >= 0 {
		urlHost = "[" + urlHost + "]"
	}

	url := fmt.Sprintf("%s://%s%s", p.protocol, hostWithPort(urlHost, port), relURLForTarget(target, p.url))

	// Prepare request body
	var body io.Reader
	if len(p.requestBody) > 0 {
		body = &requestBody{p.requestBody}
	}
	req, err := http.NewRequest(p.method, url, body)
	if err != nil {
		p.l.Error("target: ", target.Name, ", error creating HTTP request: ", err.Error())
		return nil
	}

	req.ContentLength = int64(len(p.requestBody))

	var probeHostHeader string
	for _, header := range p.c.GetHeaders() {
		if header.GetName() == "Host" {
			probeHostHeader = header.GetValue()
			continue
		}
		req.Header.Set(header.GetName(), header.GetValue())
	}

	// Host header is set by http.NewRequest based on the URL, update it based
	// on various conditions.
	req.Host = hostHeaderForTarget(target, probeHostHeader, port)

	return req
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"flag"
	"time"
	"sync/atomic"
	"runtime"
)

import (
	"crypto"
	_ "crypto/md5"
	_ "crypto/sha1"
	_ "crypto/sha256"
	_ "crypto/sha512"
//	_ "golang.org/x/crypto/blake2b"
//	_ "golang.org/x/crypto/blake2s"
//	_ "golang.org/x/crypto/md4"
//	_ "golang.org/x/crypto/ripemd160"
//	_ "golang.org/x/crypto/sha3"
	"hash"
	"reflect"
)

import "github.com/splace/varbinary"
import "github.com/splace/fsflags"

const timeoutStatusCode = 124

// use variable binary encoding as nonce
type hashIndex struct{
	varbinary.Uint64
}

func main() {
	var leadingBitCount uint
	flag.UintVar(&leadingBitCount, "bits", 1, "Number of leading bits being searched for.")
	var bitState bool
	flag.BoolVar(&bitState, "set", false, "Leading bits set.")
	var hashType string
	flag.StringVar(&hashType, "hash", "SHA1", "hash type. one of \"MD4,MD5,SHA1,SHA224,SHA256,SHA384,SHA512,RIPEMD160,SHA3_224,SHA3_256,SHA3_384,SHA3_512,SHA512_224,SHA512_256\"")
	//flag.StringVar(&hashType, "hash", "SHA1", "hash type. one of MD4,MD5,SHA1,SHA224,SHA256,SHA384,SHA512,RIPEMD160,SHA3_224,SHA3_256,SHA3_384,SHA3_512,SHA512_224,SHA512_256,BLAKE2s_256,BLAKE2b_256,BLAKE2b_384,BLAKE2b_512")
	var startHashIndex uint64
	flag.Uint64Var(&startHashIndex, "start", 0, "Hash index to start search from.(default:#0)")
	var stopHashIndex uint64
	flag.Uint64Var(&stopHashIndex, "stop", 0, "Hash index to stop search at.(default:#0 = unlimited)")
	var logInterval time.Duration
	flag.DurationVar(&logInterval, "interval", time.Second, "time between log status reports.")
	var limit time.Duration
	flag.DurationVar(&limit, "end", 0, "search time limit.")
	var help bool
	flag.BoolVar(&help, "help", false, "display help/usage.")
	flag.BoolVar(&help, "h", false, "display help/usage.")
	var quiet bool
	flag.BoolVar(&quiet, "quiet", false, "no progress logging.")
	flag.BoolVar(&quiet, "q", false, "no progress logging.")
	var source fsflags.FileValue
	flag.Var(&source, "i", "input source bytes.(default:<Stdin>)")
	flag.Var(&source, "input", "input source bytes.(default:<Stdin>)")
	var sink fsflags.CreateFileValue
	flag.Var(&sink, "o", "output file, written with input file + nonce appended.(default:Stdout just written with nonce.)")
	flag.Var(&sink, "output", "output file, written with input file + nonce appended.(default:Stdout just written with nonce.)")
	var logToo fsflags.CreateFileValue
	flag.Var(&logToo, "log", "progress log destination.(default:Stderr)")
	flag.Parse()

	if help {
		flag.Usage()
		os.Exit(0)
	}

	// for optimisation #1.1
	var arrayOfBytePerms = [0x100][]byte{}
	for i:=range(arrayOfBytePerms){
		arrayOfBytePerms[i]=[]byte{byte(i)}
	}

	var baseHasher hash.Hash
	switch hashType {
	case "MD4":
		baseHasher = crypto.MD4.New()
	case "MD5":
		baseHasher = crypto.MD5.New()
	case "SHA1":
		baseHasher = crypto.SHA1.New()
	case "SHA224":
		baseHasher = crypto.SHA224.New()
	case "SHA256":
		baseHasher = crypto.SHA256.New()
	case "SHA384":
		baseHasher = crypto.SHA384.New()
	case "SHA512":
		baseHasher = crypto.SHA512.New()
	case "SHA512_224":
		baseHasher = crypto.SHA512_224.New()
	case "SHA512_256":
		baseHasher = crypto.SHA512_256.New()
//	case "RIPEMD160":
//		baseHasher = crypto.RIPEMD160.New()
//	case "SHA3_224":
//		baseHasher = crypto.SHA3_224.New()
//	case "SHA3_256":
//		baseHasher = crypto.SHA3_256.New()
//	case "SHA3_384":
//		baseHasher = crypto.SHA3_384.New()
//	case "SHA3_512":
//		baseHasher = crypto.SHA3_512.New()
//	case "BLAKE2s_256":
//		baseHasher = crypto.BLAKE2s_256.New()
//	case "BLAKE2b_256":
//		baseHasher = crypto.BLAKE2b_256.New()
//	case "BLAKE2b_384":
//		baseHasher = crypto.BLAKE2b_384.New()
//	case "BLAKE2b_512":
//		baseHasher = crypto.BLAKE2b_512.New()
	default:
		log.Fatalf("Aborting, Unknown Hash Scheme:" + hashType)
	}

	if logToo.File == nil {
		logToo.File = os.Stderr
	}

	var progressLog *log.Logger
	if quiet {
		progressLog=log.New(ioutil.Discard, "", log.LstdFlags)
		}else{
		progressLog=log.New(logToo, "", log.LstdFlags)
	}

	if source.File == nil {
		source.File = os.Stdin
	}

	progressLog.Printf("Loading:%q", &source)
	if sink.File == nil {
		sink.File = os.Stdout
		if _, err := io.Copy(baseHasher, source); err != nil {
			log.Fatal(err)
		}
		}else{
		// if not stdout write the source to the sink now, so later writing of the nonce then is appended.
		if _, err := io.Copy(baseHasher, io.TeeReader(source, sink)); err != nil {
			log.Fatal(err)
		}
	}
	source.Close()

	if leadingBitCount > 128 {
		log.Fatalf("Aborting, leading bits matching over 128 not supported (%v)", leadingBitCount)
	}

	// because of optimisation#1, need to find hash index with one byte removed.
	if startHashIndex>0{
		startHashIndex=uint64(hashIndexTruncate(hashIndex{Uint64:varbinary.Uint64(startHashIndex)},1).Uint64)
	}
	if stopHashIndex>0{
		stopHashIndex=uint64(hashIndexTruncate(hashIndex{Uint64:varbinary.Uint64(stopHashIndex)},1).Uint64)+1
		}else{
		stopHashIndex=1<<64-1
	}

	startTime := time.Now()
	doLog := time.NewTicker(logInterval)
	go func() {
		lhashIndex := startHashIndex
		for t := range doLog.C {
			runningFor := t.Sub(startTime)
			if limit > 0 && runningFor > limit {
				progressLog.Print("Aborting: time limit reached")
				os.Exit(timeoutStatusCode)
			}
			progressLog.Printf("#%d @%v\t%.0f#/s\tMean Match:%v", startHashIndex, runningFor/time.Second*time.Second, float64(startHashIndex-lhashIndex)/logInterval.Seconds(), (logInterval / time.Duration(startHashIndex-lhashIndex) * (1 << leadingBitCount) / time.Second * time.Second))
			lhashIndex = startHashIndex
		}
	}()

	var matchCondition func([]byte) bool
	if bitState{
		matchCondition = leadingSetBits(leadingBitCount)
		}else{
		matchCondition = leadingZeroBits(leadingBitCount)
	}

	if sum := baseHasher.Sum(nil); matchCondition(sum) {
		progressLog.Printf("Match on Source file as-is. Hash(%s):[% x]", hashType, sum)
		os.Exit(0)
	}

	stride := uint64(runtime.NumCPU())

	searchStripe := func(start uint64) {
		progressLog.Printf("Starting thread @ #%d", start)
		var hasher, branchHasher hash.Hash
		var sum []byte
		var n int
		buf := make([]byte,8,8)
		for hi := start; hi<=stopHashIndex ; hi += stride {
			hasher = cloneHash(baseHasher)
			n=varbinary.Uint64Put(varbinary.Uint64(hi),buf)
			hasher.Write(buf[:n])
			// optimisation#1: rather than check hash, with existing nonce, copy it and check all possible single bytes added to it. (+20% intel core2)
			for i := range arrayOfBytePerms { // optimisation#1.1: use pre-generated array of []byte for added byte. (+5% intel core2)
				branchHasher = cloneHash(hasher)
				branchHasher.Write(arrayOfBytePerms[i])
				sum = branchHasher.Sum(nil)
				if matchCondition(sum) {
					doLog.Stop()
					progressLog.Printf("#%d @%.1fs\tMatch:%q+[%s %x] Saving:%q Hash(%s):[% x]", uint64(hashIndexAppend(hashIndex{Uint64:varbinary.Uint64(hi)},arrayOfBytePerms[i][0]).Uint64), time.Since(startTime).Seconds(), &source, varbinary.Uint64(hi), arrayOfBytePerms[i], &sink, hashType, sum)
					n=varbinary.Uint64Put(varbinary.Uint64(hi),buf)
					sink.Write(buf[:n])
					sink.Write(arrayOfBytePerms[i])
					sink.Close()
					os.Exit(0)
				}
			}
			atomic.AddUint64(&startHashIndex, 0x100) //keep track of number checked, optimisation#1: each byte tested
		}
	}

	// start a thread for each core, each searching from different start indexes, but striding the same so always missing each other.
	for t := 0; t < runtime.NumCPU()-1; t++ {
		go func(s uint64) {
			searchStripe(s)
		}(startHashIndex)
		startHashIndex++
	}
	searchStripe(startHashIndex)
	progressLog.Printf("#%d @%.1fs Stopping Search of:%q", startHashIndex, time.Since(startTime).Seconds(), &source)

}

//  need to clone hash because what we want isn't exposed, and forking hashing isn't possible due to its calls to the runtime.
func cloneHash(h hash.Hash) hash.Hash{
	return clone(h).(hash.Hash)
}

// copy an interface value using reflection
func clone(i interface{}) interface{} {
	hv := reflect.ValueOf(i)
	hv1 := reflect.New(hv.Type().Elem())
	hv1.Elem().Set(hv.Elem())
	return hv1.Interface()
}

// return new hashindex whose representation is as the source hashindex but with added byte(s)
func hashIndexAppend(hi hashIndex, b ...byte) (nhi hashIndex) {
	buf,_:=hi.MarshalBinary()
	(&nhi).UnmarshalBinary(append(buf,b...))
	return
}

// return new hashindex whose representation is as the source hashindex but with removed byte(s)
func hashIndexTruncate(hi hashIndex,c int) (nhi hashIndex) {
	buf,_:=hi.MarshalBinary()
	(&nhi).UnmarshalBinary(buf[:len(buf)-c])
	return
}

<commit_msg>format<commit_after>package main

import (
	"flag"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"sync/atomic"
	"time"
)

import (
	"crypto"
	_ "crypto/md5"
	_ "crypto/sha1"
	_ "crypto/sha256"
	_ "crypto/sha512"
	//	_ "golang.org/x/crypto/blake2b"
	//	_ "golang.org/x/crypto/blake2s"
	//	_ "golang.org/x/crypto/md4"
	//	_ "golang.org/x/crypto/ripemd160"
	//	_ "golang.org/x/crypto/sha3"
	"hash"
	"reflect"
)

import "github.com/splace/varbinary"
import "github.com/splace/fsflags"

const timeoutStatusCode = 124

// use variable binary encoding as nonce
type hashIndex struct {
	varbinary.Uint64
}

func main() {
	var leadingBitCount uint
	flag.UintVar(&leadingBitCount, "bits", 1, "Number of leading bits being searched for.")
	var bitState bool
	flag.BoolVar(&bitState, "set", false, "Leading bits set.")
	var hashType string
	flag.StringVar(&hashType, "hash", "SHA1", "hash type. one of \"MD4,MD5,SHA1,SHA224,SHA256,SHA384,SHA512,RIPEMD160,SHA3_224,SHA3_256,SHA3_384,SHA3_512,SHA512_224,SHA512_256\"")
	//flag.StringVar(&hashType, "hash", "SHA1", "hash type. one of MD4,MD5,SHA1,SHA224,SHA256,SHA384,SHA512,RIPEMD160,SHA3_224,SHA3_256,SHA3_384,SHA3_512,SHA512_224,SHA512_256,BLAKE2s_256,BLAKE2b_256,BLAKE2b_384,BLAKE2b_512")
	var startHashIndex uint64
	flag.Uint64Var(&startHashIndex, "start", 0, "Hash index to start search from.(default:#0)")
	var stopHashIndex uint64
	flag.Uint64Var(&stopHashIndex, "stop", 0, "Hash index to stop search at.(default:#0 = unlimited)")
	var logInterval time.Duration
	flag.DurationVar(&logInterval, "interval", time.Second, "time between log status reports.")
	var limit time.Duration
	flag.DurationVar(&limit, "end", 0, "search time limit.")
	var help bool
	flag.BoolVar(&help, "help", false, "display help/usage.")
	flag.BoolVar(&help, "h", false, "display help/usage.")
	var quiet bool
	flag.BoolVar(&quiet, "quiet", false, "no progress logging.")
	flag.BoolVar(&quiet, "q", false, "no progress logging.")
	var source fsflags.FileValue
	flag.Var(&source, "i", "input source bytes.(default:<Stdin>)")
	flag.Var(&source, "input", "input source bytes.(default:<Stdin>)")
	var sink fsflags.CreateFileValue
	flag.Var(&sink, "o", "output file, written with input file + nonce appended.(default:Stdout just written with nonce.)")
	flag.Var(&sink, "output", "output file, written with input file + nonce appended.(default:Stdout just written with nonce.)")
	var logToo fsflags.CreateFileValue
	flag.Var(&logToo, "log", "progress log destination.(default:Stderr)")
	flag.Parse()

	if help {
		flag.Usage()
		os.Exit(0)
	}

	// for optimisation #1.1
	var arrayOfBytePerms = [0x100][]byte{}
	for i := range arrayOfBytePerms {
		arrayOfBytePerms[i] = []byte{byte(i)}
	}

	var baseHasher hash.Hash
	switch hashType {
	case "MD4":
		baseHasher = crypto.MD4.New()
	case "MD5":
		baseHasher = crypto.MD5.New()
	case "SHA1":
		baseHasher = crypto.SHA1.New()
	case "SHA224":
		baseHasher = crypto.SHA224.New()
	case "SHA256":
		baseHasher = crypto.SHA256.New()
	case "SHA384":
		baseHasher = crypto.SHA384.New()
	case "SHA512":
		baseHasher = crypto.SHA512.New()
	case "SHA512_224":
		baseHasher = crypto.SHA512_224.New()
	case "SHA512_256":
		baseHasher = crypto.SHA512_256.New()
		//	case "RIPEMD160":
		//		baseHasher = crypto.RIPEMD160.New()
		//	case "SHA3_224":
		//		baseHasher = crypto.SHA3_224.New()
		//	case "SHA3_256":
		//		baseHasher = crypto.SHA3_256.New()
		//	case "SHA3_384":
		//		baseHasher = crypto.SHA3_384.New()
		//	case "SHA3_512":
		//		baseHasher = crypto.SHA3_512.New()
		//	case "BLAKE2s_256":
		//		baseHasher = crypto.BLAKE2s_256.New()
		//	case "BLAKE2b_256":
		//		baseHasher = crypto.BLAKE2b_256.New()
		//	case "BLAKE2b_384":
		//		baseHasher = crypto.BLAKE2b_384.New()
		//	case "BLAKE2b_512":
		//		baseHasher = crypto.BLAKE2b_512.New()
	default:
		log.Fatalf("Aborting, Unknown Hash Scheme:" + hashType)
	}

	if logToo.File == nil {
		logToo.File = os.Stderr
	}

	var progressLog *log.Logger
	if quiet {
		progressLog = log.New(ioutil.Discard, "", log.LstdFlags)
	} else {
		progressLog = log.New(logToo, "", log.LstdFlags)
	}

	if source.File == nil {
		source.File = os.Stdin
	}

	progressLog.Printf("Loading:%q", &source)
	if sink.File == nil {
		sink.File = os.Stdout
		if _, err := io.Copy(baseHasher, source); err != nil {
			log.Fatal(err)
		}
	} else {
		// if not stdout write the source to the sink now, so later writing of the nonce then is appended.
		if _, err := io.Copy(baseHasher, io.TeeReader(source, sink)); err != nil {
			log.Fatal(err)
		}
	}
	source.Close()

	if leadingBitCount > 128 {
		log.Fatalf("Aborting, leading bits matching over 128 not supported (%v)", leadingBitCount)
	}

	// because of optimisation#1, need to find hash index with one byte removed.
	if startHashIndex > 0 {
		startHashIndex = uint64(hashIndexTruncate(hashIndex{Uint64: varbinary.Uint64(startHashIndex)}, 1).Uint64)
	}
	if stopHashIndex > 0 {
		stopHashIndex = uint64(hashIndexTruncate(hashIndex{Uint64: varbinary.Uint64(stopHashIndex)}, 1).Uint64) + 1
	} else {
		stopHashIndex = 1<<64 - 1
	}

	startTime := time.Now()
	doLog := time.NewTicker(logInterval)
	go func() {
		lhashIndex := startHashIndex
		for t := range doLog.C {
			runningFor := t.Sub(startTime)
			if limit > 0 && runningFor > limit {
				progressLog.Print("Aborting: time limit reached")
				os.Exit(timeoutStatusCode)
			}
			progressLog.Printf("#%d @%v\t%.0f#/s\tMean Match:%v", startHashIndex, runningFor/time.Second*time.Second, float64(startHashIndex-lhashIndex)/logInterval.Seconds(), (logInterval / time.Duration(startHashIndex-lhashIndex) * (1 << leadingBitCount) / time.Second * time.Second))
			lhashIndex = startHashIndex
		}
	}()

	var matchCondition func([]byte) bool
	if bitState {
		matchCondition = leadingSetBits(leadingBitCount)
	} else {
		matchCondition = leadingZeroBits(leadingBitCount)
	}

	if sum := baseHasher.Sum(nil); matchCondition(sum) {
		progressLog.Printf("Match on Source file as-is. Hash(%s):[% x]", hashType, sum)
		os.Exit(0)
	}

	stride := uint64(runtime.NumCPU())

	searchStripe := func(start uint64) {
		progressLog.Printf("Starting thread @ #%d", start)
		var hasher, branchHasher hash.Hash
		var sum []byte
		var n int
		buf := make([]byte, 8, 8)
		for hi := start; hi <= stopHashIndex; hi += stride {
			hasher = cloneHash(baseHasher)
			n = varbinary.Uint64Put(varbinary.Uint64(hi), buf)
			hasher.Write(buf[:n])
			// optimisation#1: rather than check hash, with existing nonce, copy it and check all possible single bytes added to it. (+20% intel core2)
			for i := range arrayOfBytePerms { // optimisation#1.1: use pre-generated array of []byte for added byte. (+5% intel core2)
				branchHasher = cloneHash(hasher)
				branchHasher.Write(arrayOfBytePerms[i])
				sum = branchHasher.Sum(nil)
				if matchCondition(sum) {
					doLog.Stop()
					progressLog.Printf("#%d @%.1fs\tMatch:%q+[%s %x] Saving:%q Hash(%s):[% x]", uint64(hashIndexAppend(hashIndex{Uint64: varbinary.Uint64(hi)}, arrayOfBytePerms[i][0]).Uint64), time.Since(startTime).Seconds(), &source, varbinary.Uint64(hi), arrayOfBytePerms[i], &sink, hashType, sum)
					n = varbinary.Uint64Put(varbinary.Uint64(hi), buf)
					sink.Write(buf[:n])
					sink.Write(arrayOfBytePerms[i])
					sink.Close()
					os.Exit(0)
				}
			}
			atomic.AddUint64(&startHashIndex, 0x100) //keep track of number checked, optimisation#1: each byte tested
		}
	}

	// start a thread for each core, each searching from different start indexes, but striding the same so always missing each other.
	for t := 0; t < runtime.NumCPU()-1; t++ {
		go func(s uint64) {
			searchStripe(s)
		}(startHashIndex)
		startHashIndex++
	}
	searchStripe(startHashIndex)
	progressLog.Printf("#%d @%.1fs Stopping Search of:%q", startHashIndex, time.Since(startTime).Seconds(), &source)

}

//  need to clone hash because what we want isn't exposed, and forking hashing isn't possible due to its calls to the runtime.
func cloneHash(h hash.Hash) hash.Hash {
	return clone(h).(hash.Hash)
}

// copy an interface value using reflection
func clone(i interface{}) interface{} {
	hv := reflect.ValueOf(i)
	hv1 := reflect.New(hv.Type().Elem())
	hv1.Elem().Set(hv.Elem())
	return hv1.Interface()
}

// return new hashindex whose representation is as the source hashindex but with added byte(s)
func hashIndexAppend(hi hashIndex, b ...byte) (nhi hashIndex) {
	buf, _ := hi.MarshalBinary()
	(&nhi).UnmarshalBinary(append(buf, b...))
	return
}

// return new hashindex whose representation is as the source hashindex but with removed byte(s)
func hashIndexTruncate(hi hashIndex, c int) (nhi hashIndex) {
	buf, _ := hi.MarshalBinary()
	(&nhi).UnmarshalBinary(buf[:len(buf)-c])
	return
}

<|endoftext|>


<commit_before>package tcp

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/bbangert/toml"
	"github.com/millken/kaman/metrics"
	"github.com/millken/kaman/plugins"
)

// Input plugin implementation that listens for Heka proself.col messages on a
// specified TCP sockeself.g Creates a separate goroutine for each TCP connection.
type TcpInput struct {
	keepAliveDuration time.Duration
	listener          net.Listener
	wg                sync.WaitGroup
	stopChan          chan bool
	config            *TcpInputConfig
	common            *plugins.PluginCommonConfig
	runner            plugins.InputRunner
}

type TcpInputConfig struct {
	// Network type (e.g. "tcp", "tcp4", "tcp6", "unix" or "unixpacket").
	// Needs to match the input type.
	Net string
	// String representation of the address of the network connection on which
	// the listener should be listening (e.g. "127.0.0.1:5565").
	Address string
	// Set to true if TCP Keep Alive should be used.
	KeepAlive bool `toml:"keep_alive"`
	// Integer indicating seconds between keep alives.
	KeepAlivePeriod int `toml:"keep_alive_period"`
}

func (self *TcpInput) Init(pcf *plugins.PluginCommonConfig, conf toml.Primitive) (err error) {

	log.Println("TcpInput Init")
	self.common = pcf
	self.config = &TcpInputConfig{
		Net: "tcp",
	}
	if err := toml.PrimitiveDecode(conf, self.config); err != nil {
		return fmt.Errorf("Can't unmarshal TcpInput config: %s", err)
	}
	address, err := net.ResolveTCPAddr(self.config.Net, self.config.Address)
	if err != nil {
		return fmt.Errorf("ResolveTCPAddress failed: %s\n", err.Error())
	}
	self.listener, err = net.ListenTCP(self.config.Net, address)
	if err != nil {
		return fmt.Errorf("ListenTCP failed: %s\n", err.Error())
	}
	// We're already listening, make sure we clean up if init fails later on.
	closeIt := true
	defer func() {
		if closeIt {
			self.listener.Close()
		}
	}()
	if self.config.KeepAlivePeriod != 0 {
		self.keepAliveDuration = time.Duration(self.config.KeepAlivePeriod) * time.Second
	}
	self.stopChan = make(chan bool)
	closeIt = false
	return nil
}

// Listen on the provided TCP connection, extracting messages from the incoming
// data until the connection is closed or Stop is called on the input.
func (self *TcpInput) handleConnection(conn net.Conn) {
	var (frag []byte
		 err error = nil
		 pack *plugins.PipelinePack
		)
	//raddr := conn.RemoteAddr().String()
	//host, _, err := net.SplitHostPort(raddr)
	//if err != nil {
	//	host = raddr
	//}
	//log.Printf("handle conn: %s, host: %s", raddr, host)
	counter := fmt.Sprintf("Tag:%s,Type:%s", self.common.Tag, self.common.Type)
	mc := metrics.NewCounter(counter)
	defer func() {
		conn.Close()
		self.wg.Done()
	}()

	buf := make([]byte, 1024)
	b1 := []byte{}
	count := 0
	stopped := false
	reader := bufio.NewReaderSize(conn, 8192)

	ticker := time.NewTicker(time.Duration(1) * time.Minute)
	defer ticker.Stop()
	for !stopped {
		conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		select {
		case <-self.stopChan:
			stopped = true
		case <-ticker.C:
			if count == 0 {
				//log.Printf("remove unused conn : %s", raddr)
				stopped = true
			}
			count = 0
		default:
			frag, err = reader.ReadSlice('\n')
			if err != nil {
				//log.Printf("disconnect : %s", raddr)
				stopped = true
			}

			if len(frag) == 0 {
				continue
			}
			buf = append(b1, frag...)
			count++
			pack = <-self.runner.InChan()
			pack.MsgBytes = bytes.TrimSpace(buf[:])
			pack.Msg.Tag = self.common.Tag
			pack.Msg.Timestamp = time.Now().Unix()
			mc.Add(1)
			self.runner.RouterChan() <- pack
			buf = buf[:0]
		}
	}
	buf = nil
	reader = nil
}

func (self *TcpInput) Run(runner plugins.InputRunner) error {
	var (
		conn net.Conn
		e    error
	)
	self.runner = runner
	for {
		if conn, e = self.listener.Accept(); e != nil {
			if netErr, ok := e.(net.Error); ok && netErr.Temporary() {
				log.Print(fmt.Errorf("TCP accept failed: %s", e))
				continue
			} else {
				break
			}
		}
		if self.config.KeepAlive {
			tcpConn, ok := conn.(*net.TCPConn)
			if !ok {
				return errors.New("KeepAlive only supported for TCP Connections.")
			}
			tcpConn.SetKeepAlive(self.config.KeepAlive)
			if self.keepAliveDuration != 0 {
				tcpConn.SetKeepAlivePeriod(self.keepAliveDuration)
			}
		}
		self.wg.Add(1)
		go self.handleConnection(conn)
	}
	self.wg.Wait()
	return e
}

func init() {
	plugins.RegisterInput("TcpInput", func() interface{} {
		return new(TcpInput)
	})
}
<commit_msg>limit goroutine run time<commit_after>package tcp

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/bbangert/toml"
	"github.com/millken/kaman/metrics"
	"github.com/millken/kaman/plugins"
)

// Input plugin implementation that listens for Heka proself.col messages on a
// specified TCP sockeself.g Creates a separate goroutine for each TCP connection.
type TcpInput struct {
	keepAliveDuration time.Duration
	listener          net.Listener
	wg                sync.WaitGroup
	stopChan          chan bool
	config            *TcpInputConfig
	common            *plugins.PluginCommonConfig
	runner            plugins.InputRunner
}

type TcpInputConfig struct {
	// Network type (e.g. "tcp", "tcp4", "tcp6", "unix" or "unixpacket").
	// Needs to match the input type.
	Net string
	// String representation of the address of the network connection on which
	// the listener should be listening (e.g. "127.0.0.1:5565").
	Address string
	// Set to true if TCP Keep Alive should be used.
	KeepAlive bool `toml:"keep_alive"`
	// Integer indicating seconds between keep alives.
	KeepAlivePeriod int `toml:"keep_alive_period"`
}

func (self *TcpInput) Init(pcf *plugins.PluginCommonConfig, conf toml.Primitive) (err error) {

	log.Println("TcpInput Init")
	self.common = pcf
	self.config = &TcpInputConfig{
		Net: "tcp",
	}
	if err := toml.PrimitiveDecode(conf, self.config); err != nil {
		return fmt.Errorf("Can't unmarshal TcpInput config: %s", err)
	}
	address, err := net.ResolveTCPAddr(self.config.Net, self.config.Address)
	if err != nil {
		return fmt.Errorf("ResolveTCPAddress failed: %s\n", err.Error())
	}
	self.listener, err = net.ListenTCP(self.config.Net, address)
	if err != nil {
		return fmt.Errorf("ListenTCP failed: %s\n", err.Error())
	}
	// We're already listening, make sure we clean up if init fails later on.
	closeIt := true
	defer func() {
		if closeIt {
			self.listener.Close()
		}
	}()
	if self.config.KeepAlivePeriod != 0 {
		self.keepAliveDuration = time.Duration(self.config.KeepAlivePeriod) * time.Second
	}
	self.stopChan = make(chan bool)
	closeIt = false
	return nil
}

// Listen on the provided TCP connection, extracting messages from the incoming
// data until the connection is closed or Stop is called on the input.
func (self *TcpInput) handleConnection(conn net.Conn) {
	var (frag []byte
		 err error = nil
		 pack *plugins.PipelinePack
		)
	//raddr := conn.RemoteAddr().String()
	//host, _, err := net.SplitHostPort(raddr)
	//if err != nil {
	//	host = raddr
	//}
	//log.Printf("handle conn: %s, host: %s", raddr, host)
	counter := fmt.Sprintf("Tag:%s,Type:%s", self.common.Tag, self.common.Type)
	mc := metrics.NewCounter(counter)
	defer func() {
		conn.Close()
		self.wg.Done()
	}()

	buf := make([]byte, 1024)
	b1 := []byte{}
	count := 0
	limit_run_times := 60
	stopped := false
	reader := bufio.NewReaderSize(conn, 8192)

	ticker := time.NewTicker(time.Duration(1) * time.Minute)
	defer ticker.Stop()
	for !stopped {
		conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		select {
		case <-self.stopChan:
			stopped = true
		case <-ticker.C:
			if count == 0 || limit_run_times <= 0 {
				//log.Printf("remove unused conn : %s", raddr)
				stopped = true
			}
			limit_run_times -= 1
			count = 0
		default:
			frag, err = reader.ReadSlice('\n')
			if err != nil {
				//log.Printf("disconnect : %s", raddr)
				stopped = true
			}

			if len(frag) == 0 {
				continue
			}
			buf = append(b1, frag...)
			count++
			pack = <-self.runner.InChan()
			pack.MsgBytes = bytes.TrimSpace(buf[:])
			pack.Msg.Tag = self.common.Tag
			pack.Msg.Timestamp = time.Now().Unix()
			mc.Add(1)
			self.runner.RouterChan() <- pack
			buf = buf[:0]
		}
	}
	buf = nil
	reader = nil
}

func (self *TcpInput) Run(runner plugins.InputRunner) error {
	var (
		conn net.Conn
		e    error
	)
	self.runner = runner
	for {
		if conn, e = self.listener.Accept(); e != nil {
			if netErr, ok := e.(net.Error); ok && netErr.Temporary() {
				log.Print(fmt.Errorf("TCP accept failed: %s", e))
				continue
			} else {
				break
			}
		}
		if self.config.KeepAlive {
			tcpConn, ok := conn.(*net.TCPConn)
			if !ok {
				return errors.New("KeepAlive only supported for TCP Connections.")
			}
			tcpConn.SetKeepAlive(self.config.KeepAlive)
			if self.keepAliveDuration != 0 {
				tcpConn.SetKeepAlivePeriod(self.keepAliveDuration)
			}
		}
		self.wg.Add(1)
		go self.handleConnection(conn)
	}
	self.wg.Wait()
	return e
}

func init() {
	plugins.RegisterInput("TcpInput", func() interface{} {
		return new(TcpInput)
	})
}
<|endoftext|>


<commit_before>package goprocnet

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

const (
	PROC_NET_TCP = "/proc/net/tcp"
	PROC_NET_UDP = "/proc/net/udp"
)

var STATE = map[string]string{
	"01": "ESTABLISHED",
	"02": "SYN_SENT",
	"03": "SYN_RECV",
	"04": "FIN_WAIT1",
	"05": "FIN_WAIT2",
	"06": "TIME_WAIT",
	"07": "CLOSE",
	"08": "CLOSE_WAIT",
	"09": "LAST_ACK",
	"0A": "LISTEN",
	"0B": "CLOSING",
}

var Version string

type Socket struct {
	PID        string
	UID        string
	Bin        string
	Name       string
	State      string
	LocalIP    string
	LocalPort  string
	RemoteIP   string
	RemotePort string
}

func (s Socket) String() string {
	return fmt.Sprintf("%s\t%s\t%s\t%s\t%s:%s\t%s:%s",
		s.PID,
		s.UID,
		s.Bin,
		s.State,
		s.LocalIP,
		s.LocalPort,
		s.RemoteIP,
		s.RemotePort)
}

func GetTCPSockets() []Socket {
	sockets, err := netstat("tcp")
	if err != nil {
		log.Panic(err)
	}
	return sockets
}

func GetUDPSockets() []Socket {
	sockets, err := netstat("udp")
	if err != nil {
		log.Panic(err)
	}
	return sockets
}

func getFilename(t string) string {
	switch {
	case t == "tcp":
		return PROC_NET_TCP
	case t == "udp":
		return PROC_NET_UDP
	}
	return ""
}

func readFile(t string) ([]string, error) {
	filename := getFilename(t)
	if filename == "" {
	}

	data, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Println("Unable to read file", err)
		return nil, err
	}
	lines := strings.Split(string(data), "\n")

	// return lines without header or tailing blank line
	return lines[1 : len(lines)-1], nil
}

func convertHexToDec(hex string) int64 {
	d, err := strconv.ParseInt(hex, 16, 64)
	if err != nil {
		log.Panic("Unable to convert hex to dec", err)
	}
	return d
}

func getIP(hexip string) string {
	// we are going to assume IPv4 addresses for now...
	// We are also assuming address is in little endian
	return fmt.Sprintf("%v.%v.%v.%v",
		convertHexToDec(hexip[6:8]),
		convertHexToDec(hexip[4:6]),
		convertHexToDec(hexip[2:4]),
		convertHexToDec(hexip[0:2]))
}

func getPort(hex string) string {
	return fmt.Sprintf("%v", convertHexToDec(hex))
}

func removeInnerSpace(values []string) []string {
	var newValues []string
	for _, value := range values {
		if value != "" {
			newValues = append(newValues, value)
		}
	}
	return newValues
}

func netstat(t string) ([]Socket, error) {
	lines, err := readFile(t)
	if err != nil {
		log.Println("Failed to retrieve proc data", err)
		return nil, err
	}

	var sockets []Socket

	for _, line := range lines {
		values := removeInnerSpace(strings.Split(strings.TrimSpace(line), " "))

		ipPort := strings.Split(values[1], ":")
		localIP := getIP(ipPort[0])
		localPort := getPort(ipPort[1])

		remIpPort := strings.Split(values[2], ":")
		remoteIP := getIP(remIpPort[0])
		remotePort := getPort(remIpPort[1])

		state := STATE[values[3]]
		uid := values[7]
		pid := getProcessID(values[9])
		exe := getProcessExecutable(pid)
		name := getProcessName(exe)
		s := Socket{pid, uid, exe, name, state, localIP, localPort, remoteIP, remotePort}
		sockets = append(sockets, s)
	}
	return sockets, nil
}

func getProcessID(inode string) string {
	// Loop over the fd directories under /proc/pid for the right inode
	pid := "-"

	procDirs, err := filepath.Glob("/proc/[0-9]*/fd/[0-9]*")
	if err != nil {
		log.Panic(err)
	}

	re := regexp.MustCompile(inode)
	for _, item := range procDirs {
		path, _ := os.Readlink(item)
		out := re.FindString(path)
		if len(out) != 0 {
			pid = strings.Split(item, "/")[2]
		}
	}
	return pid
}

func getProcessExecutable(pid string) string {
	path, _ := os.Readlink(fmt.Sprintf("/proc/%s/exe", pid))
	return path
}

func getProcessName(exe string) string {
	n := strings.Split(exe, "/")
	name := n[len(n)-1]
	return strings.Title(name)
}
<commit_msg>Change Socket Port Types to `int32`<commit_after>package goprocnet

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

const (
	PROC_NET_TCP = "/proc/net/tcp"
	PROC_NET_UDP = "/proc/net/udp"
)

var STATE = map[string]string{
	"01": "ESTABLISHED",
	"02": "SYN_SENT",
	"03": "SYN_RECV",
	"04": "FIN_WAIT1",
	"05": "FIN_WAIT2",
	"06": "TIME_WAIT",
	"07": "CLOSE",
	"08": "CLOSE_WAIT",
	"09": "LAST_ACK",
	"0A": "LISTEN",
	"0B": "CLOSING",
}

var Version string

type Socket struct {
	PID        string
	UID        string
	Bin        string
	Name       string
	State      string
	LocalIP    string
	LocalPort  int32
	RemoteIP   string
	RemotePort int32
}

func (s Socket) String() string {
	return fmt.Sprintf("%s\t%s\t%s\t%s\t%s:%d\t%s:%d",
		s.PID,
		s.UID,
		s.Bin,
		s.State,
		s.LocalIP,
		s.LocalPort,
		s.RemoteIP,
		s.RemotePort)
}

func GetTCPSockets() []Socket {
	sockets, err := netstat("tcp")
	if err != nil {
		log.Panic(err)
	}
	return sockets
}

func GetUDPSockets() []Socket {
	sockets, err := netstat("udp")
	if err != nil {
		log.Panic(err)
	}
	return sockets
}

func getFilename(t string) string {
	switch {
	case t == "tcp":
		return PROC_NET_TCP
	case t == "udp":
		return PROC_NET_UDP
	}
	return ""
}

func readFile(t string) ([]string, error) {
	filename := getFilename(t)
	if filename == "" {
	}

	data, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Println("Unable to read file", err)
		return nil, err
	}
	lines := strings.Split(string(data), "\n")

	// return lines without header or tailing blank line
	return lines[1 : len(lines)-1], nil
}

func convertHexToDec(hex string) int64 {
	d, err := strconv.ParseInt(hex, 16, 64)
	if err != nil {
		log.Panic("Unable to convert hex to dec", err)
	}
	return d
}

func getIP(hexip string) string {
	// we are going to assume IPv4 addresses for now...
	// We are also assuming address is in little endian
	return fmt.Sprintf("%v.%v.%v.%v",
		convertHexToDec(hexip[6:8]),
		convertHexToDec(hexip[4:6]),
		convertHexToDec(hexip[2:4]),
		convertHexToDec(hexip[0:2]))
}

func getPort(hex string) string {
	return fmt.Sprintf("%v", convertHexToDec(hex))
}

func removeInnerSpace(values []string) []string {
	var newValues []string
	for _, value := range values {
		if value != "" {
			newValues = append(newValues, value)
		}
	}
	return newValues
}

func netstat(t string) ([]Socket, error) {
	lines, err := readFile(t)
	if err != nil {
		log.Println("Failed to retrieve proc data", err)
		return nil, err
	}

	var sockets []Socket

	mustParse := func(p string) int32 {
		i, err := strconv.ParseInt(p, 10, 32)
		if err != nil {
			log.Panic(err)
		}
		return int32(i)
	}

	for _, line := range lines {
		values := removeInnerSpace(strings.Split(strings.TrimSpace(line), " "))

		ipPort := strings.Split(values[1], ":")
		localIP := getIP(ipPort[0])
		localPort := mustParse(getPort(ipPort[1]))

		remIpPort := strings.Split(values[2], ":")
		remoteIP := getIP(remIpPort[0])
		remotePort := mustParse(getPort(remIpPort[1]))

		state := STATE[values[3]]
		uid := values[7]
		pid := getProcessID(values[9])
		exe := getProcessExecutable(pid)
		name := getProcessName(exe)
		s := Socket{pid, uid, exe, name, state, localIP, localPort, remoteIP, remotePort}
		sockets = append(sockets, s)
	}
	return sockets, nil
}

func getProcessID(inode string) string {
	// Loop over the fd directories under /proc/pid for the right inode
	pid := "-"

	procDirs, err := filepath.Glob("/proc/[0-9]*/fd/[0-9]*")
	if err != nil {
		log.Panic(err)
	}

	re := regexp.MustCompile(inode)
	for _, item := range procDirs {
		path, _ := os.Readlink(item)
		out := re.FindString(path)
		if len(out) != 0 {
			pid = strings.Split(item, "/")[2]
		}
	}
	return pid
}

func getProcessExecutable(pid string) string {
	path, _ := os.Readlink(fmt.Sprintf("/proc/%s/exe", pid))
	return path
}

func getProcessName(exe string) string {
	n := strings.Split(exe, "/")
	name := n[len(n)-1]
	return strings.Title(name)
}
<|endoftext|>


<commit_before>// Copyright 2017 Stratumn SAS. 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 postgresstore

import "database/sql"

const (
	sqlCreateLink = `
		INSERT INTO links (
			link_hash,
			priority,
			map_id,
			prev_link_hash,
			tags,
			data,
			process
		)
		VALUES ($1, $2, $3, $4, $5, $6, $7)
		ON CONFLICT (link_hash)
		DO UPDATE SET
			priority = $2,
			map_id = $3,
			prev_link_hash = $4,
			tags = $5,
			data = $6,
			process = $7
	`
	sqlGetSegment = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = $1
	`
	sqlDeleteLink = `
		DELETE FROM links
		WHERE link_hash = $1
		RETURNING data
	`
	sqlFindSegments = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE (length($3) = 0 OR process = $3)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $1 LIMIT $2
	`
	sqlFindSegmentsWithLinkHashes = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = any($1::bytea[])
		AND (length($4) = 0 OR process = $4)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $2 LIMIT $3
	`
	sqlFindSegmentsWithLinkHashesAndMapIDs = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = any($1::bytea[])
		AND map_id = any($2::text[])
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithLinkHashesAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = any($1::bytea[])
		AND tags @> $2
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithLinkHashesAndMapIDsAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = any($1::bytea[])
		AND map_id = any($2::text[]) AND tags @> $3
		AND (length($6) = 0 OR process = $6)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $4 LIMIT $5
	`
	sqlFindSegmentsWithMapIDs = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE map_id = any($1::text[])
		AND (length($4) = 0 OR process = $4)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $2 LIMIT $3
	`
	sqlFindSegmentsWithPrevLinkHash = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE prev_link_hash = $1
		AND (length($4) = 0 OR process = $4)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $2 LIMIT $3
	`
	sqlFindSegmentsWithTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE tags @> $1
		AND (length($4) = 0 OR process = $4)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $2 LIMIT $3
	`
	sqlFindSegmentsWithMapIDsAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE map_id = any($1::text[]) AND tags @> $2
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithPrevLinkHashAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE prev_link_hash = $1 AND tags @> $2
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithPrevLinkHashAndMapIDs = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE prev_link_hash = $1
		AND map_id = any($2::text[])
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithPrevLinkHashAndMapIDsAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE prev_link_hash = $1
		AND map_id = any($2::text[]) AND tags @> $3
		AND (length($6) = 0 OR process = $6)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $4 LIMIT $5
	`
	sqlGetMapIDs = `
		SELECT DISTINCT map_id FROM links
		WHERE (length($3) = 0 OR process = $3)
		ORDER BY map_id
		OFFSET $1 LIMIT $2
	`
	sqlSaveValue = `
		INSERT INTO values (
			key,
			value
		)
		VALUES ($1, $2)
		ON CONFLICT (key)
		DO UPDATE SET
			value = $2
	`
	sqlGetValue = `
		SELECT value FROM values
		WHERE key = $1
	`
	sqlDeleteValue = `
		DELETE FROM values
		WHERE key = $1
		RETURNING value
	`
	sqlGetEvidences = `
		SELECT data FROM evidences
		WHERE link_hash = $1
	`
	sqlAddEvidence = `
		INSERT INTO evidences (
			link_hash,
			provider,
			data
		)
		VALUES ($1, $2, $3)
		ON CONFLICT (link_hash, provider)
		DO NOTHING
	`
)

var sqlCreate = []string{
	`
		CREATE TABLE links (
			id BIGSERIAL PRIMARY KEY,
			link_hash bytea NOT NULL,
			priority double precision NOT NULL,
			map_id text NOT NULL,
			prev_link_hash bytea DEFAULT NULL,
			tags text[] DEFAULT NULL,
			data jsonb NOT NULL,
			process text NOT NULL,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
		)
	`,
	`
		CREATE UNIQUE INDEX links_link_hash_idx
		ON links (link_hash)
	`,
	`
		CREATE INDEX links_priority_created_at_idx
		ON links (priority DESC, created_at DESC)
	`,
	`
		CREATE INDEX links_map_id_idx
		ON links (map_id)
	`,
	`
		CREATE INDEX links_map_id_priority_created_at_idx
		ON links (map_id, priority DESC, created_at DESC)
	`,
	`
		CREATE INDEX links_prev_link_hash_priority_created_at_idx
		ON links (prev_link_hash, priority DESC, created_at DESC)
	`,
	`
		CREATE INDEX links_tags_idx
		ON links USING gin(tags)
	`,
	`
		CREATE TABLE evidences (
			id BIGSERIAL PRIMARY KEY,
			link_hash bytea NOT NULL,
			provider text NOT NULL,
			data jsonb NOT NULL,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
	)
	`,
	`
		CREATE UNIQUE INDEX evidences_link_hash_provider_idx
		ON evidences (link_hash, provider)
	`,
	`
		CREATE INDEX evidences_link_hash_idx
		ON evidences (link_hash)
	`,
	`
		CREATE TABLE values (
			id BIGSERIAL PRIMARY KEY,
			key bytea NOT NULL,
			value bytea NOT NULL,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
		)
	`,
	`
		CREATE UNIQUE INDEX values_key_idx
		ON values (key)
	`,
}

var sqlDrop = []string{
	"DROP TABLE links, evidences, values",
}

type writeStmts struct {
	CreateLink  *sql.Stmt
	DeleteLink  *sql.Stmt
	SaveValue   *sql.Stmt
	DeleteValue *sql.Stmt
	AddEvidence *sql.Stmt
}

type readStmts struct {
	GetSegment   *sql.Stmt
	FindSegments *sql.Stmt
	GetMapIDs    *sql.Stmt
	GetValue     *sql.Stmt
	GetEvidences *sql.Stmt

	FindSegmentsWithMapIDs                       *sql.Stmt
	FindSegmentsWithPrevLinkHash                 *sql.Stmt
	FindSegmentsWithTags                         *sql.Stmt
	FindSegmentsWithLinkHashes                   *sql.Stmt
	FindSegmentsWithLinkHashesAndMapIDs          *sql.Stmt
	FindSegmentsWithLinkHashesAndTags            *sql.Stmt
	FindSegmentsWithLinkHashesAndMapIDsAndTags   *sql.Stmt
	FindSegmentsWithTagsAndLimit                 *sql.Stmt
	FindSegmentsWithMapIDsAndTags                *sql.Stmt
	FindSegmentsWithPrevLinkHashAndTags          *sql.Stmt
	FindSegmentsWithPrevLinkHashAndMapIDs        *sql.Stmt
	FindSegmentsWithPrevLinkHashAndMapIDsAndTags *sql.Stmt
}

type stmts struct {
	readStmts
	writeStmts
}

type batchStmts stmts

func newStmts(db *sql.DB) (*stmts, error) {
	var (
		s   stmts
		err error
	)

	prepare := func(str string) (stmt *sql.Stmt) {
		if err == nil {
			stmt, err = db.Prepare(str)
		}
		return
	}

	s.GetSegment = prepare(sqlGetSegment)
	s.FindSegments = prepare(sqlFindSegments)
	s.GetMapIDs = prepare(sqlGetMapIDs)
	s.GetValue = prepare(sqlGetValue)
	s.GetEvidences = prepare(sqlGetEvidences)

	s.FindSegmentsWithMapIDs = prepare(sqlFindSegmentsWithMapIDs)
	s.FindSegmentsWithPrevLinkHash = prepare(sqlFindSegmentsWithPrevLinkHash)
	s.FindSegmentsWithTags = prepare(sqlFindSegmentsWithTags)

	s.FindSegmentsWithLinkHashes = prepare(sqlFindSegmentsWithLinkHashes)
	s.FindSegmentsWithLinkHashesAndMapIDs = prepare(sqlFindSegmentsWithLinkHashesAndMapIDs)
	s.FindSegmentsWithLinkHashesAndTags = prepare(sqlFindSegmentsWithLinkHashesAndTags)
	s.FindSegmentsWithLinkHashesAndMapIDsAndTags = prepare(sqlFindSegmentsWithLinkHashesAndMapIDsAndTags)

	s.FindSegmentsWithMapIDsAndTags = prepare(sqlFindSegmentsWithMapIDsAndTags)
	s.FindSegmentsWithPrevLinkHashAndTags = prepare(sqlFindSegmentsWithPrevLinkHashAndTags)
	s.FindSegmentsWithPrevLinkHashAndMapIDs = prepare(sqlFindSegmentsWithPrevLinkHashAndMapIDs)
	s.FindSegmentsWithPrevLinkHashAndMapIDsAndTags = prepare(sqlFindSegmentsWithPrevLinkHashAndMapIDsAndTags)

	s.CreateLink = prepare(sqlCreateLink)
	s.DeleteLink = prepare(sqlDeleteLink)
	s.SaveValue = prepare(sqlSaveValue)
	s.DeleteValue = prepare(sqlDeleteValue)
	s.AddEvidence = prepare(sqlAddEvidence)

	if err != nil {
		return nil, err
	}

	return &s, nil
}

func newBatchStmts(tx *sql.Tx) (*batchStmts, error) {
	var (
		s   batchStmts
		err error
	)

	prepare := func(str string) (stmt *sql.Stmt) {
		if err == nil {
			stmt, err = tx.Prepare(str)
		}
		return
	}

	s.GetSegment = prepare(sqlGetSegment)
	s.FindSegments = prepare(sqlFindSegments)
	s.GetMapIDs = prepare(sqlGetMapIDs)
	s.GetValue = prepare(sqlGetValue)

	s.CreateLink = prepare(sqlCreateLink)
	s.DeleteLink = prepare(sqlDeleteLink)
	s.SaveValue = prepare(sqlSaveValue)
	s.DeleteValue = prepare(sqlDeleteValue)

	if err != nil {
		return nil, err
	}

	return &s, nil
}
<commit_msg>postgresstore: sort mapIds by last update<commit_after>// Copyright 2017 Stratumn SAS. 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 postgresstore

import "database/sql"

const (
	sqlCreateLink = `
		INSERT INTO links (
			link_hash,
			priority,
			map_id,
			prev_link_hash,
			tags,
			data,
			process
		)
		VALUES ($1, $2, $3, $4, $5, $6, $7)
		ON CONFLICT (link_hash)
		DO UPDATE SET
			priority = $2,
			map_id = $3,
			prev_link_hash = $4,
			tags = $5,
			data = $6,
			process = $7
	`
	sqlGetSegment = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = $1
	`
	sqlDeleteLink = `
		DELETE FROM links
		WHERE link_hash = $1
		RETURNING data
	`
	sqlFindSegments = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE (length($3) = 0 OR process = $3)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $1 LIMIT $2
	`
	sqlFindSegmentsWithLinkHashes = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = any($1::bytea[])
		AND (length($4) = 0 OR process = $4)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $2 LIMIT $3
	`
	sqlFindSegmentsWithLinkHashesAndMapIDs = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = any($1::bytea[])
		AND map_id = any($2::text[])
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithLinkHashesAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = any($1::bytea[])
		AND tags @> $2
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithLinkHashesAndMapIDsAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE l.link_hash = any($1::bytea[])
		AND map_id = any($2::text[]) AND tags @> $3
		AND (length($6) = 0 OR process = $6)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $4 LIMIT $5
	`
	sqlFindSegmentsWithMapIDs = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE map_id = any($1::text[])
		AND (length($4) = 0 OR process = $4)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $2 LIMIT $3
	`
	sqlFindSegmentsWithPrevLinkHash = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE prev_link_hash = $1
		AND (length($4) = 0 OR process = $4)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $2 LIMIT $3
	`
	sqlFindSegmentsWithTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE tags @> $1
		AND (length($4) = 0 OR process = $4)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $2 LIMIT $3
	`
	sqlFindSegmentsWithMapIDsAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE map_id = any($1::text[]) AND tags @> $2
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithPrevLinkHashAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE prev_link_hash = $1 AND tags @> $2
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithPrevLinkHashAndMapIDs = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE prev_link_hash = $1
		AND map_id = any($2::text[])
		AND (length($5) = 0 OR process = $5)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $3 LIMIT $4
	`
	sqlFindSegmentsWithPrevLinkHashAndMapIDsAndTags = `
		SELECT l.link_hash, l.data, e.data FROM links l
		LEFT JOIN evidences e ON l.link_hash = e.link_hash
		WHERE prev_link_hash = $1
		AND map_id = any($2::text[]) AND tags @> $3
		AND (length($6) = 0 OR process = $6)
		ORDER BY priority DESC, l.created_at DESC
		OFFSET $4 LIMIT $5
	`
	sqlGetMapIDs = `
		SELECT l.map_id FROM links l
		WHERE (length($3) = 0 OR process = $3)
		GROUP BY l.map_id
		ORDER BY MAX(l.updated_at) DESC
		OFFSET $1 LIMIT $2
	`
	sqlSaveValue = `
		INSERT INTO values (
			key,
			value
		)
		VALUES ($1, $2)
		ON CONFLICT (key)
		DO UPDATE SET
			value = $2
	`
	sqlGetValue = `
		SELECT value FROM values
		WHERE key = $1
	`
	sqlDeleteValue = `
		DELETE FROM values
		WHERE key = $1
		RETURNING value
	`
	sqlGetEvidences = `
		SELECT data FROM evidences
		WHERE link_hash = $1
	`
	sqlAddEvidence = `
		INSERT INTO evidences (
			link_hash,
			provider,
			data
		)
		VALUES ($1, $2, $3)
		ON CONFLICT (link_hash, provider)
		DO NOTHING
	`
)

var sqlCreate = []string{
	`
		CREATE TABLE links (
			id BIGSERIAL PRIMARY KEY,
			link_hash bytea NOT NULL,
			priority double precision NOT NULL,
			map_id text NOT NULL,
			prev_link_hash bytea DEFAULT NULL,
			tags text[] DEFAULT NULL,
			data jsonb NOT NULL,
			process text NOT NULL,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
		)
	`,
	`
		CREATE UNIQUE INDEX links_link_hash_idx
		ON links (link_hash)
	`,
	`
		CREATE INDEX links_priority_created_at_idx
		ON links (priority DESC, created_at DESC)
	`,
	`
		CREATE INDEX links_map_id_idx
		ON links (map_id)
	`,
	`
		CREATE INDEX links_map_id_priority_created_at_idx
		ON links (map_id, priority DESC, created_at DESC)
	`,
	`
		CREATE INDEX links_prev_link_hash_priority_created_at_idx
		ON links (prev_link_hash, priority DESC, created_at DESC)
	`,
	`
		CREATE INDEX links_tags_idx
		ON links USING gin(tags)
	`,
	`
		CREATE TABLE evidences (
			id BIGSERIAL PRIMARY KEY,
			link_hash bytea NOT NULL,
			provider text NOT NULL,
			data jsonb NOT NULL,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
	)
	`,
	`
		CREATE UNIQUE INDEX evidences_link_hash_provider_idx
		ON evidences (link_hash, provider)
	`,
	`
		CREATE INDEX evidences_link_hash_idx
		ON evidences (link_hash)
	`,
	`
		CREATE TABLE values (
			id BIGSERIAL PRIMARY KEY,
			key bytea NOT NULL,
			value bytea NOT NULL,
			created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
		)
	`,
	`
		CREATE UNIQUE INDEX values_key_idx
		ON values (key)
	`,
}

var sqlDrop = []string{
	"DROP TABLE links, evidences, values",
}

type writeStmts struct {
	CreateLink  *sql.Stmt
	DeleteLink  *sql.Stmt
	SaveValue   *sql.Stmt
	DeleteValue *sql.Stmt
	AddEvidence *sql.Stmt
}

type readStmts struct {
	GetSegment   *sql.Stmt
	FindSegments *sql.Stmt
	GetMapIDs    *sql.Stmt
	GetValue     *sql.Stmt
	GetEvidences *sql.Stmt

	FindSegmentsWithMapIDs                       *sql.Stmt
	FindSegmentsWithPrevLinkHash                 *sql.Stmt
	FindSegmentsWithTags                         *sql.Stmt
	FindSegmentsWithLinkHashes                   *sql.Stmt
	FindSegmentsWithLinkHashesAndMapIDs          *sql.Stmt
	FindSegmentsWithLinkHashesAndTags            *sql.Stmt
	FindSegmentsWithLinkHashesAndMapIDsAndTags   *sql.Stmt
	FindSegmentsWithTagsAndLimit                 *sql.Stmt
	FindSegmentsWithMapIDsAndTags                *sql.Stmt
	FindSegmentsWithPrevLinkHashAndTags          *sql.Stmt
	FindSegmentsWithPrevLinkHashAndMapIDs        *sql.Stmt
	FindSegmentsWithPrevLinkHashAndMapIDsAndTags *sql.Stmt
}

type stmts struct {
	readStmts
	writeStmts
}

type batchStmts stmts

func newStmts(db *sql.DB) (*stmts, error) {
	var (
		s   stmts
		err error
	)

	prepare := func(str string) (stmt *sql.Stmt) {
		if err == nil {
			stmt, err = db.Prepare(str)
		}
		return
	}

	s.GetSegment = prepare(sqlGetSegment)
	s.FindSegments = prepare(sqlFindSegments)
	s.GetMapIDs = prepare(sqlGetMapIDs)
	s.GetValue = prepare(sqlGetValue)
	s.GetEvidences = prepare(sqlGetEvidences)

	s.FindSegmentsWithMapIDs = prepare(sqlFindSegmentsWithMapIDs)
	s.FindSegmentsWithPrevLinkHash = prepare(sqlFindSegmentsWithPrevLinkHash)
	s.FindSegmentsWithTags = prepare(sqlFindSegmentsWithTags)

	s.FindSegmentsWithLinkHashes = prepare(sqlFindSegmentsWithLinkHashes)
	s.FindSegmentsWithLinkHashesAndMapIDs = prepare(sqlFindSegmentsWithLinkHashesAndMapIDs)
	s.FindSegmentsWithLinkHashesAndTags = prepare(sqlFindSegmentsWithLinkHashesAndTags)
	s.FindSegmentsWithLinkHashesAndMapIDsAndTags = prepare(sqlFindSegmentsWithLinkHashesAndMapIDsAndTags)

	s.FindSegmentsWithMapIDsAndTags = prepare(sqlFindSegmentsWithMapIDsAndTags)
	s.FindSegmentsWithPrevLinkHashAndTags = prepare(sqlFindSegmentsWithPrevLinkHashAndTags)
	s.FindSegmentsWithPrevLinkHashAndMapIDs = prepare(sqlFindSegmentsWithPrevLinkHashAndMapIDs)
	s.FindSegmentsWithPrevLinkHashAndMapIDsAndTags = prepare(sqlFindSegmentsWithPrevLinkHashAndMapIDsAndTags)

	s.CreateLink = prepare(sqlCreateLink)
	s.DeleteLink = prepare(sqlDeleteLink)
	s.SaveValue = prepare(sqlSaveValue)
	s.DeleteValue = prepare(sqlDeleteValue)
	s.AddEvidence = prepare(sqlAddEvidence)

	if err != nil {
		return nil, err
	}

	return &s, nil
}

func newBatchStmts(tx *sql.Tx) (*batchStmts, error) {
	var (
		s   batchStmts
		err error
	)

	prepare := func(str string) (stmt *sql.Stmt) {
		if err == nil {
			stmt, err = tx.Prepare(str)
		}
		return
	}

	s.GetSegment = prepare(sqlGetSegment)
	s.FindSegments = prepare(sqlFindSegments)
	s.GetMapIDs = prepare(sqlGetMapIDs)
	s.GetValue = prepare(sqlGetValue)

	s.CreateLink = prepare(sqlCreateLink)
	s.DeleteLink = prepare(sqlDeleteLink)
	s.SaveValue = prepare(sqlSaveValue)
	s.DeleteValue = prepare(sqlDeleteValue)

	if err != nil {
		return nil, err
	}

	return &s, nil
}
<|endoftext|>


<commit_before>package logtap

import (
	"encoding/json"
	"fmt"
	"log"
)

// Telemetry is the interface that represents a telemetry receiver.
type Telemetry interface {
	Value(value interface{}, name string)
	Count(value int, name string)
}

type discardTelemetry struct{}

// DiscardTelemetry discards telemetry data.
var DiscardTelemetry = &discardTelemetry{}

func (*discardTelemetry) Value(value interface{}, name string) {}
func (*discardTelemetry) Count(value int, name string)         {}

type logTelemetry struct{}

// LogTelemetry sends telemetry to the standard logger.
//
// Examples:
//
//     LogTelemetry.Count(33, "beans")  // prints `APP_METRIC {"stat": "beans", "count": 33}`
//     LogTelemetry.Value(36.6, "temp") // prints `APP_METRIC {"stat": "temp", "value": 36.6}`
var LogTelemetry = &logTelemetry{}

func printAppMetric(key string, value interface{}, name string) {
	var v, n []byte
	switch value.(type) {
	case float32, float64,
		int, int8, int16, int32, int64,
		uint, uint8, uint16, uint32, uint64:
		v, _ = json.Marshal(value)
	default:
		return
	}
	n, _ = json.Marshal(name)
	log.Print(fmt.Sprintf(`APP_METRIC {"stat": %s, "%s": %s}`, n, key, v))
}

func (*logTelemetry) Value(value interface{}, name string) {
	printAppMetric("value", value, name)
}

func (*logTelemetry) Count(value int, name string) {
	printAppMetric("count", value, name)
}
<commit_msg>Telemetry: Additional property pairs may be passed to Count/Value methods.<commit_after>package logtap

import (
	"encoding/json"
	"fmt"
	"log"
)

// Telemetry is the interface that represents a telemetry receiver.
type Telemetry interface {
	Value(value interface{}, name string, props ...interface{})
	Count(value int, name string, props ...interface{})
}

type discardTelemetry struct{}

// DiscardTelemetry discards telemetry data.
var DiscardTelemetry = &discardTelemetry{}

func (*discardTelemetry) Value(value interface{}, name string, props ...interface{}) {}
func (*discardTelemetry) Count(value int, name string, props ...interface{})         {}

type logTelemetry struct{}

// LogTelemetry sends telemetry to the standard logger.
//
// Examples:
//
//     LogTelemetry.Count(33, "beans")  // prints `APP_METRIC {"stat": "beans", "count": 33}`
//     LogTelemetry.Value(36.6, "temp") // prints `APP_METRIC {"stat": "temp", "value": 36.6}`
var LogTelemetry = &logTelemetry{}

func printAppMetric(key string, value interface{}, name string) {
	var v, n []byte
	switch value.(type) {
	case float32, float64,
		int, int8, int16, int32, int64,
		uint, uint8, uint16, uint32, uint64:
		v, _ = json.Marshal(value)
	default:
		return
	}
	n, _ = json.Marshal(name)
	log.Print(fmt.Sprintf(`APP_METRIC {"stat": %s, "%s": %s}`, n, key, v))
}

func (*logTelemetry) Value(value interface{}, name string, props ...interface{}) {
	printAppMetric("value", value, name)
}

func (*logTelemetry) Count(value int, name string, props ...interface{}) {
	printAppMetric("count", value, name)
}
<|endoftext|>


<commit_before>/* Copyright 2017 Google Inc.

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

    https://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 gnmi

import (
	"errors"
	"fmt"
	"reflect"
	"sort"

	"github.com/openconfig/goyang/pkg/yang"
	"github.com/openconfig/ygot/experimental/ygotutils"
	"github.com/openconfig/ygot/ygot"

	pb "github.com/openconfig/gnmi/proto/gnmi"
	cpb "google.golang.org/genproto/googleapis/rpc/code"
)

// JSONUnmarshaler is the signature of the Unmarshal() function in the GoStruct code generated by openconfig ygot library.
type JSONUnmarshaler func([]byte, ygot.GoStruct) error

// GoStructEnumData is the data type to maintain GoStruct enum type.
type GoStructEnumData map[string]map[int64]ygot.EnumDefinition

// Model contains the model data and GoStruct information for the device to config.
type Model struct {
	modelData       []*pb.ModelData
	structRootType  reflect.Type
	schemaTreeRoot  *yang.Entry
	jsonUnmarshaler JSONUnmarshaler
	enumData        GoStructEnumData
}

// NewModel returns an instance of Model struct.
func NewModel(m []*pb.ModelData, t reflect.Type, r *yang.Entry, f JSONUnmarshaler, e GoStructEnumData) *Model {
	return &Model{
		modelData:       m,
		structRootType:  t,
		schemaTreeRoot:  r,
		jsonUnmarshaler: f,
		enumData:        e,
	}
}

// NewConfigStruct creates a ValidatedGoStruct of this model from jsonConfig. If jsonConfig is nil, creates an empty GoStruct.
func (m *Model) NewConfigStruct(jsonConfig []byte) (ygot.ValidatedGoStruct, error) {
	rootNode, stat := ygotutils.NewNode(m.structRootType, &pb.Path{})
	if stat.GetCode() != int32(cpb.Code_OK) {
		return nil, fmt.Errorf("cannot create root node: %v", stat)
	}

	rootStruct, ok := rootNode.(ygot.ValidatedGoStruct)
	if !ok {
		return nil, errors.New("root node is not a ygot.ValidatedGoStruct")
	}
	if jsonConfig != nil {
		if err := m.jsonUnmarshaler(jsonConfig, rootStruct); err != nil {
			return nil, err
		}
		if err := rootStruct.Validate(); err != nil {
			return nil, err
		}
	}
	return rootStruct, nil
}

// SupportedModels returns a list of supported models.
func (m *Model) SupportedModels() []string {
	mDesc := make([]string, len(m.modelData))
	for i, m := range m.modelData {
		mDesc[i] = fmt.Sprintf("%s %s", m.Name, m.Version)
	}
	sort.Strings(mDesc)
	return mDesc
}
<commit_msg>Update type definition of an unmarshal function to latest ygot.<commit_after>/* Copyright 2017 Google Inc.

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

    https://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 gnmi

import (
	"errors"
	"fmt"
	"reflect"
	"sort"

	"github.com/openconfig/goyang/pkg/yang"
	"github.com/openconfig/ygot/experimental/ygotutils"
	"github.com/openconfig/ygot/ygot"

	pb "github.com/openconfig/gnmi/proto/gnmi"
	cpb "google.golang.org/genproto/googleapis/rpc/code"
)

// JSONUnmarshaler is the signature of the Unmarshal() function in the GoStruct code generated by openconfig ygot library.
type JSONUnmarshaler func([]byte, ygot.GoStruct, ...ytypes.UnmarshalOpt) error

// GoStructEnumData is the data type to maintain GoStruct enum type.
type GoStructEnumData map[string]map[int64]ygot.EnumDefinition

// Model contains the model data and GoStruct information for the device to config.
type Model struct {
	modelData       []*pb.ModelData
	structRootType  reflect.Type
	schemaTreeRoot  *yang.Entry
	jsonUnmarshaler JSONUnmarshaler
	enumData        GoStructEnumData
}

// NewModel returns an instance of Model struct.
func NewModel(m []*pb.ModelData, t reflect.Type, r *yang.Entry, f JSONUnmarshaler, e GoStructEnumData) *Model {
	return &Model{
		modelData:       m,
		structRootType:  t,
		schemaTreeRoot:  r,
		jsonUnmarshaler: f,
		enumData:        e,
	}
}

// NewConfigStruct creates a ValidatedGoStruct of this model from jsonConfig. If jsonConfig is nil, creates an empty GoStruct.
func (m *Model) NewConfigStruct(jsonConfig []byte) (ygot.ValidatedGoStruct, error) {
	rootNode, stat := ygotutils.NewNode(m.structRootType, &pb.Path{})
	if stat.GetCode() != int32(cpb.Code_OK) {
		return nil, fmt.Errorf("cannot create root node: %v", stat)
	}

	rootStruct, ok := rootNode.(ygot.ValidatedGoStruct)
	if !ok {
		return nil, errors.New("root node is not a ygot.ValidatedGoStruct")
	}
	if jsonConfig != nil {
		if err := m.jsonUnmarshaler(jsonConfig, rootStruct); err != nil {
			return nil, err
		}
		if err := rootStruct.Validate(); err != nil {
			return nil, err
		}
	}
	return rootStruct, nil
}

// SupportedModels returns a list of supported models.
func (m *Model) SupportedModels() []string {
	mDesc := make([]string, len(m.modelData))
	for i, m := range m.modelData {
		mDesc[i] = fmt.Sprintf("%s %s", m.Name, m.Version)
	}
	sort.Strings(mDesc)
	return mDesc
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 datastore

import (
	"context"
	"testing"
	"time"

	"github.com/google/uuid"
	pb "github.com/googleforgames/open-saves/api"
	m "github.com/googleforgames/open-saves/internal/pkg/metadb"
	"github.com/googleforgames/open-saves/internal/pkg/metadb/metadbtest"
	"github.com/stretchr/testify/assert"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func newDriver(ctx context.Context, t *testing.T) *Driver {
	driver, err := NewDriver(ctx, "triton-for-games-dev")
	if err != nil {
		t.Fatalf("Initializing Datastore driver: %v", err)
	}
	driver.Namespace = "datastore-unittests"
	if err := driver.Connect(ctx); err != nil {
		t.Fatalf("Failed to connect to Datastore: %v", err)
	}
	t.Cleanup(func() { driver.Disconnect(ctx) })
	return driver
}

func newStoreKey() string {
	return "unittest_store_" + uuid.New().String()
}

func newRecordKey() string {
	return "unittest_record_" + uuid.New().String()
}

func cloneRecord(r *m.Record) *m.Record {
	if r == nil {
		return nil
	}
	ret := *r
	ret.Blob = append([]byte{}, r.Blob...)
	ret.Properties = make(m.PropertyMap)
	for k, v := range r.Properties {
		ret.Properties[k] = v
	}
	ret.Tags = append([]string{}, r.Tags...)
	return &ret
}

// setupTestStoreRecord creates a new store and inserts a record into it, then registers
// cleanup functions to delete these test store and record.
// Passing a nil to record will skip the record insertion.
func setupTestStoreRecord(ctx context.Context, t *testing.T, driver *Driver, store *m.Store, record *m.Record) (*m.Store, *m.Record) {
	newStore, err := driver.CreateStore(ctx, store)
	if err != nil {
		t.Fatalf("Could not create a new store: %v", err)
	}
	t.Cleanup(func() {
		driver.DeleteStore(ctx, newStore.Key)
		assert.NoError(t, driver.DeleteStore(ctx, newStore.Key))
	})
	metadbtest.AssertEqualStore(t, store, newStore, "CreateStore should return the created store.")
	var newRecord *m.Record
	if record != nil {
		newRecord, err = driver.InsertRecord(ctx, newStore.Key, record)
		if err != nil {
			t.Fatalf("Could not create a new record: %v", err)
		}
		t.Cleanup(func() {
			assert.NoError(t, driver.DeleteRecord(ctx, newStore.Key, newRecord.Key))
		})
		metadbtest.AssertEqualRecord(t, record, newRecord, "GetRecord should return the exact same record.")
	}
	return newStore, newRecord
}

func TestDriver_ConnectDisconnect(t *testing.T) {
	ctx := context.Background()
	// Connect() is tested inside newDriver().
	driver := newDriver(ctx, t)
	if err := driver.Disconnect(ctx); err != nil {
		t.Fatalf("Failed to disconnect: %v", err)
	}
}

func TestDriver_SimpleCreateGetDeleteStore(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	storeKey := newStoreKey()
	storeName := "SimpleCreateGetDeleteStore" + uuid.New().String()
	createdAt := time.Date(1988, 4, 16, 8, 6, 5, int(1234*time.Microsecond), time.UTC)
	store := &m.Store{
		Key:     storeKey,
		Name:    storeName,
		OwnerID: "owner",
		Tags:    []string{"abc", "def"},
		Timestamps: m.Timestamps{
			CreatedAt: createdAt,
			UpdatedAt: createdAt,
			Signature: uuid.MustParse("db94be80-e036-4ca8-a9c0-2259b8a67acc"),
		},
	}
	createdStore, err := driver.CreateStore(ctx, store)
	if err != nil {
		t.Fatalf("Could not create a new store: %v", err)
	}
	metadbtest.AssertEqualStore(t, store, createdStore, "CreateStore should return the created store.")

	store2, err := driver.GetStore(ctx, storeKey)
	if err != nil {
		t.Fatalf("Could not get store (%s): %v", storeKey, err)
	}
	metadbtest.AssertEqualStore(t, store, store2, "GetStore should return the exact same store.")

	store3, err := driver.FindStoreByName(ctx, storeName)
	if err != nil {
		t.Fatalf("Could not fetch store by name (%s): %v", storeName, err)
	}
	metadbtest.AssertEqualStore(t, store, store3, "FindStoreByName should return the exact same store.")

	err = driver.DeleteStore(ctx, storeKey)
	if err != nil {
		t.Fatalf("Failed to delete a store (%s): %v", storeKey, err)
	}
	store4, err := driver.GetStore(ctx, storeKey)
	if err == nil {
		t.Fatalf("GetStore didn't return an error after deleting a store: %v", err)
	}
	assert.Nil(t, store4, "GetStore should return a nil store.")
}

func TestDriver_SimpleCreateGetDeleteRecord(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	storeKey := newStoreKey()
	store := &m.Store{
		Key:     storeKey,
		Name:    "SimpleCreateGetDeleteRecord",
		OwnerID: "owner",
	}
	setupTestStoreRecord(ctx, t, driver, store, nil)

	recordKey := newRecordKey()
	blob := []byte{0x54, 0x72, 0x69, 0x74, 0x6f, 0x6e}
	createdAt := time.Date(1988, 4, 16, 8, 6, 5, int(1234*time.Microsecond), time.UTC)
	record := &m.Record{
		Key:      recordKey,
		Blob:     blob,
		BlobSize: int64(len(blob)),
		OwnerID:  "record owner",
		Tags:     []string{"abc", "def"},
		Properties: m.PropertyMap{
			"BoolTP":   {Type: pb.Property_BOOLEAN, BooleanValue: false},
			"IntTP":    {Type: pb.Property_INTEGER, IntegerValue: 42},
			"StringTP": {Type: pb.Property_STRING, StringValue: "a string value"},
		},
		Timestamps: m.Timestamps{
			CreatedAt: createdAt,
			UpdatedAt: createdAt,
			Signature: uuid.MustParse("89223949-0414-438e-8f5e-3fd9e2d11c1e"),
		},
	}
	expected := cloneRecord(record)

	createdRecord, err := driver.InsertRecord(ctx, storeKey, record)
	if err != nil {
		t.Fatalf("Failed to create a new record (%s) in store (%s): %v", recordKey, storeKey, err)
	}
	metadbtest.AssertEqualRecord(t, expected, createdRecord, "InsertRecord should return the created record.")

	record2, err := driver.GetRecord(ctx, storeKey, recordKey)
	if err != nil {
		t.Fatalf("Failed to get a record (%s) in store (%s): %v", recordKey, storeKey, err)
	}
	metadbtest.AssertEqualRecord(t, expected, record2, "GetRecord should return the exact same record.")

	record3, err := driver.InsertRecord(ctx, storeKey, record)
	assert.NotNil(t, err, "Insert should fail if a record with the same already exists.")
	assert.Nil(t, record3, "Insert should fail and return nil.")

	err = driver.DeleteRecord(ctx, storeKey, recordKey)
	assert.Nilf(t, err, "Failed to delete a record (%s) in store (%s): %v", recordKey, storeKey, err)

	record4, err := driver.GetRecord(ctx, storeKey, recordKey)
	assert.NotNilf(t, err, "GetRecord didn't return an error after deleting a record (%s)", recordKey)
	assert.Nil(t, record4, "GetRecord should return a nil with error")

	// The test store is deleted inside the cleanup function.
}

func TestDriver_InsertRecordShouldFailWithNonExistentStore(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	storeKey := newStoreKey()
	record := &m.Record{
		Key: newRecordKey(),
	}
	record1, err := driver.InsertRecord(ctx, storeKey, record)
	if err == nil {
		t.Error("InsertRecord should fail if the store doesn't exist.")
		driver.DeleteRecord(ctx, storeKey, record.Key)
	} else {
		assert.Equalf(t, codes.FailedPrecondition, status.Code(err),
			"InsertRecord should return FailedPrecondition if the store doesn't exist: %v", err)
	}
	assert.Nil(t, record1, "InsertRecord should return nil with error")
}

func TestDriver_DeleteStoreShouldFailWhenNotEmpty(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	store := &m.Store{Key: newStoreKey()}
	record := &m.Record{Key: newRecordKey(), Properties: make(m.PropertyMap)}

	setupTestStoreRecord(ctx, t, driver, store, record)

	if err := driver.DeleteStore(ctx, store.Key); err == nil {
		t.Error("DeleteStore should fail if the store is not empty.")
	} else {
		assert.Equalf(t, codes.FailedPrecondition, status.Code(err),
			"DeleteStore should return FailedPrecondition if the store is not empty: %v", err)
	}
}

func TestDriver_UpdateRecord(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)

	storeKey := newStoreKey()
	store := &m.Store{
		Key:     storeKey,
		Name:    "UpdateRecord",
		OwnerID: "owner",
	}

	recordKey := newRecordKey()
	blob := []byte{0x54, 0x72, 0x69, 0x74, 0x6f, 0x6e}
	createdAt := time.Date(1988, 4, 16, 8, 6, 5, int(1234*time.Microsecond), time.UTC)
	record := &m.Record{
		Key:        recordKey,
		Blob:       blob,
		BlobSize:   int64(len(blob)),
		Properties: make(m.PropertyMap),
		OwnerID:    "record owner",
		Tags:       []string{"abc", "def"},
		Timestamps: m.Timestamps{
			CreatedAt: createdAt,
			UpdatedAt: createdAt,
			Signature: uuid.MustParse("e4a677f6-8f1c-4765-be45-11b6400cc43b"),
		},
	}
	expected := cloneRecord(record)

	updatedRecord, err := driver.UpdateRecord(ctx, storeKey, record)
	assert.NotNil(t, err, "UpdateRecord should return an error if the specified record doesn't exist.")
	assert.Nil(t, updatedRecord, "UpdateRecord should return nil with error")

	setupTestStoreRecord(ctx, t, driver, store, record)

	record.Tags = append(record.Tags, "ghi")
	expected.Tags = append(expected.Tags, "ghi")
	record.OwnerID = "NewOwner"
	expected.OwnerID = record.OwnerID
	record.Timestamps.UpdatedAt = time.Date(1988, 5, 17, 5, 6, 8, int(4321*time.Microsecond), time.UTC)
	record.Timestamps.Signature = uuid.MustParse("3c1dc762-8f22-4d85-b729-b90393f45ca6")
	expected.Timestamps = record.Timestamps

	// Make sure UpdateRecord doesn't update CreatedAt
	record.Timestamps.CreatedAt = time.Unix(0, 0)
	updatedRecord, err = driver.UpdateRecord(ctx, storeKey, record)
	assert.Nilf(t, err, "Failed to update a record (%s) in store (%s): %v", recordKey, storeKey, err)
	metadbtest.AssertEqualRecord(t, expected, updatedRecord, "Updated record is not the same as original.")

	record2, err := driver.GetRecord(ctx, storeKey, recordKey)
	if err != nil {
		t.Fatalf("Failed to get a record (%s) in store (%s): %v", recordKey, storeKey, err)
	}
	metadbtest.AssertEqualRecord(t, expected, record2, "GetRecord should fetch the updated record.")
}

func TestDriver_DeleteShouldNotFailWithNonExistentKey(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	storeKey := newStoreKey()
	if err := driver.DeleteStore(ctx, storeKey); err != nil {
		t.Fatalf("DeleteStore failed with a non-existent key: %v", err)
	}
	recordKey := newRecordKey()
	if err := driver.DeleteRecord(ctx, storeKey, recordKey); err != nil {
		t.Fatalf("DeleteRecord failed with a non-existent key: %v", err)
	}
}

func TestDriver_TimestampPrecision(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	assert.Equal(t, timestampPrecision, driver.TimestampPrecision())
}
<commit_msg>fix: Change test keys to uuid + suffix (#152)<commit_after>// Copyright 2020 Google LLC
//
// 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 datastore

import (
	"context"
	"testing"
	"time"

	"github.com/google/uuid"
	pb "github.com/googleforgames/open-saves/api"
	m "github.com/googleforgames/open-saves/internal/pkg/metadb"
	"github.com/googleforgames/open-saves/internal/pkg/metadb/metadbtest"
	"github.com/stretchr/testify/assert"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func newDriver(ctx context.Context, t *testing.T) *Driver {
	driver, err := NewDriver(ctx, "triton-for-games-dev")
	if err != nil {
		t.Fatalf("Initializing Datastore driver: %v", err)
	}
	driver.Namespace = "datastore-unittests"
	if err := driver.Connect(ctx); err != nil {
		t.Fatalf("Failed to connect to Datastore: %v", err)
	}
	t.Cleanup(func() { driver.Disconnect(ctx) })
	return driver
}

func newStoreKey() string {
	return uuid.New().String() + "_unittest_store"
}

func newRecordKey() string {
	return uuid.New().String() + "_unittest_record"
}

func cloneRecord(r *m.Record) *m.Record {
	if r == nil {
		return nil
	}
	ret := *r
	ret.Blob = append([]byte{}, r.Blob...)
	ret.Properties = make(m.PropertyMap)
	for k, v := range r.Properties {
		ret.Properties[k] = v
	}
	ret.Tags = append([]string{}, r.Tags...)
	return &ret
}

// setupTestStoreRecord creates a new store and inserts a record into it, then registers
// cleanup functions to delete these test store and record.
// Passing a nil to record will skip the record insertion.
func setupTestStoreRecord(ctx context.Context, t *testing.T, driver *Driver, store *m.Store, record *m.Record) (*m.Store, *m.Record) {
	newStore, err := driver.CreateStore(ctx, store)
	if err != nil {
		t.Fatalf("Could not create a new store: %v", err)
	}
	t.Cleanup(func() {
		driver.DeleteStore(ctx, newStore.Key)
		assert.NoError(t, driver.DeleteStore(ctx, newStore.Key))
	})
	metadbtest.AssertEqualStore(t, store, newStore, "CreateStore should return the created store.")
	var newRecord *m.Record
	if record != nil {
		newRecord, err = driver.InsertRecord(ctx, newStore.Key, record)
		if err != nil {
			t.Fatalf("Could not create a new record: %v", err)
		}
		t.Cleanup(func() {
			assert.NoError(t, driver.DeleteRecord(ctx, newStore.Key, newRecord.Key))
		})
		metadbtest.AssertEqualRecord(t, record, newRecord, "GetRecord should return the exact same record.")
	}
	return newStore, newRecord
}

func TestDriver_ConnectDisconnect(t *testing.T) {
	ctx := context.Background()
	// Connect() is tested inside newDriver().
	driver := newDriver(ctx, t)
	if err := driver.Disconnect(ctx); err != nil {
		t.Fatalf("Failed to disconnect: %v", err)
	}
}

func TestDriver_SimpleCreateGetDeleteStore(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	storeKey := newStoreKey()
	storeName := "SimpleCreateGetDeleteStore" + uuid.New().String()
	createdAt := time.Date(1988, 4, 16, 8, 6, 5, int(1234*time.Microsecond), time.UTC)
	store := &m.Store{
		Key:     storeKey,
		Name:    storeName,
		OwnerID: "owner",
		Tags:    []string{"abc", "def"},
		Timestamps: m.Timestamps{
			CreatedAt: createdAt,
			UpdatedAt: createdAt,
			Signature: uuid.MustParse("db94be80-e036-4ca8-a9c0-2259b8a67acc"),
		},
	}
	createdStore, err := driver.CreateStore(ctx, store)
	if err != nil {
		t.Fatalf("Could not create a new store: %v", err)
	}
	metadbtest.AssertEqualStore(t, store, createdStore, "CreateStore should return the created store.")

	store2, err := driver.GetStore(ctx, storeKey)
	if err != nil {
		t.Fatalf("Could not get store (%s): %v", storeKey, err)
	}
	metadbtest.AssertEqualStore(t, store, store2, "GetStore should return the exact same store.")

	store3, err := driver.FindStoreByName(ctx, storeName)
	if err != nil {
		t.Fatalf("Could not fetch store by name (%s): %v", storeName, err)
	}
	metadbtest.AssertEqualStore(t, store, store3, "FindStoreByName should return the exact same store.")

	err = driver.DeleteStore(ctx, storeKey)
	if err != nil {
		t.Fatalf("Failed to delete a store (%s): %v", storeKey, err)
	}
	store4, err := driver.GetStore(ctx, storeKey)
	if err == nil {
		t.Fatalf("GetStore didn't return an error after deleting a store: %v", err)
	}
	assert.Nil(t, store4, "GetStore should return a nil store.")
}

func TestDriver_SimpleCreateGetDeleteRecord(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	storeKey := newStoreKey()
	store := &m.Store{
		Key:     storeKey,
		Name:    "SimpleCreateGetDeleteRecord",
		OwnerID: "owner",
	}
	setupTestStoreRecord(ctx, t, driver, store, nil)

	recordKey := newRecordKey()
	blob := []byte{0x54, 0x72, 0x69, 0x74, 0x6f, 0x6e}
	createdAt := time.Date(1988, 4, 16, 8, 6, 5, int(1234*time.Microsecond), time.UTC)
	record := &m.Record{
		Key:      recordKey,
		Blob:     blob,
		BlobSize: int64(len(blob)),
		OwnerID:  "record owner",
		Tags:     []string{"abc", "def"},
		Properties: m.PropertyMap{
			"BoolTP":   {Type: pb.Property_BOOLEAN, BooleanValue: false},
			"IntTP":    {Type: pb.Property_INTEGER, IntegerValue: 42},
			"StringTP": {Type: pb.Property_STRING, StringValue: "a string value"},
		},
		Timestamps: m.Timestamps{
			CreatedAt: createdAt,
			UpdatedAt: createdAt,
			Signature: uuid.MustParse("89223949-0414-438e-8f5e-3fd9e2d11c1e"),
		},
	}
	expected := cloneRecord(record)

	createdRecord, err := driver.InsertRecord(ctx, storeKey, record)
	if err != nil {
		t.Fatalf("Failed to create a new record (%s) in store (%s): %v", recordKey, storeKey, err)
	}
	metadbtest.AssertEqualRecord(t, expected, createdRecord, "InsertRecord should return the created record.")

	record2, err := driver.GetRecord(ctx, storeKey, recordKey)
	if err != nil {
		t.Fatalf("Failed to get a record (%s) in store (%s): %v", recordKey, storeKey, err)
	}
	metadbtest.AssertEqualRecord(t, expected, record2, "GetRecord should return the exact same record.")

	record3, err := driver.InsertRecord(ctx, storeKey, record)
	assert.NotNil(t, err, "Insert should fail if a record with the same already exists.")
	assert.Nil(t, record3, "Insert should fail and return nil.")

	err = driver.DeleteRecord(ctx, storeKey, recordKey)
	assert.Nilf(t, err, "Failed to delete a record (%s) in store (%s): %v", recordKey, storeKey, err)

	record4, err := driver.GetRecord(ctx, storeKey, recordKey)
	assert.NotNilf(t, err, "GetRecord didn't return an error after deleting a record (%s)", recordKey)
	assert.Nil(t, record4, "GetRecord should return a nil with error")

	// The test store is deleted inside the cleanup function.
}

func TestDriver_InsertRecordShouldFailWithNonExistentStore(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	storeKey := newStoreKey()
	record := &m.Record{
		Key: newRecordKey(),
	}
	record1, err := driver.InsertRecord(ctx, storeKey, record)
	if err == nil {
		t.Error("InsertRecord should fail if the store doesn't exist.")
		driver.DeleteRecord(ctx, storeKey, record.Key)
	} else {
		assert.Equalf(t, codes.FailedPrecondition, status.Code(err),
			"InsertRecord should return FailedPrecondition if the store doesn't exist: %v", err)
	}
	assert.Nil(t, record1, "InsertRecord should return nil with error")
}

func TestDriver_DeleteStoreShouldFailWhenNotEmpty(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	store := &m.Store{Key: newStoreKey()}
	record := &m.Record{Key: newRecordKey(), Properties: make(m.PropertyMap)}

	setupTestStoreRecord(ctx, t, driver, store, record)

	if err := driver.DeleteStore(ctx, store.Key); err == nil {
		t.Error("DeleteStore should fail if the store is not empty.")
	} else {
		assert.Equalf(t, codes.FailedPrecondition, status.Code(err),
			"DeleteStore should return FailedPrecondition if the store is not empty: %v", err)
	}
}

func TestDriver_UpdateRecord(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)

	storeKey := newStoreKey()
	store := &m.Store{
		Key:     storeKey,
		Name:    "UpdateRecord",
		OwnerID: "owner",
	}

	recordKey := newRecordKey()
	blob := []byte{0x54, 0x72, 0x69, 0x74, 0x6f, 0x6e}
	createdAt := time.Date(1988, 4, 16, 8, 6, 5, int(1234*time.Microsecond), time.UTC)
	record := &m.Record{
		Key:        recordKey,
		Blob:       blob,
		BlobSize:   int64(len(blob)),
		Properties: make(m.PropertyMap),
		OwnerID:    "record owner",
		Tags:       []string{"abc", "def"},
		Timestamps: m.Timestamps{
			CreatedAt: createdAt,
			UpdatedAt: createdAt,
			Signature: uuid.MustParse("e4a677f6-8f1c-4765-be45-11b6400cc43b"),
		},
	}
	expected := cloneRecord(record)

	updatedRecord, err := driver.UpdateRecord(ctx, storeKey, record)
	assert.NotNil(t, err, "UpdateRecord should return an error if the specified record doesn't exist.")
	assert.Nil(t, updatedRecord, "UpdateRecord should return nil with error")

	setupTestStoreRecord(ctx, t, driver, store, record)

	record.Tags = append(record.Tags, "ghi")
	expected.Tags = append(expected.Tags, "ghi")
	record.OwnerID = "NewOwner"
	expected.OwnerID = record.OwnerID
	record.Timestamps.UpdatedAt = time.Date(1988, 5, 17, 5, 6, 8, int(4321*time.Microsecond), time.UTC)
	record.Timestamps.Signature = uuid.MustParse("3c1dc762-8f22-4d85-b729-b90393f45ca6")
	expected.Timestamps = record.Timestamps

	// Make sure UpdateRecord doesn't update CreatedAt
	record.Timestamps.CreatedAt = time.Unix(0, 0)
	updatedRecord, err = driver.UpdateRecord(ctx, storeKey, record)
	assert.Nilf(t, err, "Failed to update a record (%s) in store (%s): %v", recordKey, storeKey, err)
	metadbtest.AssertEqualRecord(t, expected, updatedRecord, "Updated record is not the same as original.")

	record2, err := driver.GetRecord(ctx, storeKey, recordKey)
	if err != nil {
		t.Fatalf("Failed to get a record (%s) in store (%s): %v", recordKey, storeKey, err)
	}
	metadbtest.AssertEqualRecord(t, expected, record2, "GetRecord should fetch the updated record.")
}

func TestDriver_DeleteShouldNotFailWithNonExistentKey(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	storeKey := newStoreKey()
	if err := driver.DeleteStore(ctx, storeKey); err != nil {
		t.Fatalf("DeleteStore failed with a non-existent key: %v", err)
	}
	recordKey := newRecordKey()
	if err := driver.DeleteRecord(ctx, storeKey, recordKey); err != nil {
		t.Fatalf("DeleteRecord failed with a non-existent key: %v", err)
	}
}

func TestDriver_TimestampPrecision(t *testing.T) {
	ctx := context.Background()
	driver := newDriver(ctx, t)
	assert.Equal(t, timestampPrecision, driver.TimestampPrecision())
}
<|endoftext|>