<commit_before>package main

import (
	"fmt"
)

// #cgo CFLAGS: -I/usr/include/chewing
// #cgo LDFLAGS: -lchewing
// #include <chewing.h>
import "C"

type ChewingBenchmarkContext struct {
	accuracy []Accuracy
	ctx      *[0]byte
}

var BOPOMOFO_START = map[rune]uint8{
	'ㄅ': '1',
	'ㄆ': 'q',
	'ㄇ': 'a',
	'ㄈ': 'z',
	'ㄉ': '2',
	'ㄊ': 'w',
	'ㄋ': 's',
	'ㄌ': 'x',
	'ㄍ': 'e',
	'ㄎ': 'd',
	'ㄏ': 'c',
	'ㄐ': 'r',
	'ㄑ': 'f',
	'ㄒ': 'v',
	'ㄓ': '5',
	'ㄔ': 't',
	'ㄕ': 'g',
	'ㄖ': 'b',
	'ㄗ': 'y',
	'ㄘ': 'h',
	'ㄙ': 'n',
}

var BOPOMOFO_END = map[rune]uint8{
	'ㄧ': 'u',
	'ㄨ': 'j',
	'ㄩ': 'm',

	'ㄚ': '8',
	'ㄛ': 'i',
	'ㄜ': 'k',
	'ㄝ': ',',
	'ㄞ': '9',
	'ㄟ': 'o',
	'ㄠ': 'l',
	'ㄡ': '.',
	'ㄢ': '0',
	'ㄣ': 'p',
	'ㄤ': ';',
	'ㄥ': '/',
	'ㄦ': '-',
}

var BOPOMOFO_TONE = map[rune]uint8{
	'˙': '7',
	'ˊ': '6',
	'ˇ': '3',
	'ˋ': '4',
}

func newChewingBenchmarkItem() *ChewingBenchmarkContext {
	ctx := new(ChewingBenchmarkContext)

	ctx.ctx = C.chewing_new2(nil, nil, nil, nil)
	if ctx.ctx == nil {
		panic("chewing_new2 returns NULL")
	}

	return ctx
}

func (ctx *ChewingBenchmarkContext) deinit() {
	if ctx.ctx == nil {
		return
	}

	C.chewing_delete(ctx.ctx)
	ctx.ctx = nil
}

func (ctx *ChewingBenchmarkContext) getName() string {
	return "chewing"
}

func (ctx *ChewingBenchmarkContext) enterBenchmarkInput(input *BenchmarkInput) {
	if ctx.ctx == nil {
		return
	}

	C.chewing_clean_bopomofo_buf(ctx.ctx)
	C.chewing_clean_preedit_buf(ctx.ctx)

	for _, key := range bopomofoToKey(input.inputBopomofo) {
		C.chewing_handle_Default(ctx.ctx, C.int(key))
	}

	var accuracy Accuracy

	result := C.GoString(C.chewing_buffer_String_static(ctx.ctx))

	if len(result) != len(input.inputString) {
		panic("len(result) != len(input.inputString)")
	}

	for i := range result {
		if result[i] == input.inputString[i] {
			accuracy.correctCount++
		}
		accuracy.wordCount++
	}

	ctx.accuracy = append(ctx.accuracy, accuracy)
}

func (ctx *ChewingBenchmarkContext) getAccuracy() []Accuracy {
	return ctx.accuracy
}

func bopomofoToKey(bopomofo string) (keySequence []uint8) {
	terminated := true
	var key uint8
	var ok bool

	for _, runeValue := range bopomofo {
		key, ok = BOPOMOFO_START[runeValue]
		if ok {
			if !terminated {
				keySequence = append(keySequence, ' ')
				terminated = true
			}
			keySequence = append(keySequence, key)
			continue
		}

		key, ok = BOPOMOFO_END[runeValue]
		if ok {
			terminated = false
			keySequence = append(keySequence, key)
			continue
		}

		key, ok = BOPOMOFO_TONE[runeValue]
		if ok {
			terminated = true
			keySequence = append(keySequence, key)
			continue
		}

		panic(fmt.Sprintf("Unknown bopomofo: %c", runeValue))
	}

	if !terminated {
		keySequence = append(keySequence, ' ')
	}

	return keySequence
}
<commit_msg>Add enterBopomofo and computeAccuracy<commit_after>package main

import (
	"fmt"
)

// #cgo CFLAGS: -I/usr/include/chewing
// #cgo LDFLAGS: -lchewing
// #include <chewing.h>
import "C"

type ChewingBenchmarkContext struct {
	accuracy []Accuracy
	ctx      *[0]byte
}

var BOPOMOFO_START = map[rune]uint8{
	'ㄅ': '1',
	'ㄆ': 'q',
	'ㄇ': 'a',
	'ㄈ': 'z',
	'ㄉ': '2',
	'ㄊ': 'w',
	'ㄋ': 's',
	'ㄌ': 'x',
	'ㄍ': 'e',
	'ㄎ': 'd',
	'ㄏ': 'c',
	'ㄐ': 'r',
	'ㄑ': 'f',
	'ㄒ': 'v',
	'ㄓ': '5',
	'ㄔ': 't',
	'ㄕ': 'g',
	'ㄖ': 'b',
	'ㄗ': 'y',
	'ㄘ': 'h',
	'ㄙ': 'n',
}

var BOPOMOFO_END = map[rune]uint8{
	'ㄧ': 'u',
	'ㄨ': 'j',
	'ㄩ': 'm',

	'ㄚ': '8',
	'ㄛ': 'i',
	'ㄜ': 'k',
	'ㄝ': ',',
	'ㄞ': '9',
	'ㄟ': 'o',
	'ㄠ': 'l',
	'ㄡ': '.',
	'ㄢ': '0',
	'ㄣ': 'p',
	'ㄤ': ';',
	'ㄥ': '/',
	'ㄦ': '-',
}

var BOPOMOFO_TONE = map[rune]uint8{
	'˙': '7',
	'ˊ': '6',
	'ˇ': '3',
	'ˋ': '4',
}

func newChewingBenchmarkItem() *ChewingBenchmarkContext {
	ctx := new(ChewingBenchmarkContext)

	ctx.ctx = C.chewing_new2(nil, nil, nil, nil)
	if ctx.ctx == nil {
		panic("chewing_new2 returns NULL")
	}

	return ctx
}

func (ctx *ChewingBenchmarkContext) deinit() {
	if ctx.ctx == nil {
		return
	}

	C.chewing_delete(ctx.ctx)
	ctx.ctx = nil
}

func (ctx *ChewingBenchmarkContext) getName() string {
	return "chewing"
}

func (ctx *ChewingBenchmarkContext) enterBenchmarkInput(input *BenchmarkInput) {
	if ctx.ctx == nil {
		return
	}

	ctx.enterBopomofo(input)
	ctx.computeAccuracy(input)
}

func (ctx *ChewingBenchmarkContext) enterBopomofo(input *BenchmarkInput) {
	C.chewing_clean_bopomofo_buf(ctx.ctx)
	C.chewing_clean_preedit_buf(ctx.ctx)

	for _, key := range bopomofoToKey(input.inputBopomofo) {
		C.chewing_handle_Default(ctx.ctx, C.int(key))
	}
}

func (ctx *ChewingBenchmarkContext) computeAccuracy(input *BenchmarkInput) {
	var accuracy Accuracy

	result := C.GoString(C.chewing_buffer_String_static(ctx.ctx))

	if len(result) != len(input.inputString) {
		panic("len(result) != len(input.inputString)")
	}

	for i := range result {
		if result[i] == input.inputString[i] {
			accuracy.correctCount++
		}
		accuracy.wordCount++
	}

	ctx.accuracy = append(ctx.accuracy, accuracy)
}

func (ctx *ChewingBenchmarkContext) getAccuracy() []Accuracy {
	return ctx.accuracy
}

func bopomofoToKey(bopomofo string) (keySequence []uint8) {
	terminated := true
	var key uint8
	var ok bool

	for _, runeValue := range bopomofo {
		key, ok = BOPOMOFO_START[runeValue]
		if ok {
			if !terminated {
				keySequence = append(keySequence, ' ')
				terminated = true
			}
			keySequence = append(keySequence, key)
			continue
		}

		key, ok = BOPOMOFO_END[runeValue]
		if ok {
			terminated = false
			keySequence = append(keySequence, key)
			continue
		}

		key, ok = BOPOMOFO_TONE[runeValue]
		if ok {
			terminated = true
			keySequence = append(keySequence, key)
			continue
		}

		panic(fmt.Sprintf("Unknown bopomofo: %c", runeValue))
	}

	if !terminated {
		keySequence = append(keySequence, ' ')
	}

	return keySequence
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"bytes"
	"crypto/ecdsa"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"hash"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	"go.mozilla.org/autograph/signer/apk"
	"go.mozilla.org/autograph/signer/contentsignature"
	"go.mozilla.org/autograph/signer/mar"
	"go.mozilla.org/autograph/signer/xpi"
	"go.mozilla.org/hawk"
)

type signaturerequest struct {
	Input   string `json:"input"`
	KeyID   string `json:"keyid"`
	Options interface{}
}

type signatureresponse struct {
	Ref        string `json:"ref"`
	Type       string `json:"type"`
	SignerID   string `json:"signer_id"`
	PublicKey  string `json:"public_key,omitempty"`
	Signature  string `json:"signature"`
	SignedFile string `json:"signed_file"`
	X5U        string `json:"x5u,omitempty"`
}

type requestType int

const (
	requestTypeNone = iota
	requestTypeData
	requestTypeHash
	requestTypeFile
)

func urlToRequestType(url string) requestType {
	if strings.HasSuffix(url, "/sign/data") {
		return requestTypeData
	} else if strings.HasSuffix(url, "/sign/hash") {
		return requestTypeHash
	} else if strings.HasSuffix(url, "/sign/file") {
		return requestTypeFile
	} else {
		log.Fatalf("Unrecognized request type for url", url)
		return requestTypeNone
	}
}

func main() {
	var (
		userid, pass, data, hash, url, infile, outfile, keyid, cn string
		iter, maxworkers                                          int
		debug                                                     bool
		err                                                       error
		requests                                                  []signaturerequest
	)
	flag.Usage = func() {
		fmt.Print("autograph-client - simple command line client to the autograph service\n\n")
		flag.PrintDefaults()
		fmt.Print(`
examples:
* sign an APK, returns a signed APK
	$ go run client.go -f signed.apk -o test.apk -k testapp-android
	$ /opt/android-sdk/build-tools/27.0.3/apksigner verify -v test.apk
	Verifies
	Verified using v1 scheme (JAR signing): true
	Verified using v2 scheme (APK Signature Scheme v2): false
	Number of signers: 1

* issue a content signature on a hash, returns a CS header string:
	$ echo -en "Content-Signature:\0\0foo bar baz" | openssl dgst -sha256 -binary | openssl enc -base64
	$ go run client.go -a rniK3StMMdrWbXuJxVqEjALHR4cIp6mn3Coilj1kozk= -o testcs.txt
	$ cat testcs.txt
	keyid=appkey1;p384ecdsa=gf_X5JHv1KItwnpgGxmIdJ9KdjZ7EZMcleM-BTMGLnDuPpRvaGUdUDUg...

* sign an XPI, returns a PKCS7 detached signature:
	$ base64 -w 0 mozilla.sf
	U2lnbmF0dXJlLVZlcnNpb246IDEuMApNRDUt...
	$ go run client.go -d U2lnbmF0dXJlLVZlcnNpb2...  -cn cariboumaurice -k webextensions-rsa -o detachedxpisig.pkcs7

* sign an XPI file:
        $ go run client.go -f unsigned.xpi -cn cariboumaurice -k webextensions-rsa -o signed.xpi
`)
	}
	flag.StringVar(&userid, "u", "alice", "User ID")
	flag.StringVar(&pass, "p", "fs5wgcer9qj819kfptdlp8gm227ewxnzvsuj9ztycsx08hfhzu", "Secret passphrase")
	flag.StringVar(&data, "d", "base64(data)", "Base64 data to sign, will use the /sign/data endpoint")
	flag.StringVar(&hash, "a", "base64(sha256(data))", "Base64 hash to sign, will use the /sign/hash endpoint")
	flag.StringVar(&infile, "f", "/path/to/file", "Input file to sign, will use the /sign/file endpoint")
	flag.StringVar(&outfile, "o", ``, "Output file. If set, writes the signature or file to this location")
	flag.StringVar(&keyid, "k", ``, "Key ID to request a signature from a specific signer")
	flag.StringVar(&url, "t", `http://localhost:8000`, "target server, do not specific a URI or trailing slash")
	flag.IntVar(&iter, "i", 1, "number of signatures to request")
	flag.IntVar(&maxworkers, "m", 1, "maximum number of parallel workers")
	flag.StringVar(&cn, "cn", "", "when signing XPI, sets the CN to the add-on ID")
	flag.BoolVar(&debug, "D", false, "debug logs: show raw requests & responses")
	flag.Parse()

	if data != "base64(data)" {
		log.Printf("signing data %q", data)
		url = url + "/sign/data"
	} else if hash != "base64(sha256(data))" {
		log.Printf("signing hash %q", hash)
		url = url + "/sign/hash"
		data = hash
	} else if infile != "/path/to/file" {
		log.Printf("signing file %q", infile)
		url = url + "/sign/file"
		filebytes, err := ioutil.ReadFile(infile)
		if err != nil {
			log.Fatal(err)
		}
		data = base64.StdEncoding.EncodeToString(filebytes)
	}
	request := signaturerequest{
		Input: data,
		KeyID: keyid,
	}
	// if signing an xpi, the CN is set in the options
	if cn != "" {
		request.Options = xpi.Options{ID: cn}
	}
	requests = append(requests, request)
	reqBody, err := json.Marshal(requests)
	if err != nil {
		log.Fatal(err)
	}
	tr := &http.Transport{
		DisableKeepAlives: false,
	}
	cli := &http.Client{Transport: tr}

	workers := 0
	for i := 0; i < iter; i++ {
		for {
			if workers < maxworkers {
				break
			}
			time.Sleep(time.Second)
		}
		workers++
		go func() {
			// prepare the http request, with hawk token
			rdr := bytes.NewReader(reqBody)
			req, err := http.NewRequest("POST", url, rdr)
			if err != nil {
				log.Fatal(err)
			}
			req.Header.Set("Content-Type", "application/json")
			authheader := getAuthHeader(req, userid, pass, sha256.New, fmt.Sprintf("%d", time.Now().Nanosecond()), "application/json", reqBody)
			req.Header.Set("Authorization", authheader)
			if debug {
				fmt.Printf("DEBUG: sending request\nDEBUG: %+v\nDEBUG: %s\n", req, reqBody)
			}
			resp, err := cli.Do(req)
			if err != nil || resp == nil {
				log.Fatal(err)
			}
			if debug {
				fmt.Printf("DEBUG: received response\nDEBUG: %+v\n", resp)
			}
			defer resp.Body.Close()
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				log.Fatal(err)
			}
			if debug {
				fmt.Printf("DEBUG: %s\n", body)
			}
			if resp.StatusCode != http.StatusCreated {
				log.Fatalf("%s %s", resp.Status, body)
			}
			// verify that we got a proper signature response, with a valid signature
			var responses []signatureresponse
			err = json.Unmarshal(body, &responses)
			if err != nil {
				log.Fatal(err)
			}
			if len(requests) != len(responses) {
				log.Fatalf("sent %d signature requests and got %d responses, something's wrong", len(requests), len(responses))
			}
			reqType := urlToRequestType(url)
			for i, response := range responses {
				input, err := base64.StdEncoding.DecodeString(requests[i].Input)
				if err != nil {
					log.Fatal(err)
				}
				var (
					sigStatus bool
					sigData   []byte
				)
				switch response.Type {
				case contentsignature.Type:
					sigStatus = verifyContentSignature(input, response, req.URL.RequestURI())
					sig, err := contentsignature.Unmarshal(response.Signature)
					if err != nil {
						log.Fatal(err)
					}
					var sigStr string
					if response.X5U != "" {
						sigStr = "x5u=" + response.X5U + ";"
					} else {
						sigStr = "keyid=" + response.SignerID + ";"
					}
					sigStr += sig.Mode + "=" + response.Signature + "\n"
					sigData = []byte(sigStr)
				case xpi.Type:
					sigStatus = verifyXPI(input, request, response, reqType)
					switch reqType {
					case requestTypeData:
						sigData, err = base64.StdEncoding.DecodeString(response.Signature)
					case requestTypeFile:
						sigData, err = base64.StdEncoding.DecodeString(response.SignedFile)
					default:
						err = fmt.Errorf("Cannot decode signature data for request type %s", reqType)
					}
					if err != nil {
						log.Fatal(err)
					}
				case apk.Type:
					sigStatus = verifyAPK(input)
					sigData, err = base64.StdEncoding.DecodeString(response.SignedFile)
					if err != nil {
						log.Fatal(err)
					}
				case mar.Type:
					sigStatus = verifyMAR(input)
					sigData, err = base64.StdEncoding.DecodeString(response.SignedFile)
					if err != nil {
						log.Fatal(err)
					}
				default:
					log.Fatal("unsupported signature type", response.Type)
				}
				if sigStatus {
					log.Printf("signature %d from signer %q passes", i, response.SignerID)
				} else {
					log.Fatalf("response %d from signer %q does not pass!", i, response.SignerID)
				}
				if outfile != "" {
					err = ioutil.WriteFile(outfile, sigData, 0644)
					if err != nil {
						log.Fatal(err)
					}
					log.Println("response written to", outfile)
				}
			}
			workers--
		}()
	}
	for {
		if workers <= 0 {
			break
		}
		time.Sleep(time.Second)
	}
}

func getAuthHeader(req *http.Request, user, token string, hash func() hash.Hash, ext, contenttype string, payload []byte) string {
	auth := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   user,
			Key:  token,
			Hash: hash},
		0)
	auth.Ext = ext
	payloadhash := auth.PayloadHash(contenttype)
	payloadhash.Write(payload)
	auth.SetHash(payloadhash)
	return auth.RequestHeader()
}

// verify an ecdsa signature
func verifyContentSignature(input []byte, resp signatureresponse, endpoint string) bool {
	keyBytes, err := base64.StdEncoding.DecodeString(resp.PublicKey)
	if err != nil {
		log.Fatal(err)
	}
	keyInterface, err := x509.ParsePKIXPublicKey(keyBytes)
	if err != nil {
		log.Fatal(err)
	}
	pubKey := keyInterface.(*ecdsa.PublicKey)
	if endpoint == "/sign/data" {
		var templated []byte
		templated = make([]byte, len("Content-Signature:\x00")+len(input))
		copy(templated[:len("Content-Signature:\x00")], []byte("Content-Signature:\x00"))
		copy(templated[len("Content-Signature:\x00"):], input)

		var md hash.Hash
		switch pubKey.Params().Name {
		case "P-256":
			md = sha256.New()
		case "P-384":
			md = sha512.New384()
		case "P-521":
			md = sha512.New()
		default:
			log.Fatalf("unsupported curve algorithm %q", pubKey.Params().Name)
		}
		md.Write(templated)
		input = md.Sum(nil)
	}
	sig, err := contentsignature.Unmarshal(resp.Signature)
	if err != nil {
		log.Fatal(err)
	}
	return ecdsa.Verify(pubKey, input, sig.R, sig.S)
}

func verifyXPI(input []byte, req signaturerequest, resp signatureresponse, reqType requestType) bool {
	switch reqType {
	case requestTypeData:
		sig, err := xpi.Unmarshal(resp.Signature, input)
		if err != nil {
			log.Fatal(err)
		}
		err = sig.VerifyWithChain(nil)
		if err != nil {
			log.Fatal(err)
		}
		return true
	case requestTypeFile:
		signedFile, err := base64.StdEncoding.DecodeString(resp.SignedFile)
		if err != nil {
			log.Fatal(err)
		}
		err = xpi.VerifySignedFile(signedFile, nil)
		if err != nil {
			log.Fatal(err)
		}
		return true
	default:
		return false
	}
}

func verifyAPK(signedAPK []byte) bool {
	zipReader := bytes.NewReader(signedAPK)
	r, err := zip.NewReader(zipReader, int64(len(signedAPK)))
	if err != nil {
		log.Fatal(err)
	}
	var (
		sigstr  string
		sigdata []byte
	)
	for _, f := range r.File {
		switch f.Name {
		case "META-INF/SIGNATURE.SF":
			rc, err := f.Open()
			defer rc.Close()
			if err != nil {
				log.Fatal(err)
			}
			sigdata, err = ioutil.ReadAll(rc)
			if err != nil {
				log.Fatal(err)
			}
		case "META-INF/SIGNATURE.RSA":
			rc, err := f.Open()
			defer rc.Close()
			if err != nil {
				log.Fatal(err)
			}
			rawsig, err := ioutil.ReadAll(rc)
			if err != nil {
				log.Fatal(err)
			}
			sigstr = base64.StdEncoding.EncodeToString(rawsig)
		}
	}
	// convert string format back to signature
	sig, err := apk.Unmarshal(sigstr, sigdata)
	if err != nil {
		log.Fatalf("failed to unmarshal signature: %v", err)
	}
	// verify signature on input data
	if sig.Verify() != nil {
		log.Fatalf("failed to verify apk signature: %v", sig.Verify())
	}
	return true
}

func verifyMAR(signedMAR []byte) bool {
	return true
}
<commit_msg>autograph-client: verify apk after decoding the response<commit_after>package main

import (
	"archive/zip"
	"bytes"
	"crypto/ecdsa"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"hash"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	"go.mozilla.org/autograph/signer/apk"
	"go.mozilla.org/autograph/signer/contentsignature"
	"go.mozilla.org/autograph/signer/mar"
	"go.mozilla.org/autograph/signer/xpi"
	"go.mozilla.org/hawk"
)

type signaturerequest struct {
	Input   string `json:"input"`
	KeyID   string `json:"keyid"`
	Options interface{}
}

type signatureresponse struct {
	Ref        string `json:"ref"`
	Type       string `json:"type"`
	SignerID   string `json:"signer_id"`
	PublicKey  string `json:"public_key,omitempty"`
	Signature  string `json:"signature"`
	SignedFile string `json:"signed_file"`
	X5U        string `json:"x5u,omitempty"`
}

type requestType int

const (
	requestTypeNone = iota
	requestTypeData
	requestTypeHash
	requestTypeFile
)

func urlToRequestType(url string) requestType {
	if strings.HasSuffix(url, "/sign/data") {
		return requestTypeData
	} else if strings.HasSuffix(url, "/sign/hash") {
		return requestTypeHash
	} else if strings.HasSuffix(url, "/sign/file") {
		return requestTypeFile
	} else {
		log.Fatalf("Unrecognized request type for url", url)
		return requestTypeNone
	}
}

func main() {
	var (
		userid, pass, data, hash, url, infile, outfile, keyid, cn string
		iter, maxworkers                                          int
		debug                                                     bool
		err                                                       error
		requests                                                  []signaturerequest
	)
	flag.Usage = func() {
		fmt.Print("autograph-client - simple command line client to the autograph service\n\n")
		flag.PrintDefaults()
		fmt.Print(`
examples:
* sign an APK, returns a signed APK
	$ go run client.go -f signed.apk -o test.apk -k testapp-android
	$ /opt/android-sdk/build-tools/27.0.3/apksigner verify -v test.apk
	Verifies
	Verified using v1 scheme (JAR signing): true
	Verified using v2 scheme (APK Signature Scheme v2): false
	Number of signers: 1

* issue a content signature on a hash, returns a CS header string:
	$ echo -en "Content-Signature:\0\0foo bar baz" | openssl dgst -sha256 -binary | openssl enc -base64
	$ go run client.go -a rniK3StMMdrWbXuJxVqEjALHR4cIp6mn3Coilj1kozk= -o testcs.txt
	$ cat testcs.txt
	keyid=appkey1;p384ecdsa=gf_X5JHv1KItwnpgGxmIdJ9KdjZ7EZMcleM-BTMGLnDuPpRvaGUdUDUg...

* sign an XPI, returns a PKCS7 detached signature:
	$ base64 -w 0 mozilla.sf
	U2lnbmF0dXJlLVZlcnNpb246IDEuMApNRDUt...
	$ go run client.go -d U2lnbmF0dXJlLVZlcnNpb2...  -cn cariboumaurice -k webextensions-rsa -o detachedxpisig.pkcs7

* sign an XPI file:
        $ go run client.go -f unsigned.xpi -cn cariboumaurice -k webextensions-rsa -o signed.xpi
`)
	}
	flag.StringVar(&userid, "u", "alice", "User ID")
	flag.StringVar(&pass, "p", "fs5wgcer9qj819kfptdlp8gm227ewxnzvsuj9ztycsx08hfhzu", "Secret passphrase")
	flag.StringVar(&data, "d", "base64(data)", "Base64 data to sign, will use the /sign/data endpoint")
	flag.StringVar(&hash, "a", "base64(sha256(data))", "Base64 hash to sign, will use the /sign/hash endpoint")
	flag.StringVar(&infile, "f", "/path/to/file", "Input file to sign, will use the /sign/file endpoint")
	flag.StringVar(&outfile, "o", ``, "Output file. If set, writes the signature or file to this location")
	flag.StringVar(&keyid, "k", ``, "Key ID to request a signature from a specific signer")
	flag.StringVar(&url, "t", `http://localhost:8000`, "target server, do not specific a URI or trailing slash")
	flag.IntVar(&iter, "i", 1, "number of signatures to request")
	flag.IntVar(&maxworkers, "m", 1, "maximum number of parallel workers")
	flag.StringVar(&cn, "cn", "", "when signing XPI, sets the CN to the add-on ID")
	flag.BoolVar(&debug, "D", false, "debug logs: show raw requests & responses")
	flag.Parse()

	if data != "base64(data)" {
		log.Printf("signing data %q", data)
		url = url + "/sign/data"
	} else if hash != "base64(sha256(data))" {
		log.Printf("signing hash %q", hash)
		url = url + "/sign/hash"
		data = hash
	} else if infile != "/path/to/file" {
		log.Printf("signing file %q", infile)
		url = url + "/sign/file"
		filebytes, err := ioutil.ReadFile(infile)
		if err != nil {
			log.Fatal(err)
		}
		data = base64.StdEncoding.EncodeToString(filebytes)
	}
	request := signaturerequest{
		Input: data,
		KeyID: keyid,
	}
	// if signing an xpi, the CN is set in the options
	if cn != "" {
		request.Options = xpi.Options{ID: cn}
	}
	requests = append(requests, request)
	reqBody, err := json.Marshal(requests)
	if err != nil {
		log.Fatal(err)
	}
	tr := &http.Transport{
		DisableKeepAlives: false,
	}
	cli := &http.Client{Transport: tr}

	workers := 0
	for i := 0; i < iter; i++ {
		for {
			if workers < maxworkers {
				break
			}
			time.Sleep(time.Second)
		}
		workers++
		go func() {
			// prepare the http request, with hawk token
			rdr := bytes.NewReader(reqBody)
			req, err := http.NewRequest("POST", url, rdr)
			if err != nil {
				log.Fatal(err)
			}
			req.Header.Set("Content-Type", "application/json")
			authheader := getAuthHeader(req, userid, pass, sha256.New, fmt.Sprintf("%d", time.Now().Nanosecond()), "application/json", reqBody)
			req.Header.Set("Authorization", authheader)
			if debug {
				fmt.Printf("DEBUG: sending request\nDEBUG: %+v\nDEBUG: %s\n", req, reqBody)
			}
			resp, err := cli.Do(req)
			if err != nil || resp == nil {
				log.Fatal(err)
			}
			if debug {
				fmt.Printf("DEBUG: received response\nDEBUG: %+v\n", resp)
			}
			defer resp.Body.Close()
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				log.Fatal(err)
			}
			if debug {
				fmt.Printf("DEBUG: %s\n", body)
			}
			if resp.StatusCode != http.StatusCreated {
				log.Fatalf("%s %s", resp.Status, body)
			}
			// verify that we got a proper signature response, with a valid signature
			var responses []signatureresponse
			err = json.Unmarshal(body, &responses)
			if err != nil {
				log.Fatal(err)
			}
			if len(requests) != len(responses) {
				log.Fatalf("sent %d signature requests and got %d responses, something's wrong", len(requests), len(responses))
			}
			reqType := urlToRequestType(url)
			for i, response := range responses {
				input, err := base64.StdEncoding.DecodeString(requests[i].Input)
				if err != nil {
					log.Fatal(err)
				}
				var (
					sigStatus bool
					sigData   []byte
				)
				switch response.Type {
				case contentsignature.Type:
					sigStatus = verifyContentSignature(input, response, req.URL.RequestURI())
					sig, err := contentsignature.Unmarshal(response.Signature)
					if err != nil {
						log.Fatal(err)
					}
					var sigStr string
					if response.X5U != "" {
						sigStr = "x5u=" + response.X5U + ";"
					} else {
						sigStr = "keyid=" + response.SignerID + ";"
					}
					sigStr += sig.Mode + "=" + response.Signature + "\n"
					sigData = []byte(sigStr)
				case xpi.Type:
					sigStatus = verifyXPI(input, request, response, reqType)
					switch reqType {
					case requestTypeData:
						sigData, err = base64.StdEncoding.DecodeString(response.Signature)
					case requestTypeFile:
						sigData, err = base64.StdEncoding.DecodeString(response.SignedFile)
					default:
						err = fmt.Errorf("Cannot decode signature data for request type %s", reqType)
					}
					if err != nil {
						log.Fatal(err)
					}
				case apk.Type:
					sigData, err = base64.StdEncoding.DecodeString(response.SignedFile)
					if err != nil {
						log.Fatal(err)
					}
					sigStatus = verifyAPK(sigData)
				case mar.Type:
					sigStatus = verifyMAR(input)
					sigData, err = base64.StdEncoding.DecodeString(response.SignedFile)
					if err != nil {
						log.Fatal(err)
					}
				default:
					log.Fatal("unsupported signature type", response.Type)
				}
				if sigStatus {
					log.Printf("signature %d from signer %q passes", i, response.SignerID)
				} else {
					log.Fatalf("response %d from signer %q does not pass!", i, response.SignerID)
				}
				if outfile != "" {
					err = ioutil.WriteFile(outfile, sigData, 0644)
					if err != nil {
						log.Fatal(err)
					}
					log.Println("response written to", outfile)
				}
			}
			workers--
		}()
	}
	for {
		if workers <= 0 {
			break
		}
		time.Sleep(time.Second)
	}
}

func getAuthHeader(req *http.Request, user, token string, hash func() hash.Hash, ext, contenttype string, payload []byte) string {
	auth := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   user,
			Key:  token,
			Hash: hash},
		0)
	auth.Ext = ext
	payloadhash := auth.PayloadHash(contenttype)
	payloadhash.Write(payload)
	auth.SetHash(payloadhash)
	return auth.RequestHeader()
}

// verify an ecdsa signature
func verifyContentSignature(input []byte, resp signatureresponse, endpoint string) bool {
	keyBytes, err := base64.StdEncoding.DecodeString(resp.PublicKey)
	if err != nil {
		log.Fatal(err)
	}
	keyInterface, err := x509.ParsePKIXPublicKey(keyBytes)
	if err != nil {
		log.Fatal(err)
	}
	pubKey := keyInterface.(*ecdsa.PublicKey)
	if endpoint == "/sign/data" {
		var templated []byte
		templated = make([]byte, len("Content-Signature:\x00")+len(input))
		copy(templated[:len("Content-Signature:\x00")], []byte("Content-Signature:\x00"))
		copy(templated[len("Content-Signature:\x00"):], input)

		var md hash.Hash
		switch pubKey.Params().Name {
		case "P-256":
			md = sha256.New()
		case "P-384":
			md = sha512.New384()
		case "P-521":
			md = sha512.New()
		default:
			log.Fatalf("unsupported curve algorithm %q", pubKey.Params().Name)
		}
		md.Write(templated)
		input = md.Sum(nil)
	}
	sig, err := contentsignature.Unmarshal(resp.Signature)
	if err != nil {
		log.Fatal(err)
	}
	return ecdsa.Verify(pubKey, input, sig.R, sig.S)
}

func verifyXPI(input []byte, req signaturerequest, resp signatureresponse, reqType requestType) bool {
	switch reqType {
	case requestTypeData:
		sig, err := xpi.Unmarshal(resp.Signature, input)
		if err != nil {
			log.Fatal(err)
		}
		err = sig.VerifyWithChain(nil)
		if err != nil {
			log.Fatal(err)
		}
		return true
	case requestTypeFile:
		signedFile, err := base64.StdEncoding.DecodeString(resp.SignedFile)
		if err != nil {
			log.Fatal(err)
		}
		err = xpi.VerifySignedFile(signedFile, nil)
		if err != nil {
			log.Fatal(err)
		}
		return true
	default:
		return false
	}
}

func verifyAPK(signedAPK []byte) bool {
	zipReader := bytes.NewReader(signedAPK)
	r, err := zip.NewReader(zipReader, int64(len(signedAPK)))
	if err != nil {
		log.Fatal(err)
	}
	var (
		sigstr  string
		sigdata []byte
	)
	for _, f := range r.File {
		switch f.Name {
		case "META-INF/SIGNATURE.SF":
			rc, err := f.Open()
			defer rc.Close()
			if err != nil {
				log.Fatal(err)
			}
			sigdata, err = ioutil.ReadAll(rc)
			if err != nil {
				log.Fatal(err)
			}
		case "META-INF/SIGNATURE.RSA":
			rc, err := f.Open()
			defer rc.Close()
			if err != nil {
				log.Fatal(err)
			}
			rawsig, err := ioutil.ReadAll(rc)
			if err != nil {
				log.Fatal(err)
			}
			sigstr = base64.StdEncoding.EncodeToString(rawsig)
		}
	}
	// convert string format back to signature
	sig, err := apk.Unmarshal(sigstr, sigdata)
	if err != nil {
		log.Fatalf("failed to unmarshal signature: %v", err)
	}
	// verify signature on input data
	if sig.Verify() != nil {
		log.Fatalf("failed to verify apk signature: %v", sig.Verify())
	}
	return true
}

func verifyMAR(signedMAR []byte) bool {
	return true
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes 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 cache

import (
	"fmt"
	"sync"

	"k8s.io/apimachinery/pkg/util/sets"
)

// ThreadSafeStore is an interface that allows concurrent access to a storage backend.
// TL;DR caveats: you must not modify anything returned by Get or List as it will break
// the indexing feature in addition to not being thread safe.
//
// The guarantees of thread safety provided by List/Get are only valid if the caller
// treats returned items as read-only. For example, a pointer inserted in the store
// through `Add` will be returned as is by `Get`. Multiple clients might invoke `Get`
// on the same key and modify the pointer in a non-thread-safe way. Also note that
// modifying objects stored by the indexers (if any) will *not* automatically lead
// to a re-index. So it's not a good idea to directly modify the objects returned by
// Get/List, in general.
type ThreadSafeStore interface {
	Add(key string, obj interface{})
	Update(key string, obj interface{})
	Delete(key string)
	Get(key string) (item interface{}, exists bool)
	List() []interface{}
	ListKeys() []string
	Replace(map[string]interface{}, string)
	Index(indexName string, obj interface{}) ([]interface{}, error)
	IndexKeys(indexName, indexKey string) ([]string, error)
	ListIndexFuncValues(name string) []string
	ByIndex(indexName, indexKey string) ([]interface{}, error)
	GetIndexers() Indexers

	// AddIndexers adds more indexers to this store.  If you call this after you already have data
	// in the store, the results are undefined.
	AddIndexers(newIndexers Indexers) error
	Resync() error
}

// threadSafeMap implements ThreadSafeStore
type threadSafeMap struct {
	lock  sync.RWMutex
	items map[string]interface{}

	// indexers maps a name to an IndexFunc
	indexers Indexers
	// indices maps a name to an Index
	indices Indices
}

func (c *threadSafeMap) Add(key string, obj interface{}) {
	c.lock.Lock()
	defer c.lock.Unlock()
	oldObject := c.items[key]
	c.items[key] = obj
	c.updateIndices(oldObject, obj, key)
}

func (c *threadSafeMap) Update(key string, obj interface{}) {
	c.lock.Lock()
	defer c.lock.Unlock()
	oldObject := c.items[key]
	c.items[key] = obj
	c.updateIndices(oldObject, obj, key)
}

func (c *threadSafeMap) Delete(key string) {
	c.lock.Lock()
	defer c.lock.Unlock()
	if obj, exists := c.items[key]; exists {
		c.deleteFromIndices(obj, key)
		delete(c.items, key)
	}
}

func (c *threadSafeMap) Get(key string) (item interface{}, exists bool) {
	c.lock.RLock()
	defer c.lock.RUnlock()
	item, exists = c.items[key]
	return item, exists
}

func (c *threadSafeMap) List() []interface{} {
	c.lock.RLock()
	defer c.lock.RUnlock()
	list := make([]interface{}, 0, len(c.items))
	for _, item := range c.items {
		list = append(list, item)
	}
	return list
}

// ListKeys returns a list of all the keys of the objects currently
// in the threadSafeMap.
func (c *threadSafeMap) ListKeys() []string {
	c.lock.RLock()
	defer c.lock.RUnlock()
	list := make([]string, 0, len(c.items))
	for key := range c.items {
		list = append(list, key)
	}
	return list
}

func (c *threadSafeMap) Replace(items map[string]interface{}, resourceVersion string) {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.items = items

	// rebuild any index
	c.indices = Indices{}
	for key, item := range c.items {
		c.updateIndices(nil, item, key)
	}
}

// Index returns a list of items that match on the index function
// Index is thread-safe so long as you treat all items as immutable
func (c *threadSafeMap) Index(indexName string, obj interface{}) ([]interface{}, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()

	indexFunc := c.indexers[indexName]
	if indexFunc == nil {
		return nil, fmt.Errorf("Index with name %s does not exist", indexName)
	}

	indexKeys, err := indexFunc(obj)
	if err != nil {
		return nil, err
	}
	index := c.indices[indexName]

	var returnKeySet sets.String
	if len(indexKeys) == 1 {
		// In majority of cases, there is exactly one value matching.
		// Optimize the most common path - deduping is not needed here.
		returnKeySet = index[indexKeys[0]]
	} else {
		// Need to de-dupe the return list.
		// Since multiple keys are allowed, this can happen.
		returnKeySet = sets.String{}
		for _, indexKey := range indexKeys {
			for key := range index[indexKey] {
				returnKeySet.Insert(key)
			}
		}
	}

	list := make([]interface{}, 0, returnKeySet.Len())
	for absoluteKey := range returnKeySet {
		list = append(list, c.items[absoluteKey])
	}
	return list, nil
}

// ByIndex returns a list of items that match an exact value on the index function
func (c *threadSafeMap) ByIndex(indexName, indexKey string) ([]interface{}, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()

	indexFunc := c.indexers[indexName]
	if indexFunc == nil {
		return nil, fmt.Errorf("Index with name %s does not exist", indexName)
	}

	index := c.indices[indexName]

	set := index[indexKey]
	list := make([]interface{}, 0, set.Len())
	for _, key := range set.List() {
		list = append(list, c.items[key])
	}

	return list, nil
}

// IndexKeys returns a list of keys that match on the index function.
// IndexKeys is thread-safe so long as you treat all items as immutable.
func (c *threadSafeMap) IndexKeys(indexName, indexKey string) ([]string, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()

	indexFunc := c.indexers[indexName]
	if indexFunc == nil {
		return nil, fmt.Errorf("Index with name %s does not exist", indexName)
	}

	index := c.indices[indexName]

	set := index[indexKey]
	return set.List(), nil
}

func (c *threadSafeMap) ListIndexFuncValues(indexName string) []string {
	c.lock.RLock()
	defer c.lock.RUnlock()

	index := c.indices[indexName]
	names := make([]string, 0, len(index))
	for key := range index {
		names = append(names, key)
	}
	return names
}

func (c *threadSafeMap) GetIndexers() Indexers {
	return c.indexers
}

func (c *threadSafeMap) AddIndexers(newIndexers Indexers) error {
	c.lock.Lock()
	defer c.lock.Unlock()

	if len(c.items) > 0 {
		return fmt.Errorf("cannot add indexers to running index")
	}

	oldKeys := sets.StringKeySet(c.indexers)
	newKeys := sets.StringKeySet(newIndexers)

	if oldKeys.HasAny(newKeys.List()...) {
		return fmt.Errorf("indexer conflict: %v", oldKeys.Intersection(newKeys))
	}

	for k, v := range newIndexers {
		c.indexers[k] = v
	}
	return nil
}

// updateIndices modifies the objects location in the managed indexes, if this is an update, you must provide an oldObj
// updateIndices must be called from a function that already has a lock on the cache
func (c *threadSafeMap) updateIndices(oldObj interface{}, newObj interface{}, key string) {
	// if we got an old object, we need to remove it before we add it again
	if oldObj != nil {
		c.deleteFromIndices(oldObj, key)
	}
	for name, indexFunc := range c.indexers {
		indexValues, err := indexFunc(newObj)
		if err != nil {
			panic(fmt.Errorf("unable to calculate an index entry for key %q on index %q: %v", key, name, err))
		}
		index := c.indices[name]
		if index == nil {
			index = Index{}
			c.indices[name] = index
		}

		for _, indexValue := range indexValues {
			set := index[indexValue]
			if set == nil {
				set = sets.String{}
				index[indexValue] = set
			}
			set.Insert(key)
		}
	}
}

// deleteFromIndices removes the object from each of the managed indexes
// it is intended to be called from a function that already has a lock on the cache
func (c *threadSafeMap) deleteFromIndices(obj interface{}, key string) {
	for name, indexFunc := range c.indexers {
		indexValues, err := indexFunc(obj)
		if err != nil {
			panic(fmt.Errorf("unable to calculate an index entry for key %q on index %q: %v", key, name, err))
		}

		index := c.indices[name]
		if index == nil {
			continue
		}
		for _, indexValue := range indexValues {
			set := index[indexValue]
			if set != nil {
				set.Delete(key)
			}
		}
	}
}

func (c *threadSafeMap) Resync() error {
	// Nothing to do
	return nil
}

func NewThreadSafeStore(indexers Indexers, indices Indices) ThreadSafeStore {
	return &threadSafeMap{
		items:    map[string]interface{}{},
		indexers: indexers,
		indices:  indices,
	}
}
<commit_msg>Avoid allocations in ByIndex() function<commit_after>/*
Copyright 2014 The Kubernetes 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 cache

import (
	"fmt"
	"sync"

	"k8s.io/apimachinery/pkg/util/sets"
)

// ThreadSafeStore is an interface that allows concurrent access to a storage backend.
// TL;DR caveats: you must not modify anything returned by Get or List as it will break
// the indexing feature in addition to not being thread safe.
//
// The guarantees of thread safety provided by List/Get are only valid if the caller
// treats returned items as read-only. For example, a pointer inserted in the store
// through `Add` will be returned as is by `Get`. Multiple clients might invoke `Get`
// on the same key and modify the pointer in a non-thread-safe way. Also note that
// modifying objects stored by the indexers (if any) will *not* automatically lead
// to a re-index. So it's not a good idea to directly modify the objects returned by
// Get/List, in general.
type ThreadSafeStore interface {
	Add(key string, obj interface{})
	Update(key string, obj interface{})
	Delete(key string)
	Get(key string) (item interface{}, exists bool)
	List() []interface{}
	ListKeys() []string
	Replace(map[string]interface{}, string)
	Index(indexName string, obj interface{}) ([]interface{}, error)
	IndexKeys(indexName, indexKey string) ([]string, error)
	ListIndexFuncValues(name string) []string
	ByIndex(indexName, indexKey string) ([]interface{}, error)
	GetIndexers() Indexers

	// AddIndexers adds more indexers to this store.  If you call this after you already have data
	// in the store, the results are undefined.
	AddIndexers(newIndexers Indexers) error
	Resync() error
}

// threadSafeMap implements ThreadSafeStore
type threadSafeMap struct {
	lock  sync.RWMutex
	items map[string]interface{}

	// indexers maps a name to an IndexFunc
	indexers Indexers
	// indices maps a name to an Index
	indices Indices
}

func (c *threadSafeMap) Add(key string, obj interface{}) {
	c.lock.Lock()
	defer c.lock.Unlock()
	oldObject := c.items[key]
	c.items[key] = obj
	c.updateIndices(oldObject, obj, key)
}

func (c *threadSafeMap) Update(key string, obj interface{}) {
	c.lock.Lock()
	defer c.lock.Unlock()
	oldObject := c.items[key]
	c.items[key] = obj
	c.updateIndices(oldObject, obj, key)
}

func (c *threadSafeMap) Delete(key string) {
	c.lock.Lock()
	defer c.lock.Unlock()
	if obj, exists := c.items[key]; exists {
		c.deleteFromIndices(obj, key)
		delete(c.items, key)
	}
}

func (c *threadSafeMap) Get(key string) (item interface{}, exists bool) {
	c.lock.RLock()
	defer c.lock.RUnlock()
	item, exists = c.items[key]
	return item, exists
}

func (c *threadSafeMap) List() []interface{} {
	c.lock.RLock()
	defer c.lock.RUnlock()
	list := make([]interface{}, 0, len(c.items))
	for _, item := range c.items {
		list = append(list, item)
	}
	return list
}

// ListKeys returns a list of all the keys of the objects currently
// in the threadSafeMap.
func (c *threadSafeMap) ListKeys() []string {
	c.lock.RLock()
	defer c.lock.RUnlock()
	list := make([]string, 0, len(c.items))
	for key := range c.items {
		list = append(list, key)
	}
	return list
}

func (c *threadSafeMap) Replace(items map[string]interface{}, resourceVersion string) {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.items = items

	// rebuild any index
	c.indices = Indices{}
	for key, item := range c.items {
		c.updateIndices(nil, item, key)
	}
}

// Index returns a list of items that match on the index function
// Index is thread-safe so long as you treat all items as immutable
func (c *threadSafeMap) Index(indexName string, obj interface{}) ([]interface{}, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()

	indexFunc := c.indexers[indexName]
	if indexFunc == nil {
		return nil, fmt.Errorf("Index with name %s does not exist", indexName)
	}

	indexKeys, err := indexFunc(obj)
	if err != nil {
		return nil, err
	}
	index := c.indices[indexName]

	var returnKeySet sets.String
	if len(indexKeys) == 1 {
		// In majority of cases, there is exactly one value matching.
		// Optimize the most common path - deduping is not needed here.
		returnKeySet = index[indexKeys[0]]
	} else {
		// Need to de-dupe the return list.
		// Since multiple keys are allowed, this can happen.
		returnKeySet = sets.String{}
		for _, indexKey := range indexKeys {
			for key := range index[indexKey] {
				returnKeySet.Insert(key)
			}
		}
	}

	list := make([]interface{}, 0, returnKeySet.Len())
	for absoluteKey := range returnKeySet {
		list = append(list, c.items[absoluteKey])
	}
	return list, nil
}

// ByIndex returns a list of items that match an exact value on the index function
func (c *threadSafeMap) ByIndex(indexName, indexKey string) ([]interface{}, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()

	indexFunc := c.indexers[indexName]
	if indexFunc == nil {
		return nil, fmt.Errorf("Index with name %s does not exist", indexName)
	}

	index := c.indices[indexName]

	set := index[indexKey]
	list := make([]interface{}, 0, set.Len())
	for key := range set {
		list = append(list, c.items[key])
	}

	return list, nil
}

// IndexKeys returns a list of keys that match on the index function.
// IndexKeys is thread-safe so long as you treat all items as immutable.
func (c *threadSafeMap) IndexKeys(indexName, indexKey string) ([]string, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()

	indexFunc := c.indexers[indexName]
	if indexFunc == nil {
		return nil, fmt.Errorf("Index with name %s does not exist", indexName)
	}

	index := c.indices[indexName]

	set := index[indexKey]
	return set.List(), nil
}

func (c *threadSafeMap) ListIndexFuncValues(indexName string) []string {
	c.lock.RLock()
	defer c.lock.RUnlock()

	index := c.indices[indexName]
	names := make([]string, 0, len(index))
	for key := range index {
		names = append(names, key)
	}
	return names
}

func (c *threadSafeMap) GetIndexers() Indexers {
	return c.indexers
}

func (c *threadSafeMap) AddIndexers(newIndexers Indexers) error {
	c.lock.Lock()
	defer c.lock.Unlock()

	if len(c.items) > 0 {
		return fmt.Errorf("cannot add indexers to running index")
	}

	oldKeys := sets.StringKeySet(c.indexers)
	newKeys := sets.StringKeySet(newIndexers)

	if oldKeys.HasAny(newKeys.List()...) {
		return fmt.Errorf("indexer conflict: %v", oldKeys.Intersection(newKeys))
	}

	for k, v := range newIndexers {
		c.indexers[k] = v
	}
	return nil
}

// updateIndices modifies the objects location in the managed indexes, if this is an update, you must provide an oldObj
// updateIndices must be called from a function that already has a lock on the cache
func (c *threadSafeMap) updateIndices(oldObj interface{}, newObj interface{}, key string) {
	// if we got an old object, we need to remove it before we add it again
	if oldObj != nil {
		c.deleteFromIndices(oldObj, key)
	}
	for name, indexFunc := range c.indexers {
		indexValues, err := indexFunc(newObj)
		if err != nil {
			panic(fmt.Errorf("unable to calculate an index entry for key %q on index %q: %v", key, name, err))
		}
		index := c.indices[name]
		if index == nil {
			index = Index{}
			c.indices[name] = index
		}

		for _, indexValue := range indexValues {
			set := index[indexValue]
			if set == nil {
				set = sets.String{}
				index[indexValue] = set
			}
			set.Insert(key)
		}
	}
}

// deleteFromIndices removes the object from each of the managed indexes
// it is intended to be called from a function that already has a lock on the cache
func (c *threadSafeMap) deleteFromIndices(obj interface{}, key string) {
	for name, indexFunc := range c.indexers {
		indexValues, err := indexFunc(obj)
		if err != nil {
			panic(fmt.Errorf("unable to calculate an index entry for key %q on index %q: %v", key, name, err))
		}

		index := c.indices[name]
		if index == nil {
			continue
		}
		for _, indexValue := range indexValues {
			set := index[indexValue]
			if set != nil {
				set.Delete(key)
			}
		}
	}
}

func (c *threadSafeMap) Resync() error {
	// Nothing to do
	return nil
}

func NewThreadSafeStore(indexers Indexers, indices Indices) ThreadSafeStore {
	return &threadSafeMap{
		items:    map[string]interface{}{},
		indexers: indexers,
		indices:  indices,
	}
}
<|endoftext|>


<commit_before>package pg

import (
	"log"
	"os"
	"strconv"

	"gopkg.in/pg.v4/internal/pool"
	"gopkg.in/pg.v4/orm"
	"gopkg.in/pg.v4/types"
)

// Discard is be used with Query and QueryOne to discard rows.
var Discard orm.Discard

// Scan returns ColumnScanner that copies the columns in the
// row into the values.
var Scan = orm.Scan

// Q returns a ValueAppender that represents safe SQL query.
var Q = orm.Q

// F returns a ValueAppender that represents SQL identifier,
// e.g. table or column name.
var F = orm.F

// Array returns an Array type that represents PostgreSQL array of any type.
func Array(v interface{}) *types.Array {
	return types.NewArray(v)
}

func init() {
	SetLogger(log.New(os.Stderr, "pg: ", log.LstdFlags))
}

func SetLogger(logger *log.Logger) {
	pool.Logger = logger
}

//------------------------------------------------------------------------------

type Strings []string

var _ orm.Model = (*Strings)(nil)
var _ types.ValueAppender = (*Strings)(nil)

func (strings *Strings) NewModel() orm.ColumnScanner {
	return strings
}

func (Strings) AddModel(_ orm.ColumnScanner) error {
	return nil
}

func (strings *Strings) ScanColumn(colIdx int, _ string, b []byte) error {
	*strings = append(*strings, string(b))
	return nil
}

func (strings Strings) AppendValue(dst []byte, quote int) ([]byte, error) {
	if len(strings) <= 0 {
		return dst, nil
	}

	for _, s := range strings {
		dst = types.AppendString(dst, s, 1)
		dst = append(dst, ',')
	}
	dst = dst[:len(dst)-1]
	return dst, nil
}

//------------------------------------------------------------------------------

type Ints []int64

var _ orm.Model = (*Ints)(nil)
var _ types.ValueAppender = (*Ints)(nil)

func (ints *Ints) NewModel() orm.ColumnScanner {
	return ints
}

func (Ints) AddModel(_ orm.ColumnScanner) error {
	return nil
}

func (ints *Ints) ScanColumn(colIdx int, colName string, b []byte) error {
	n, err := strconv.ParseInt(string(b), 10, 64)
	if err != nil {
		return err
	}
	*ints = append(*ints, n)
	return nil
}

func (ints Ints) AppendValue(dst []byte, quote int) ([]byte, error) {
	if len(ints) <= 0 {
		return dst, nil
	}

	for _, v := range ints {
		dst = strconv.AppendInt(dst, v, 10)
		dst = append(dst, ',')
	}
	dst = dst[:len(dst)-1]
	return dst, nil
}

//------------------------------------------------------------------------------

type IntSet map[int64]struct{}

var _ orm.Model = (*IntSet)(nil)

func (set *IntSet) NewModel() orm.ColumnScanner {
	return set
}

func (IntSet) AddModel(_ orm.ColumnScanner) error {
	return nil
}

func (setptr *IntSet) ScanColumn(colIdx int, colName string, b []byte) error {
	set := *setptr
	if set == nil {
		*setptr = make(IntSet)
		set = *setptr
	}

	n, err := strconv.ParseInt(string(b), 10, 64)
	if err != nil {
		return err
	}
	set[n] = struct{}{}
	return nil
}
<commit_msg>Fix doc.<commit_after>package pg

import (
	"log"
	"os"
	"strconv"

	"gopkg.in/pg.v4/internal/pool"
	"gopkg.in/pg.v4/orm"
	"gopkg.in/pg.v4/types"
)

// Discard is used with Query and QueryOne to discard rows.
var Discard orm.Discard

// Scan returns ColumnScanner that copies the columns in the
// row into the values.
var Scan = orm.Scan

// Q returns a ValueAppender that represents safe SQL query.
var Q = orm.Q

// F returns a ValueAppender that represents SQL identifier,
// e.g. table or column name.
var F = orm.F

// Array returns an Array type that represents PostgreSQL array of any type.
func Array(v interface{}) *types.Array {
	return types.NewArray(v)
}

func init() {
	SetLogger(log.New(os.Stderr, "pg: ", log.LstdFlags))
}

func SetLogger(logger *log.Logger) {
	pool.Logger = logger
}

//------------------------------------------------------------------------------

type Strings []string

var _ orm.Model = (*Strings)(nil)
var _ types.ValueAppender = (*Strings)(nil)

func (strings *Strings) NewModel() orm.ColumnScanner {
	return strings
}

func (Strings) AddModel(_ orm.ColumnScanner) error {
	return nil
}

func (strings *Strings) ScanColumn(colIdx int, _ string, b []byte) error {
	*strings = append(*strings, string(b))
	return nil
}

func (strings Strings) AppendValue(dst []byte, quote int) ([]byte, error) {
	if len(strings) <= 0 {
		return dst, nil
	}

	for _, s := range strings {
		dst = types.AppendString(dst, s, 1)
		dst = append(dst, ',')
	}
	dst = dst[:len(dst)-1]
	return dst, nil
}

//------------------------------------------------------------------------------

type Ints []int64

var _ orm.Model = (*Ints)(nil)
var _ types.ValueAppender = (*Ints)(nil)

func (ints *Ints) NewModel() orm.ColumnScanner {
	return ints
}

func (Ints) AddModel(_ orm.ColumnScanner) error {
	return nil
}

func (ints *Ints) ScanColumn(colIdx int, colName string, b []byte) error {
	n, err := strconv.ParseInt(string(b), 10, 64)
	if err != nil {
		return err
	}
	*ints = append(*ints, n)
	return nil
}

func (ints Ints) AppendValue(dst []byte, quote int) ([]byte, error) {
	if len(ints) <= 0 {
		return dst, nil
	}

	for _, v := range ints {
		dst = strconv.AppendInt(dst, v, 10)
		dst = append(dst, ',')
	}
	dst = dst[:len(dst)-1]
	return dst, nil
}

//------------------------------------------------------------------------------

type IntSet map[int64]struct{}

var _ orm.Model = (*IntSet)(nil)

func (set *IntSet) NewModel() orm.ColumnScanner {
	return set
}

func (IntSet) AddModel(_ orm.ColumnScanner) error {
	return nil
}

func (setptr *IntSet) ScanColumn(colIdx int, colName string, b []byte) error {
	set := *setptr
	if set == nil {
		*setptr = make(IntSet)
		set = *setptr
	}

	n, err := strconv.ParseInt(string(b), 10, 64)
	if err != nil {
		return err
	}
	set[n] = struct{}{}
	return nil
}
<|endoftext|>


<commit_before>package factorlib

import (
	"fmt"
	"github.com/randall77/factorlib/linear"
	"math/big"
	"math/rand"
)

func init() {
	factorizers["qs"] = qs
}

func qs(n big.Int, rnd *rand.Rand) []big.Int {
	// Quadratic sieve:
	// Define f(x) = x^2 - n.  We wish to find x such that f(x) is smooth.
	// We will take x starting at ceil(sqrt(n)).

	// is f(x) divisible by p^k?
	// calculating mod p^k, f(x + p^k) = (x+p^k)^2 - n = x^2 - n = f(x)
	// so if f(x) is divisible by p^k, so is f(x + i * p^k) for all i.
	//
	// if x_1 and x_2 are the two solutions to f(x) == 0 mod p^k, then all x which
	// have f(x) divisible by p^k are x_1 + i * p^k and x_2 + i * p^k.
	//
	// f(x) == 0 mod p^k is x^2 == n mod p^k.

	// So the factor base contains small primes and powers of primes.
	// For each factor base element p^k, we store two offsets satisfying
	// f(start + off1) == 0 mod p^k and f(start + off2) == 0 mod p^k.
	// Then to sieve, we just add log_p to sieve[off1], sieve[off1+p^k], sieve[off1+2*p^k], ...
	// and the same for off2.  When we reach the end of the sieve, we update off1 and off2
	// for the next sieve.

	// the sieve returns x such that f(x) = x^2 - n = prod_i factors[i]
	// all factors <= B
	// factors are sorted in increasing order.

	// first, pick a factor base
	fb, a := makeFactorBase(n)
	if a.Sign() != 0 {
		var b big.Int
		b.Div(&n, &a)
		return []big.Int{a, b}
	}

	maxp := fb[len(fb)-1]
	var bigmaxp2 big.Int
	bigmaxp2.SetInt64(maxp * maxp)
	fmt.Printf("maxp=%d maxp2=%d len(fb)=%d\n", maxp, &bigmaxp2, len(fb))

	// The x we start with
	start := *sqrtCeil(&n)

	si := makeSieveInfo(n, start, fb, rnd)

	sievelen := int(fb[len(fb)-1])
	if sievelen < 10000 {
		sievelen = 10000
	}

	sieve := make([]byte, sievelen)

	// temporary storage
	var x, y, bigz, r, bigp, t big.Int
	var factors []uint

	// matrix is used to do gaussian elimination on mod 2 exponents.
	m := linear.NewMatrix(uint(len(fb)))

	// largeprimes records instances of the equation
	//   x^2 = prod(f) * p mod n
	// where f are small primes (indexes into factor base) and
	// p is a large prime.
	type largerecord struct {
		x big.Int
		f []uint
	}
	largeprimes := map[int64]largerecord{}

	for {
		// clear sieve
		for i := 0; i < sievelen; i++ {
			sieve[i] = 0
		}

		// increment sieve entries for f(x) that are divisible
		// by each factor base prime.
		for i := range si {
			f := &si[i]
			pk := int(f.pk)
			lg_p := f.lg_p
			j := int(f.off1)
			for ; j < sievelen; j += pk {
				sieve[j] += lg_p
			}
			f.off1 = int32(j - sievelen) // for next time

			// same for off2
			j = int(f.off2)
			if j < 0 {
				continue // special case needed for p==2,k==1
			}
			for ; j < sievelen; j += pk {
				sieve[j] += lg_p
			}
			f.off2 = int32(j - sievelen) // for next time
		}

		// check sieve entries for big numbers, indicating smooth f(x).
		t.Mul(&start, &start)
		t.Sub(&t, &n)
		threshold := byte(t.BitLen()) - 2*log2(maxp)
		for i := 0; i < sievelen; i++ {
			if sieve[i] < threshold {
				// TODO: in testing mode, check if x^2-n is smooth.
				continue
			}
			//fmt.Printf("%d+%d %d %d\n", start, i, sieve[i], threshold)

			// compute x, y=f(x)
			x.SetInt64(int64(i))
			x.Add(&start, &x)
			y.Mul(&x, &x)
			y.Sub(&y, &n)

			// trial divide y by the factor base
			// accumulate factor base indexes of factors
			bigz.Set(&y)
			factors = factors[:0]
			for i, p := range fb {
				bigp.SetInt64(p)
				for r.Mod(&bigz, &bigp).Sign() == 0 {
					bigz.Div(&bigz, &bigp)
					factors = append(factors, uint(i))
				}
			}

			// if remainder > B^2, it's too big
			if bigz.Cmp(&bigmaxp2) > 0 {
				//fmt.Printf("  false positive y=%d z=%d threshold=%d sieve[i]=%d log2(y)=%d log2(y/z)=%d\n", y, bigz, threshold, sieve[i], y.BitLen(), x.Div(y, bigz).BitLen())
				continue
			}

			z := bigz.Int64()
			if z > 1 {
				// try to find another record with the same largeprime
				lr, ok := largeprimes[z]
				if !ok {
					// haven't seen this large prime yet.  Save record for later
					largeprimes[z] = largerecord{*new(big.Int).Set(&x), dup(factors)}
					//fmt.Printf("  savelarge %d %v\n", z, factors)
					continue
				}
				// combine current equation with other largeprime equation
				// x1^2 === prod(f1) * largeprime
				// x2^2 === prod(f2) * largeprime
				//fmt.Printf("  largeprime %d match\n", bigz)
				x.Mul(&x, &lr.x)
				x.Mod(&x, &n)
				y.ModInverse(&bigz, &n) // TODO: could bigz divide n?
				x.Mul(&x, &y)
				x.Mod(&x, &n)
				factors = append(factors, lr.f...)
			}
			fmt.Printf("eqn%d/%d %d^2 === ", m.Rows(), len(fb), &x)
			for j, i := range factors {
				if j > 0 {
					fmt.Printf("·")
				}
				fmt.Printf("%d", fb[i])
			}
			fmt.Printf("\n")
			idlist := m.AddRow(factors, eqn{*new(big.Int).Set(&x), dup(factors)})
			if idlist == nil {
				continue
			}

			// we found a set of equations with all even powers
			// compute a and b where a^2 === b^2 mod n
			var a, b big.Int
			a.SetInt64(1)
			b.SetInt64(1)
			odd := make([]bool, len(fb))
			for _, id := range idlist {
				e := id.(eqn)
				a.Mul(&a, &e.x)
				a.Mod(&a, &n)
				for _, i := range e.f {
					if !odd[i] {
						// first occurrence of this factor
						odd[i] = true
						continue
					}
					// second occurrence of this factor
					t.SetInt64(fb[i])
					b.Mul(&b, &t)
					b.Mod(&b, &n)
					odd[i] = false
				}
			}

			for i, p := range fb {
				if odd[i] {
					fmt.Printf("prime i=%d p=%d\n", i, p)
					panic("gauss elim failed")
				}
			}

			if a.Cmp(&b) == 0 {
				// trivial equation, ignore it
				fmt.Println("triv A")
				continue
			}
			t.Add(&a, &b)
			if t.Cmp(&n) == 0 {
				// trivial equation, ignore it
				fmt.Println("triv B")
				continue
			}

			t.GCD(nil, nil, &t, &n)
			r.Div(&n, &t)
			return []big.Int{t, r}
		}

		start.Add(&start, x.SetInt64(int64(sievelen)))
	}
	return nil
}

// x^2 === prod(f) mod n
// f is a list of indexes into the factor base
type eqn struct {
	x big.Int
	f []uint
}

// pick some prime factors for our factor base.  If we happen
// upon a factor of n, return it instead.
func makeFactorBase(n big.Int) ([]int64, big.Int) {
	// upper limit on prime factors (TODO: dependent on n) that we sieve with
	const B = 50000
	var fb []int64
	var biga, bigp big.Int
	for i := 0; ; i++ {
		p := getPrime(i)
		if p > B {
			return fb, zero
		}
		bigp.SetInt64(p)
		biga.Mod(&n, &bigp)
		a := biga.Int64()
		if a == 0 {
			return nil, bigp
		}
		if quadraticResidue(a, p) {
			// if x^2 == n mod p has no solutions, then
			// x^2-n will never be divisible by p.  So it
			// doesn't need to be in the factor base.
			fb = append(fb, p)
		}
	}
}

type sieveinfo struct {
	pk         int32 // p^k for this factor base entry
	lg_p       uint8 // ~log_2(p)
	off1, off2 int32 // starting offsets in sieve array
}

func makeSieveInfo(n big.Int, start big.Int, fb []int64, rnd *rand.Rand) []sieveinfo {
	var si []sieveinfo
	var bigpk, biga, m big.Int

	for _, p := range fb {
		pk := p
		for k := uint(1); ; k++ {
			if pk > fb[len(fb)-1] {
				break
			}
			bigpk.SetInt64(pk)
			biga.Mod(&n, &bigpk)
			z1 := sqrtModPK(biga.Int64(), p, k, rnd) // find solution to x^2 == a mod p^k
			if z1 == 0 {
				panic("factor base divides n")
			}
			z2 := pk - z1

			// adjust to start
			s := m.Mod(&start, &bigpk).Int64()
			off1 := z1 - s
			if off1 < 0 {
				off1 += pk
			}
			off2 := z2 - s
			if off2 < 0 {
				off2 += pk
			}
			if off1 == off2 {
				// should only happen for pk==2
				if pk != 2 {
					panic("duplicate offsets")
				}
				off2 = -1
			}

			si = append(si, sieveinfo{int32(pk), log2(p), int32(off1), int32(off2)})
			if p == 2 && k == 1 {
				// TODO: handle powers of 2 correctly.  Hansel's lemma
				// doesn't work for them.
				break
			}
			pk *= p
		}
	}
	return si
}

func dup(x []uint) []uint {
	y := make([]uint, len(x))
	copy(y, x)
	return y
}
<commit_msg>Build results with big.Int in them separately<commit_after>package factorlib

import (
	"fmt"
	"github.com/randall77/factorlib/linear"
	"math/big"
	"math/rand"
)

func init() {
	factorizers["qs"] = qs
}

func qs(n big.Int, rnd *rand.Rand) []big.Int {
	// Quadratic sieve:
	// Define f(x) = x^2 - n.  We wish to find x such that f(x) is smooth.
	// We will take x starting at ceil(sqrt(n)).

	// is f(x) divisible by p^k?
	// calculating mod p^k, f(x + p^k) = (x+p^k)^2 - n = x^2 - n = f(x)
	// so if f(x) is divisible by p^k, so is f(x + i * p^k) for all i.
	//
	// if x_1 and x_2 are the two solutions to f(x) == 0 mod p^k, then all x which
	// have f(x) divisible by p^k are x_1 + i * p^k and x_2 + i * p^k.
	//
	// f(x) == 0 mod p^k is x^2 == n mod p^k.

	// So the factor base contains small primes and powers of primes.
	// For each factor base element p^k, we store two offsets satisfying
	// f(start + off1) == 0 mod p^k and f(start + off2) == 0 mod p^k.
	// Then to sieve, we just add log_p to sieve[off1], sieve[off1+p^k], sieve[off1+2*p^k], ...
	// and the same for off2.  When we reach the end of the sieve, we update off1 and off2
	// for the next sieve.

	// the sieve returns x such that f(x) = x^2 - n = prod_i factors[i]
	// all factors <= B
	// factors are sorted in increasing order.

	// first, pick a factor base
	fb, a := makeFactorBase(n)
	if a.Sign() != 0 {
		var b big.Int
		b.Div(&n, &a)
		return []big.Int{a, b}
	}

	maxp := fb[len(fb)-1]
	var bigmaxp2 big.Int
	bigmaxp2.SetInt64(maxp * maxp)
	fmt.Printf("maxp=%d maxp2=%d len(fb)=%d\n", maxp, &bigmaxp2, len(fb))

	// The x we start with
	start := *sqrtCeil(&n)

	si := makeSieveInfo(n, start, fb, rnd)

	sievelen := int(fb[len(fb)-1])
	if sievelen < 10000 {
		sievelen = 10000
	}

	sieve := make([]byte, sievelen)

	// temporary storage
	var x, y, bigz, r, bigp, t big.Int
	var factors []uint

	// matrix is used to do gaussian elimination on mod 2 exponents.
	m := linear.NewMatrix(uint(len(fb)))

	// largeprimes records instances of the equation
	//   x^2 = prod(f) * p mod n
	// where f are small primes (indexes into factor base) and
	// p is a large prime.
	type largerecord struct {
		x big.Int
		f []uint
	}
	largeprimes := map[int64]largerecord{}

	for {
		// clear sieve
		for i := 0; i < sievelen; i++ {
			sieve[i] = 0
		}

		// increment sieve entries for f(x) that are divisible
		// by each factor base prime.
		for i := range si {
			f := &si[i]
			pk := int(f.pk)
			lg_p := f.lg_p
			j := int(f.off1)
			for ; j < sievelen; j += pk {
				sieve[j] += lg_p
			}
			f.off1 = int32(j - sievelen) // for next time

			// same for off2
			j = int(f.off2)
			if j < 0 {
				continue // special case needed for p==2,k==1
			}
			for ; j < sievelen; j += pk {
				sieve[j] += lg_p
			}
			f.off2 = int32(j - sievelen) // for next time
		}

		// check sieve entries for big numbers, indicating smooth f(x).
		t.Mul(&start, &start)
		t.Sub(&t, &n)
		threshold := byte(t.BitLen()) - 2*log2(maxp)
		for i := 0; i < sievelen; i++ {
			if sieve[i] < threshold {
				// TODO: in testing mode, check if x^2-n is smooth.
				continue
			}
			//fmt.Printf("%d+%d %d %d\n", start, i, sieve[i], threshold)

			// compute x, y=f(x)
			x.SetInt64(int64(i))
			x.Add(&start, &x)
			y.Mul(&x, &x)
			y.Sub(&y, &n)

			// trial divide y by the factor base
			// accumulate factor base indexes of factors
			bigz.Set(&y)
			factors = factors[:0]
			for i, p := range fb {
				bigp.SetInt64(p)
				for r.Mod(&bigz, &bigp).Sign() == 0 {
					bigz.Div(&bigz, &bigp)
					factors = append(factors, uint(i))
				}
			}

			// if remainder > B^2, it's too big
			if bigz.Cmp(&bigmaxp2) > 0 {
				//fmt.Printf("  false positive y=%d z=%d threshold=%d sieve[i]=%d log2(y)=%d log2(y/z)=%d\n", y, bigz, threshold, sieve[i], y.BitLen(), x.Div(y, bigz).BitLen())
				continue
			}

			z := bigz.Int64()
			if z > 1 {
				// try to find another record with the same largeprime
				lr, ok := largeprimes[z]
				if !ok {
					var lr largerecord
					lr.x.Set(&x)
					lr.f = dup(factors)
					// haven't seen this large prime yet.  Save record for later
					largeprimes[z] = lr
					//fmt.Printf("  savelarge %d %v\n", z, factors)
					continue
				}
				// combine current equation with other largeprime equation
				// x1^2 === prod(f1) * largeprime
				// x2^2 === prod(f2) * largeprime
				//fmt.Printf("  largeprime %d match\n", bigz)
				x.Mul(&x, &lr.x)
				x.Mod(&x, &n)
				y.ModInverse(&bigz, &n) // TODO: could bigz divide n?
				x.Mul(&x, &y)
				x.Mod(&x, &n)
				factors = append(factors, lr.f...)
			}
			fmt.Printf("eqn%d/%d %d^2 === ", m.Rows(), len(fb), &x)
			for j, i := range factors {
				if j > 0 {
					fmt.Printf("·")
				}
				fmt.Printf("%d", fb[i])
			}
			fmt.Printf("\n")
			var e eqn
			e.x.Set(&x)
			e.f = dup(factors)
			idlist := m.AddRow(factors, e)
			if idlist == nil {
				continue
			}

			// we found a set of equations with all even powers
			// compute a and b where a^2 === b^2 mod n
			var a, b big.Int
			a.SetInt64(1)
			b.SetInt64(1)
			odd := make([]bool, len(fb))
			for _, id := range idlist {
				e := id.(eqn)
				a.Mul(&a, &e.x)
				a.Mod(&a, &n)
				for _, i := range e.f {
					if !odd[i] {
						// first occurrence of this factor
						odd[i] = true
						continue
					}
					// second occurrence of this factor
					t.SetInt64(fb[i])
					b.Mul(&b, &t)
					b.Mod(&b, &n)
					odd[i] = false
				}
			}

			for i, p := range fb {
				if odd[i] {
					fmt.Printf("prime i=%d p=%d\n", i, p)
					panic("gauss elim failed")
				}
			}

			if a.Cmp(&b) == 0 {
				// trivial equation, ignore it
				fmt.Println("triv A")
				continue
			}
			t.Add(&a, &b)
			if t.Cmp(&n) == 0 {
				// trivial equation, ignore it
				fmt.Println("triv B")
				continue
			}

			t.GCD(nil, nil, &t, &n)
			r.Div(&n, &t)
			return []big.Int{t, r}
		}

		start.Add(&start, x.SetInt64(int64(sievelen)))
	}
	return nil
}

// x^2 === prod(f) mod n
// f is a list of indexes into the factor base
type eqn struct {
	x big.Int
	f []uint
}

// pick some prime factors for our factor base.  If we happen
// upon a factor of n, return it instead.
func makeFactorBase(n big.Int) ([]int64, big.Int) {
	// upper limit on prime factors (TODO: dependent on n) that we sieve with
	const B = 50000
	var fb []int64
	var biga, bigp big.Int
	for i := 0; ; i++ {
		p := getPrime(i)
		if p > B {
			return fb, zero
		}
		bigp.SetInt64(p)
		biga.Mod(&n, &bigp)
		a := biga.Int64()
		if a == 0 {
			return nil, bigp
		}
		if quadraticResidue(a, p) {
			// if x^2 == n mod p has no solutions, then
			// x^2-n will never be divisible by p.  So it
			// doesn't need to be in the factor base.
			fb = append(fb, p)
		}
	}
}

type sieveinfo struct {
	pk         int32 // p^k for this factor base entry
	lg_p       uint8 // ~log_2(p)
	off1, off2 int32 // starting offsets in sieve array
}

func makeSieveInfo(n big.Int, start big.Int, fb []int64, rnd *rand.Rand) []sieveinfo {
	var si []sieveinfo
	var bigpk, biga, m big.Int

	for _, p := range fb {
		pk := p
		for k := uint(1); ; k++ {
			if pk > fb[len(fb)-1] {
				break
			}
			bigpk.SetInt64(pk)
			biga.Mod(&n, &bigpk)
			z1 := sqrtModPK(biga.Int64(), p, k, rnd) // find solution to x^2 == a mod p^k
			if z1 == 0 {
				panic("factor base divides n")
			}
			z2 := pk - z1

			// adjust to start
			s := m.Mod(&start, &bigpk).Int64()
			off1 := z1 - s
			if off1 < 0 {
				off1 += pk
			}
			off2 := z2 - s
			if off2 < 0 {
				off2 += pk
			}
			if off1 == off2 {
				// should only happen for pk==2
				if pk != 2 {
					panic("duplicate offsets")
				}
				off2 = -1
			}

			si = append(si, sieveinfo{int32(pk), log2(p), int32(off1), int32(off2)})
			if p == 2 && k == 1 {
				// TODO: handle powers of 2 correctly.  Hansel's lemma
				// doesn't work for them.
				break
			}
			pk *= p
		}
	}
	return si
}

func dup(x []uint) []uint {
	y := make([]uint, len(x))
	copy(y, x)
	return y
}
<|endoftext|>


<commit_before>// Copyright 2016 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package dockermachine

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/docker/machine/libmachine"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/drivers/rpc"
	"github.com/docker/machine/libmachine/log"
	"github.com/pkg/errors"
	"github.com/tsuru/tsuru/iaas"
)

var defaultWriter = ioutil.Discard

type DockerMachine struct {
	io.Closer
	client libmachine.API
	config *DockerMachineConfig
	path   string
}

type DockerMachineConfig struct {
	CaPath                 string
	InsecureRegistry       string
	DockerEngineInstallURL string
	OutWriter              io.Writer
	ErrWriter              io.Writer
}

type dockerMachineAPI interface {
	io.Closer
	CreateMachine(string, string, map[string]interface{}) (*iaas.Machine, error)
	DeleteMachine(*iaas.Machine) error
}

func NewDockerMachine(config DockerMachineConfig) (dockerMachineAPI, error) {
	path, err := ioutil.TempDir("", "")
	if err != nil {
		return nil, errors.Wrap(err, "failed to create temp dir")
	}
	if config.CaPath != "" {
		err = copy(filepath.Join(config.CaPath, "ca.pem"), filepath.Join(path, "ca.pem"))
		if err != nil {
			return nil, errors.Wrap(err, "failed to copy ca file")
		}
		err = copy(filepath.Join(config.CaPath, "ca-key.pem"), filepath.Join(path, "ca-key.pem"))
		if err != nil {
			return nil, errors.Wrap(err, "failed to copy ca key file")
		}
	}
	if config.OutWriter != nil {
		log.SetOutWriter(config.OutWriter)
	} else {
		log.SetOutWriter(defaultWriter)
	}
	if config.ErrWriter != nil {
		log.SetOutWriter(config.ErrWriter)
	} else {
		log.SetOutWriter(defaultWriter)
	}
	return &DockerMachine{
		path:   path,
		client: libmachine.NewClient(path, path),
		config: &config,
	}, nil
}

func (d *DockerMachine) Close() error {
	os.RemoveAll(d.path)
	return d.client.Close()
}

func (d *DockerMachine) CreateMachine(name, driver string, params map[string]interface{}) (*iaas.Machine, error) {
	rawDriver, err := json.Marshal(&drivers.BaseDriver{
		MachineName: name,
		StorePath:   d.path,
	})
	if err != nil {
		return nil, errors.Wrap(err, "failed to marshal base driver")
	}
	host, err := d.client.NewHost(driver, rawDriver)
	if err != nil {
		return nil, errors.Wrap(err, "failed to initialize host")
	}
	err = configureDriver(host.Driver, params)
	if err != nil {
		return nil, errors.Wrap(err, "failed to configure driver")
	}
	engineOpts := host.HostOptions.EngineOptions
	if d.config.InsecureRegistry != "" {
		engineOpts.InsecureRegistry = []string{d.config.InsecureRegistry}
	}
	if d.config.DockerEngineInstallURL != "" {
		engineOpts.InstallURL = d.config.DockerEngineInstallURL
	}
	err = d.client.Create(host)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create host")
	}
	rawDriver, err = json.Marshal(host.Driver)
	if err != nil {
		return nil, errors.Wrap(err, "failed to marshal host driver")
	}
	var driverData map[string]interface{}
	err = json.Unmarshal(rawDriver, &driverData)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal host driver")
	}
	m := &iaas.Machine{
		Id:         host.Name,
		Port:       2376,
		Protocol:   "https",
		CustomData: driverData,
	}
	m.Address, err = host.Driver.GetIP()
	if err != nil {
		return m, errors.Wrap(err, "failed to retrive host ip")
	}
	if host.AuthOptions() != nil {
		m.CaCert, err = ioutil.ReadFile(host.AuthOptions().CaCertPath)
		if err != nil {
			return m, errors.Wrap(err, "failed to read host ca cert")
		}
		m.ClientCert, err = ioutil.ReadFile(host.AuthOptions().ClientCertPath)
		if err != nil {
			return m, errors.Wrap(err, "failed to read host client cert")
		}
		m.ClientKey, err = ioutil.ReadFile(host.AuthOptions().ClientKeyPath)
		if err != nil {
			return m, errors.Wrap(err, "failed to read host client key")
		}
	}
	return m, nil
}

func (d *DockerMachine) DeleteMachine(m *iaas.Machine) error {
	rawDriver, err := json.Marshal(m.CustomData)
	if err != nil {
		return errors.Wrap(err, "failed to marshal machine data")
	}
	host, err := d.client.NewHost(m.CreationParams["driver"], rawDriver)
	if err != nil {
		return errors.Wrap(err, "failed to initialize host")
	}
	err = host.Driver.Remove()
	if err != nil {
		return errors.Wrap(err, "failed to remove host")
	}
	return d.client.Remove(m.Id)
}

func configureDriver(driver drivers.Driver, driverOpts map[string]interface{}) error {
	opts := &rpcdriver.RPCFlags{Values: driverOpts}
	for _, c := range driver.GetCreateFlags() {
		_, ok := opts.Values[c.String()]
		if !ok {
			opts.Values[c.String()] = c.Default()
			if c.Default() == nil {
				opts.Values[c.String()] = false
			}
		}
	}
	err := driver.SetConfigFromFlags(opts)
	return errors.Wrap(err, "failed to set driver configuration")
}

func copy(src, dst string) error {
	fileSrc, err := ioutil.ReadFile(src)
	if err != nil {
		return errors.Wrapf(err, "failed to read %s", src)
	}
	err = ioutil.WriteFile(dst, fileSrc, 0644)
	return errors.Wrapf(err, "failed to write %s", dst)
}
<commit_msg>iaas/dockermachine: uses engine.DefaultPort instead of hardcoded docker port<commit_after>// Copyright 2016 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package dockermachine

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/docker/machine/libmachine"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/drivers/rpc"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/log"
	"github.com/pkg/errors"
	"github.com/tsuru/tsuru/iaas"
)

var defaultWriter = ioutil.Discard

type DockerMachine struct {
	io.Closer
	client libmachine.API
	config *DockerMachineConfig
	path   string
}

type DockerMachineConfig struct {
	CaPath                 string
	InsecureRegistry       string
	DockerEngineInstallURL string
	OutWriter              io.Writer
	ErrWriter              io.Writer
}

type dockerMachineAPI interface {
	io.Closer
	CreateMachine(string, string, map[string]interface{}) (*iaas.Machine, error)
	DeleteMachine(*iaas.Machine) error
}

func NewDockerMachine(config DockerMachineConfig) (dockerMachineAPI, error) {
	path, err := ioutil.TempDir("", "")
	if err != nil {
		return nil, errors.Wrap(err, "failed to create temp dir")
	}
	if config.CaPath != "" {
		err = copy(filepath.Join(config.CaPath, "ca.pem"), filepath.Join(path, "ca.pem"))
		if err != nil {
			return nil, errors.Wrap(err, "failed to copy ca file")
		}
		err = copy(filepath.Join(config.CaPath, "ca-key.pem"), filepath.Join(path, "ca-key.pem"))
		if err != nil {
			return nil, errors.Wrap(err, "failed to copy ca key file")
		}
	}
	if config.OutWriter != nil {
		log.SetOutWriter(config.OutWriter)
	} else {
		log.SetOutWriter(defaultWriter)
	}
	if config.ErrWriter != nil {
		log.SetOutWriter(config.ErrWriter)
	} else {
		log.SetOutWriter(defaultWriter)
	}
	return &DockerMachine{
		path:   path,
		client: libmachine.NewClient(path, path),
		config: &config,
	}, nil
}

func (d *DockerMachine) Close() error {
	os.RemoveAll(d.path)
	return d.client.Close()
}

func (d *DockerMachine) CreateMachine(name, driver string, params map[string]interface{}) (*iaas.Machine, error) {
	rawDriver, err := json.Marshal(&drivers.BaseDriver{
		MachineName: name,
		StorePath:   d.path,
	})
	if err != nil {
		return nil, errors.Wrap(err, "failed to marshal base driver")
	}
	host, err := d.client.NewHost(driver, rawDriver)
	if err != nil {
		return nil, errors.Wrap(err, "failed to initialize host")
	}
	err = configureDriver(host.Driver, params)
	if err != nil {
		return nil, errors.Wrap(err, "failed to configure driver")
	}
	engineOpts := host.HostOptions.EngineOptions
	if d.config.InsecureRegistry != "" {
		engineOpts.InsecureRegistry = []string{d.config.InsecureRegistry}
	}
	if d.config.DockerEngineInstallURL != "" {
		engineOpts.InstallURL = d.config.DockerEngineInstallURL
	}
	err = d.client.Create(host)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create host")
	}
	rawDriver, err = json.Marshal(host.Driver)
	if err != nil {
		return nil, errors.Wrap(err, "failed to marshal host driver")
	}
	var driverData map[string]interface{}
	err = json.Unmarshal(rawDriver, &driverData)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal host driver")
	}
	m := &iaas.Machine{
		Id:         host.Name,
		Port:       engine.DefaultPort,
		Protocol:   "https",
		CustomData: driverData,
	}
	m.Address, err = host.Driver.GetIP()
	if err != nil {
		return m, errors.Wrap(err, "failed to retrive host ip")
	}
	if host.AuthOptions() != nil {
		m.CaCert, err = ioutil.ReadFile(host.AuthOptions().CaCertPath)
		if err != nil {
			return m, errors.Wrap(err, "failed to read host ca cert")
		}
		m.ClientCert, err = ioutil.ReadFile(host.AuthOptions().ClientCertPath)
		if err != nil {
			return m, errors.Wrap(err, "failed to read host client cert")
		}
		m.ClientKey, err = ioutil.ReadFile(host.AuthOptions().ClientKeyPath)
		if err != nil {
			return m, errors.Wrap(err, "failed to read host client key")
		}
	}
	return m, nil
}

func (d *DockerMachine) DeleteMachine(m *iaas.Machine) error {
	rawDriver, err := json.Marshal(m.CustomData)
	if err != nil {
		return errors.Wrap(err, "failed to marshal machine data")
	}
	host, err := d.client.NewHost(m.CreationParams["driver"], rawDriver)
	if err != nil {
		return errors.Wrap(err, "failed to initialize host")
	}
	err = host.Driver.Remove()
	if err != nil {
		return errors.Wrap(err, "failed to remove host")
	}
	return d.client.Remove(m.Id)
}

func configureDriver(driver drivers.Driver, driverOpts map[string]interface{}) error {
	opts := &rpcdriver.RPCFlags{Values: driverOpts}
	for _, c := range driver.GetCreateFlags() {
		_, ok := opts.Values[c.String()]
		if !ok {
			opts.Values[c.String()] = c.Default()
			if c.Default() == nil {
				opts.Values[c.String()] = false
			}
		}
	}
	err := driver.SetConfigFromFlags(opts)
	return errors.Wrap(err, "failed to set driver configuration")
}

func copy(src, dst string) error {
	fileSrc, err := ioutil.ReadFile(src)
	if err != nil {
		return errors.Wrapf(err, "failed to read %s", src)
	}
	err = ioutil.WriteFile(dst, fileSrc, 0644)
	return errors.Wrapf(err, "failed to write %s", dst)
}
<|endoftext|>


<commit_before>package testing_test

import (
	"bufio"
	"bytes"
	"context"
	"strings"
	"testing"

	"github.com/influxdata/flux"
	"github.com/influxdata/flux/ast"
	"github.com/influxdata/flux/execute"
	"github.com/influxdata/flux/lang"
	"github.com/influxdata/flux/parser"
	"github.com/influxdata/flux/stdlib"

	platform "github.com/influxdata/influxdb"
	"github.com/influxdata/influxdb/cmd/influxd/launcher"
	"github.com/influxdata/influxdb/query"
	itesting "github.com/influxdata/influxdb/query/stdlib/testing"

	_ "github.com/influxdata/flux/stdlib"           // Import the built-in functions
	_ "github.com/influxdata/influxdb/query/stdlib" // Import the stdlib
)

// Default context.
var ctx = context.Background()

func init() {
	flux.FinalizeBuiltIns()
}

func TestFluxEndToEnd(t *testing.T) {
	t.Skip("flaky test: https://github.com/influxdata/influxdb/issues/14193")
	runEndToEnd(t, stdlib.FluxTestPackages)
}
func BenchmarkFluxEndToEnd(b *testing.B) {
	b.Skip("flaky test: https://github.com/influxdata/influxdb/issues/14193")
	benchEndToEnd(b, stdlib.FluxTestPackages)
}

func runEndToEnd(t *testing.T, pkgs []*ast.Package) {
	l := launcher.RunTestLauncherOrFail(t, ctx)
	l.SetupOrFail(t)
	defer l.ShutdownOrFail(t, ctx)
	for _, pkg := range pkgs {
		pkg := pkg.Copy().(*ast.Package)
		name := pkg.Files[0].Name
		t.Run(name, func(t *testing.T) {
			if reason, ok := itesting.FluxEndToEndSkipList[strings.TrimSuffix(name, ".flux")]; ok {
				t.Skip(reason)
			}
			testFlux(t, l, pkg)
		})
	}
}

func benchEndToEnd(b *testing.B, pkgs []*ast.Package) {
	l := launcher.RunTestLauncherOrFail(b, ctx)
	l.SetupOrFail(b)
	defer l.ShutdownOrFail(b, ctx)
	for _, pkg := range pkgs {
		pkg := pkg.Copy().(*ast.Package)
		name := pkg.Files[0].Name
		b.Run(name, func(b *testing.B) {
			if reason, ok := itesting.FluxEndToEndSkipList[strings.TrimSuffix(name, ".flux")]; ok {
				b.Skip(reason)
			}
			b.ResetTimer()
			b.ReportAllocs()
			for i := 0; i < b.N; i++ {
				testFlux(b, l, pkg)
			}
		})
	}
}

var optionsSource = `
import "testing"
import c "csv"

// Options bucket and org are defined dynamically per test

option testing.loadStorage = (csv) => {
	c.from(csv: csv) |> to(bucket: bucket, org: org)
	return from(bucket: bucket)
}
`
var optionsAST *ast.File

func init() {
	pkg := parser.ParseSource(optionsSource)
	if ast.Check(pkg) > 0 {
		panic(ast.GetError(pkg))
	}
	optionsAST = pkg.Files[0]
}

func testFlux(t testing.TB, l *launcher.TestLauncher, pkg *ast.Package) {

	// Query server to ensure write persists.

	b := &platform.Bucket{
		OrgID:           l.Org.ID,
		Name:            t.Name(),
		RetentionPeriod: 0,
	}

	s := l.BucketService()
	if err := s.CreateBucket(context.Background(), b); err != nil {
		t.Fatal(err)
	}

	// Define bucket and org options
	bucketOpt := &ast.OptionStatement{
		Assignment: &ast.VariableAssignment{
			ID:   &ast.Identifier{Name: "bucket"},
			Init: &ast.StringLiteral{Value: b.Name},
		},
	}
	orgOpt := &ast.OptionStatement{
		Assignment: &ast.VariableAssignment{
			ID:   &ast.Identifier{Name: "org"},
			Init: &ast.StringLiteral{Value: l.Org.Name},
		},
	}
	options := optionsAST.Copy().(*ast.File)
	options.Body = append([]ast.Statement{bucketOpt, orgOpt}, options.Body...)

	// Add options to pkg
	pkg.Files = append(pkg.Files, options)

	// Add testing.inspect call to ensure the data is loaded
	inspectCalls := stdlib.TestingInspectCalls(pkg)
	pkg.Files = append(pkg.Files, inspectCalls)

	req := &query.Request{
		OrganizationID: l.Org.ID,
		Compiler:       lang.ASTCompiler{AST: pkg},
	}
	if r, err := l.FluxQueryService().Query(ctx, req); err != nil {
		t.Fatal(err)
	} else {
		for r.More() {
			v := r.Next()
			if err := v.Tables().Do(func(tbl flux.Table) error {
				return nil
			}); err != nil {
				t.Error(err)
			}
		}
	}

	// quirk: our execution engine doesn't guarantee the order of execution for disconnected DAGS
	// so that our function-with-side effects call to `to` may run _after_ the test instead of before.
	// running twice makes sure that `to` happens at least once before we run the test.
	// this time we use a call to `run` so that the assertion error is triggered
	runCalls := stdlib.TestingRunCalls(pkg)
	pkg.Files[len(pkg.Files)-1] = runCalls
	r, err := l.FluxQueryService().Query(ctx, req)
	if err != nil {
		t.Fatal(err)
	}

	for r.More() {
		v := r.Next()
		if err := v.Tables().Do(func(tbl flux.Table) error {
			return nil
		}); err != nil {
			t.Error(err)
		}
	}
	if err := r.Err(); err != nil {
		t.Error(err)
		// Replace the testing.run calls with testing.inspect calls.
		pkg.Files[len(pkg.Files)-1] = inspectCalls
		r, err := l.FluxQueryService().Query(ctx, req)
		if err != nil {
			t.Fatal(err)
		}
		var out bytes.Buffer
		defer func() {
			if t.Failed() {
				scanner := bufio.NewScanner(&out)
				for scanner.Scan() {
					t.Log(scanner.Text())
				}
			}
		}()
		for r.More() {
			v := r.Next()
			err := execute.FormatResult(&out, v)
			if err != nil {
				t.Error(err)
			}
		}
		if err := r.Err(); err != nil {
			t.Error(err)
		}
	}
}
<commit_msg>fix(query/stdlib): remove skip from the flux e2e tests (#14253)<commit_after>package testing_test

import (
	"bufio"
	"bytes"
	"context"
	"strings"
	"testing"

	"github.com/influxdata/flux"
	"github.com/influxdata/flux/ast"
	"github.com/influxdata/flux/execute"
	"github.com/influxdata/flux/lang"
	"github.com/influxdata/flux/parser"
	"github.com/influxdata/flux/stdlib"

	platform "github.com/influxdata/influxdb"
	"github.com/influxdata/influxdb/cmd/influxd/launcher"
	"github.com/influxdata/influxdb/query"
	itesting "github.com/influxdata/influxdb/query/stdlib/testing"

	_ "github.com/influxdata/flux/stdlib"           // Import the built-in functions
	_ "github.com/influxdata/influxdb/query/stdlib" // Import the stdlib
)

// Default context.
var ctx = context.Background()

func init() {
	flux.FinalizeBuiltIns()
}

func TestFluxEndToEnd(t *testing.T) {
	runEndToEnd(t, stdlib.FluxTestPackages)
}
func BenchmarkFluxEndToEnd(b *testing.B) {
	benchEndToEnd(b, stdlib.FluxTestPackages)
}

func runEndToEnd(t *testing.T, pkgs []*ast.Package) {
	l := launcher.RunTestLauncherOrFail(t, ctx)
	l.SetupOrFail(t)
	defer l.ShutdownOrFail(t, ctx)
	for _, pkg := range pkgs {
		pkg := pkg.Copy().(*ast.Package)
		name := pkg.Files[0].Name
		t.Run(name, func(t *testing.T) {
			if reason, ok := itesting.FluxEndToEndSkipList[strings.TrimSuffix(name, ".flux")]; ok {
				t.Skip(reason)
			}
			testFlux(t, l, pkg)
		})
	}
}

func benchEndToEnd(b *testing.B, pkgs []*ast.Package) {
	l := launcher.RunTestLauncherOrFail(b, ctx)
	l.SetupOrFail(b)
	defer l.ShutdownOrFail(b, ctx)
	for _, pkg := range pkgs {
		pkg := pkg.Copy().(*ast.Package)
		name := pkg.Files[0].Name
		b.Run(name, func(b *testing.B) {
			if reason, ok := itesting.FluxEndToEndSkipList[strings.TrimSuffix(name, ".flux")]; ok {
				b.Skip(reason)
			}
			b.ResetTimer()
			b.ReportAllocs()
			for i := 0; i < b.N; i++ {
				testFlux(b, l, pkg)
			}
		})
	}
}

var optionsSource = `
import "testing"
import c "csv"

// Options bucket and org are defined dynamically per test

option testing.loadStorage = (csv) => {
	c.from(csv: csv) |> to(bucket: bucket, org: org)
	return from(bucket: bucket)
}
`
var optionsAST *ast.File

func init() {
	pkg := parser.ParseSource(optionsSource)
	if ast.Check(pkg) > 0 {
		panic(ast.GetError(pkg))
	}
	optionsAST = pkg.Files[0]
}

func testFlux(t testing.TB, l *launcher.TestLauncher, pkg *ast.Package) {

	// Query server to ensure write persists.

	b := &platform.Bucket{
		OrgID:           l.Org.ID,
		Name:            t.Name(),
		RetentionPeriod: 0,
	}

	s := l.BucketService()
	if err := s.CreateBucket(context.Background(), b); err != nil {
		t.Fatal(err)
	}

	// Define bucket and org options
	bucketOpt := &ast.OptionStatement{
		Assignment: &ast.VariableAssignment{
			ID:   &ast.Identifier{Name: "bucket"},
			Init: &ast.StringLiteral{Value: b.Name},
		},
	}
	orgOpt := &ast.OptionStatement{
		Assignment: &ast.VariableAssignment{
			ID:   &ast.Identifier{Name: "org"},
			Init: &ast.StringLiteral{Value: l.Org.Name},
		},
	}
	options := optionsAST.Copy().(*ast.File)
	options.Body = append([]ast.Statement{bucketOpt, orgOpt}, options.Body...)

	// Add options to pkg
	pkg.Files = append(pkg.Files, options)

	// Add testing.inspect call to ensure the data is loaded
	inspectCalls := stdlib.TestingInspectCalls(pkg)
	pkg.Files = append(pkg.Files, inspectCalls)

	req := &query.Request{
		OrganizationID: l.Org.ID,
		Compiler:       lang.ASTCompiler{AST: pkg},
	}
	if r, err := l.FluxQueryService().Query(ctx, req); err != nil {
		t.Fatal(err)
	} else {
		for r.More() {
			v := r.Next()
			if err := v.Tables().Do(func(tbl flux.Table) error {
				return nil
			}); err != nil {
				t.Error(err)
			}
		}
	}

	// quirk: our execution engine doesn't guarantee the order of execution for disconnected DAGS
	// so that our function-with-side effects call to `to` may run _after_ the test instead of before.
	// running twice makes sure that `to` happens at least once before we run the test.
	// this time we use a call to `run` so that the assertion error is triggered
	runCalls := stdlib.TestingRunCalls(pkg)
	pkg.Files[len(pkg.Files)-1] = runCalls
	r, err := l.FluxQueryService().Query(ctx, req)
	if err != nil {
		t.Fatal(err)
	}

	for r.More() {
		v := r.Next()
		if err := v.Tables().Do(func(tbl flux.Table) error {
			return nil
		}); err != nil {
			t.Error(err)
		}
	}
	if err := r.Err(); err != nil {
		t.Error(err)
		// Replace the testing.run calls with testing.inspect calls.
		pkg.Files[len(pkg.Files)-1] = inspectCalls
		r, err := l.FluxQueryService().Query(ctx, req)
		if err != nil {
			t.Fatal(err)
		}
		var out bytes.Buffer
		defer func() {
			if t.Failed() {
				scanner := bufio.NewScanner(&out)
				for scanner.Scan() {
					t.Log(scanner.Text())
				}
			}
		}()
		for r.More() {
			v := r.Next()
			err := execute.FormatResult(&out, v)
			if err != nil {
				t.Error(err)
			}
		}
		if err := r.Err(); err != nil {
			t.Error(err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strconv"

	"github.com/nsf/termbox-go"
)

const (
	headerRows = 1

	pidColumnTitle = "PID"
	pidColumnWidth = 5 // 32768 is the max pid on my system

	userColumnTitle = "USER"
	userColumnWidth = 8

	cpuColumnTitle = "%CPU"
	cpuColumnWidth = 5

	timeColumnTitle = "TIME+"
	timeColumnWidth = 8

	commandColumnTitle = "Command"

	titleFG     = termbox.ColorBlack
	titleBG     = termbox.ColorGreen
	titleSortBG = termbox.ColorCyan
)

type UI struct {
	pm *ProcessMonitor

	width  int
	height int

	start    int
	selected int
}

func NewUI(pm *ProcessMonitor) *UI {
	ui := &UI{
		pm: pm,
	}
	ui.updateTerminalSize()
	return ui
}

func (ui *UI) Draw() {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)

	y := 0
	x := 0

	fg := titleFG
	bg := titleBG

	bg = bgForTitle("pid")
	writeColumn(pidColumnTitle, pidColumnWidth, true, &x, y, fg, bg)

	bg = bgForTitle("user")
	writeColumn(userColumnTitle, userColumnWidth, false, &x, y, fg, bg)

	bg = bgForTitle("cpu")
	writeColumn(cpuColumnTitle, cpuColumnWidth, true, &x, y, fg, bg)

	bg = bgForTitle("time")
	writeColumn(timeColumnTitle, timeColumnWidth, true, &x, y, fg, bg)

	bg = bgForTitle("command")
	writeLastColumn(commandColumnTitle, ui.width, x, y, fg, bg)

	y++

	for i, process := range ui.visibleProcesses() {
		x = 0

		fg = termbox.ColorDefault
		bg = termbox.ColorDefault

		if i == ui.selected {
			fg = termbox.ColorBlack
			bg = termbox.ColorCyan
		}

		// PID
		pidColumn := strconv.Itoa(process.PID)
		writeColumn(pidColumn, pidColumnWidth, true, &x, y, fg, bg)

		// User
		maxUserLen := len(process.User.Username)
		if maxUserLen > userColumnWidth {
			maxUserLen = userColumnWidth
		}
		userColumn := process.User.Username[0:maxUserLen]
		writeColumn(userColumn, userColumnWidth, false, &x, y, fg, bg)

		// CPU Percentage
		totalUsage := float64(ui.pm.CPUTimeDiff)
		userUsage := 100 * float64(process.UtimeDiff) / totalUsage
		systemUsage := 100 * float64(process.StimeDiff) / totalUsage
		cpuColumn := fmt.Sprintf("%.1f", (userUsage+systemUsage)*float64(ui.pm.NumCPUs))
		writeColumn(cpuColumn, cpuColumnWidth, true, &x, y, fg, bg)

		// Time
		hertz := uint64(100)
		// TODO: this has only been tested on my Ubuntu 14.04 system that has
		// a CLK_TICK of 100. Test on other configurations. (getconf CLK_TICK)
		totalJiffies := process.Utime + process.Stime
		totalSeconds := totalJiffies / hertz

		minutes := totalSeconds / 60
		seconds := totalSeconds % 60
		hundredths := totalJiffies % hertz

		// FIXME: this won't be pretty when minutes gets big, maybe format hours?
		timeColumn := fmt.Sprintf("%d:%02d:%02d", minutes, seconds, hundredths)
		writeColumn(timeColumn, timeColumnWidth, true, &x, y, fg, bg)

		// Command
		writeLastColumn(process.Command, ui.width, x, y, fg, bg)

		y++
	}

	termbox.Flush()
}

func (ui *UI) HandleResize() {
	ui.updateTerminalSize()
}

func (ui *UI) HandleDown() {
	if ui.shouldScrollDown() {
		ui.start++
		return
	}

	if !ui.bottomSelected() {
		ui.selected++
	}
}

func (ui *UI) HandleUp() {
	if ui.shouldScrollUp() {
		ui.start--
		return
	}

	if !ui.topSelected() {
		ui.selected--
	}
}

func (ui *UI) shouldScrollDown() bool {
	return ui.bottomSelected() && ui.moreProcessesToShow()
}

func (ui *UI) shouldScrollUp() bool {
	return ui.topSelected() && ui.start > 0
}

func (ui *UI) bottomSelected() bool {
	bottom := len(ui.pm.List) - 1
	if len(ui.pm.List) > ui.numProcessesOnScreen() {
		// Not all processes fit on the same screen
		bottom = ui.numProcessesOnScreen() - 1
	}
	return ui.selected == bottom
}

func (ui *UI) topSelected() bool {
	return ui.selected == 0
}

func (ui *UI) moreProcessesToShow() bool {
	return len(ui.pm.List)-ui.start > ui.numProcessesOnScreen()
}

func (ui *UI) numProcessesOnScreen() int {
	return ui.height - headerRows
}

func (ui *UI) updateTerminalSize() {
	ui.width, ui.height = termbox.Size()
}

func (ui *UI) visibleProcesses() []*Process {
	// Maybe all processes will fit on the same screen
	end := len(ui.pm.List)

	// Maybe they won't
	if end > ui.numProcessesOnScreen() {
		end = ui.start + ui.numProcessesOnScreen()

		// Maybe we need to scroll up because some process(es) died
		if end > len(ui.pm.List) {
			diff := end - len(ui.pm.List)
			ui.start -= diff
			end -= diff
		}
	}

	// When bottom process is selected and a process dies, update selected
	// to the new bottom process.
	if ui.selected >= end {
		ui.selected = end - 1
	}

	return ui.pm.List[ui.start:end]
}

func writeColumn(s string, columnWidth int, rightAlign bool, x *int, y int, fg, bg termbox.Attribute) {
	if rightAlign {
		for i := 0; i < columnWidth-len(s); i++ {
			termbox.SetCell(*x, y, ' ', fg, bg)
			*x++
		}
	}

	for _, ch := range s {
		termbox.SetCell(*x, y, ch, fg, bg)
		*x++
	}

	if !rightAlign {
		for i := 0; i < columnWidth-len(s); i++ {
			termbox.SetCell(*x, y, ' ', fg, bg)
			*x++
		}
	}

	// Space to separate columns
	termbox.SetCell(*x, y, ' ', fg, bg)
	*x++
}

func writeLastColumn(s string, terminalWidth, x, y int, fg, bg termbox.Attribute) {
	for _, ch := range s {
		termbox.SetCell(x, y, ch, fg, bg)
		x++
	}

	for x < terminalWidth {
		termbox.SetCell(x, y, ' ', fg, bg)
		x++
	}
}

func bgForTitle(title string) termbox.Attribute {
	if title == sortFlag {
		return titleSortBG
	}
	return titleBG
}
<commit_msg>Move mutation into separate functions<commit_after>package main

import (
	"fmt"
	"strconv"

	"github.com/nsf/termbox-go"
)

const (
	headerRows = 1

	pidColumnTitle = "PID"
	pidColumnWidth = 5 // 32768 is the max pid on my system

	userColumnTitle = "USER"
	userColumnWidth = 8

	cpuColumnTitle = "%CPU"
	cpuColumnWidth = 5

	timeColumnTitle = "TIME+"
	timeColumnWidth = 8

	commandColumnTitle = "Command"

	titleFG     = termbox.ColorBlack
	titleBG     = termbox.ColorGreen
	titleSortBG = termbox.ColorCyan
)

type UI struct {
	pm *ProcessMonitor

	width  int
	height int

	start    int
	selected int
}

func NewUI(pm *ProcessMonitor) *UI {
	ui := &UI{
		pm: pm,
	}
	ui.updateTerminalSize()
	return ui
}

func (ui *UI) Draw() {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)

	y := 0
	x := 0

	fg := titleFG
	bg := titleBG

	bg = bgForTitle("pid")
	writeColumn(pidColumnTitle, pidColumnWidth, true, &x, y, fg, bg)

	bg = bgForTitle("user")
	writeColumn(userColumnTitle, userColumnWidth, false, &x, y, fg, bg)

	bg = bgForTitle("cpu")
	writeColumn(cpuColumnTitle, cpuColumnWidth, true, &x, y, fg, bg)

	bg = bgForTitle("time")
	writeColumn(timeColumnTitle, timeColumnWidth, true, &x, y, fg, bg)

	bg = bgForTitle("command")
	writeLastColumn(commandColumnTitle, ui.width, x, y, fg, bg)

	y++

	for i, process := range ui.visibleProcesses() {
		x = 0

		fg = termbox.ColorDefault
		bg = termbox.ColorDefault

		if i == ui.selected {
			fg = termbox.ColorBlack
			bg = termbox.ColorCyan
		}

		// PID
		pidColumn := strconv.Itoa(process.PID)
		writeColumn(pidColumn, pidColumnWidth, true, &x, y, fg, bg)

		// User
		maxUserLen := len(process.User.Username)
		if maxUserLen > userColumnWidth {
			maxUserLen = userColumnWidth
		}
		userColumn := process.User.Username[0:maxUserLen]
		writeColumn(userColumn, userColumnWidth, false, &x, y, fg, bg)

		// CPU Percentage
		totalUsage := float64(ui.pm.CPUTimeDiff)
		userUsage := 100 * float64(process.UtimeDiff) / totalUsage
		systemUsage := 100 * float64(process.StimeDiff) / totalUsage
		cpuColumn := fmt.Sprintf("%.1f", (userUsage+systemUsage)*float64(ui.pm.NumCPUs))
		writeColumn(cpuColumn, cpuColumnWidth, true, &x, y, fg, bg)

		// Time
		hertz := uint64(100)
		// TODO: this has only been tested on my Ubuntu 14.04 system that has
		// a CLK_TICK of 100. Test on other configurations. (getconf CLK_TICK)
		totalJiffies := process.Utime + process.Stime
		totalSeconds := totalJiffies / hertz

		minutes := totalSeconds / 60
		seconds := totalSeconds % 60
		hundredths := totalJiffies % hertz

		// FIXME: this won't be pretty when minutes gets big, maybe format hours?
		timeColumn := fmt.Sprintf("%d:%02d:%02d", minutes, seconds, hundredths)
		writeColumn(timeColumn, timeColumnWidth, true, &x, y, fg, bg)

		// Command
		writeLastColumn(process.Command, ui.width, x, y, fg, bg)

		y++
	}

	termbox.Flush()
}

func (ui *UI) HandleResize() {
	ui.updateTerminalSize()
}

func (ui *UI) HandleDown() {
	if ui.shouldScrollDown() {
		ui.scrollDown()
		return
	}

	if !ui.bottomSelected() {
		ui.down()
	}
}

func (ui *UI) HandleUp() {
	if ui.shouldScrollUp() {
		ui.scrollUp()
		return
	}

	if !ui.topSelected() {
		ui.up()
	}
}

func (ui *UI) down() {
	ui.selected++
}

func (ui *UI) up() {
	ui.selected--
}

func (ui *UI) scrollDown() {
	ui.start++
}

func (ui *UI) scrollUp() {
	ui.start--
}

func (ui *UI) shouldScrollDown() bool {
	return ui.bottomSelected() && ui.moreProcessesToShow()
}

func (ui *UI) shouldScrollUp() bool {
	return ui.topSelected() && ui.start > 0
}

func (ui *UI) bottomSelected() bool {
	bottom := len(ui.pm.List) - 1
	if len(ui.pm.List) > ui.numProcessesOnScreen() {
		// Not all processes fit on the same screen
		bottom = ui.numProcessesOnScreen() - 1
	}
	return ui.selected == bottom
}

func (ui *UI) topSelected() bool {
	return ui.selected == 0
}

func (ui *UI) moreProcessesToShow() bool {
	return len(ui.pm.List)-ui.start > ui.numProcessesOnScreen()
}

func (ui *UI) numProcessesOnScreen() int {
	return ui.height - headerRows
}

func (ui *UI) updateTerminalSize() {
	ui.width, ui.height = termbox.Size()
}

func (ui *UI) visibleProcesses() []*Process {
	// Maybe all processes will fit on the same screen
	end := len(ui.pm.List)

	// Maybe they won't
	if end > ui.numProcessesOnScreen() {
		end = ui.start + ui.numProcessesOnScreen()

		// Maybe we need to scroll up because some process(es) died
		if end > len(ui.pm.List) {
			diff := end - len(ui.pm.List)
			ui.start -= diff
			end -= diff
		}
	}

	// When bottom process is selected and a process dies, update selected
	// to the new bottom process.
	if ui.selected >= end {
		ui.selected = end - 1
	}

	return ui.pm.List[ui.start:end]
}

func writeColumn(s string, columnWidth int, rightAlign bool, x *int, y int, fg, bg termbox.Attribute) {
	if rightAlign {
		for i := 0; i < columnWidth-len(s); i++ {
			termbox.SetCell(*x, y, ' ', fg, bg)
			*x++
		}
	}

	for _, ch := range s {
		termbox.SetCell(*x, y, ch, fg, bg)
		*x++
	}

	if !rightAlign {
		for i := 0; i < columnWidth-len(s); i++ {
			termbox.SetCell(*x, y, ' ', fg, bg)
			*x++
		}
	}

	// Space to separate columns
	termbox.SetCell(*x, y, ' ', fg, bg)
	*x++
}

func writeLastColumn(s string, terminalWidth, x, y int, fg, bg termbox.Attribute) {
	for _, ch := range s {
		termbox.SetCell(x, y, ch, fg, bg)
		x++
	}

	for x < terminalWidth {
		termbox.SetCell(x, y, ' ', fg, bg)
		x++
	}
}

func bgForTitle(title string) termbox.Attribute {
	if title == sortFlag {
		return titleSortBG
	}
	return titleBG
}
<|endoftext|>