<commit_before>package db

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/HouzuoGuo/tiedot/dberr"
)

func ensureMapHasKeys(m map[int]struct{}, keys ...int) bool {
	if len(m) != len(keys) {
		return false
	}
	for _, v := range keys {
		if _, ok := m[v]; !ok {
			return false
		}
	}
	return true
}

func runQuery(query string, col *Col) (map[int]struct{}, error) {
	result := make(map[int]struct{})
	var jq interface{}
	if err := json.Unmarshal([]byte(query), &jq); err != nil {
		fmt.Println(err)
	}
	return result, EvalQuery(jq, col, &result)
}

func TestQuery(t *testing.T) {
	os.RemoveAll(TEST_DATA_DIR)
	defer os.RemoveAll(TEST_DATA_DIR)
	if err := os.MkdirAll(TEST_DATA_DIR, 0700); err != nil {
		t.Fatal(err)
	}
	if err := ioutil.WriteFile(TEST_DATA_DIR+"/number_of_partitions", []byte("2"), 0600); err != nil {
		t.Fatal(err)
	}
	db, err := OpenDB(TEST_DATA_DIR)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	// Prepare collection and index
	if err = db.Create("col"); err != nil {
		t.Fatal(err)
	}
	col := db.Use("col")
	docs := []string{
		`{"a": {"b": [1]}, "c": 1, "d": 1, "f": 1, "g": 1, "special": {"thing": null}, "h": 1}`,
		`{"a": {"b": 1}, "c": [1], "d": 2, "f": 2, "g": 2}`,
		`{"a": [{"b": [2]}], "c": 2, "d": 1, "f": 3, "g": 3, "h": 3}`,
		`{"a": {"b": 3}, "c": [3], "d": 2, "f": 4, "g": 4}`,
		`{"a": {"b": [4]}, "c": 4, "d": 1, "f": 5, "g": 5}`,
		`{"a": [{"b": 5}, {"b": 6}], "c": 4, "d": 1, "f": 5, "g": 5, "h": 2}`,
		`{"a": [{"b": "val1"}, {"b": "val2"}]}`,
		`{"a": [{"b": "val3"}, {"b": ["val4", "val5"]}]}`}
	ids := make([]int, len(docs))
	for i, doc := range docs {
		var jsonDoc map[string]interface{}
		if err := json.Unmarshal([]byte(doc), &jsonDoc); err != nil {
			panic(err)
		}
		if ids[i], err = col.Insert(jsonDoc); err != nil {
			t.Fatal(err)
			return
		}
	}
	q, err := runQuery(`["all"]`, col)
	if err != nil {
		t.Fatal(err)
	}
	col.Index([]string{"a", "b"})
	col.Index([]string{"f"})
	col.Index([]string{"h"})
	col.Index([]string{"special"})
	col.Index([]string{"e"})
	// expand numbers
	q, err = runQuery(`["1", "2", ["3", "4"], "5"]`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, 1, 2, 3, 4, 5) {
		t.Fatal(q)
	}
	// hash scan
	q, err = runQuery(`{"eq": 1, "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[1]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"eq": 5, "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[5]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"eq": 6, "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[5]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"eq": 1, "limit": 1, "in": ["a", "b"]}`, col)
	if err != nil {
		fmt.Println(err)
	}
	if !ensureMapHasKeys(q, ids[1]) && !ensureMapHasKeys(q, ids[0]) {
		t.Fatal(q, ids[1], ids[0])
	}
	// collection scan
	q, err = runQuery(`{"eq": 1, "in": ["c"]}`, col)
	if dberr.Type(err) != dberr.ErrorNeedIndex {
		t.Fatal("Collection scan should not happen")
	}
	// lookup on "special" (null)
	q, err = runQuery(`{"eq": {"thing": null},  "in": ["special"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0]) {
		t.Fatal(q)
	}
	// lookup in list
	q, err = runQuery(`{"eq": "val1",  "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[6]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"eq": "val5",  "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[7]) {
		t.Fatal(q)
	}
	// "e" should not exist
	q, err = runQuery(`{"has": ["e"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q) {
		t.Fatal(q)
	}
	// existence test, hash scan, with limit
	q, err = runQuery(`{"has": ["h"], "limit": 2}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[2]) && !ensureMapHasKeys(q, ids[2], ids[5]) && !ensureMapHasKeys(q, ids[5], ids[0]) {
		t.Fatal(q, ids[0], ids[1], ids[2])
	}
	// existence test with incorrect input
	q, err = runQuery(`{"has": ["c"], "limit": "a"}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingInt {
		t.Fatal(err)
	}
	// existence test, collection scan & PK
	q, err = runQuery(`{"has": ["c"], "limit": 2}`, col)
	if dberr.Type(err) != dberr.ErrorNeedIndex {
		t.Fatal("Existence test should return error")
	}
	q, err = runQuery(`{"has": ["@id"], "limit": 2}`, col)
	if dberr.Type(err) != dberr.ErrorNeedIndex {
		t.Fatal("Existence test should return error")
	}
	// int range scan with incorrect input
	q, err = runQuery(`{"int-from": "a", "int-to": 4, "in": ["f"], "limit": 1}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingInt {
		t.Fatal(err)
	}
	q, err = runQuery(`{"int-from": 1, "int-to": "a", "in": ["f"], "limit": 1}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingInt {
		t.Fatal(err)
	}
	q, err = runQuery(`{"int-from": 1, "int-to": 2, "in": ["f"], "limit": "a"}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingInt {
		t.Fatal(err)
	}
	// int range scan
	q, err = runQuery(`{"int-from": 2, "int-to": 4, "in": ["f"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[1], ids[2], ids[3]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"int-from": 2, "int-to": 4, "in": ["f"], "limit": 2}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[1], ids[2]) {
		t.Fatal(q, ids[1], ids[2])
	}
	// int hash scan using reversed range and limit
	q, err = runQuery(`{"int-from": 10, "int-to": 0, "in": ["f"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[5], ids[4], ids[3], ids[2], ids[1], ids[0]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"int-from": 10, "int-to": 0, "in": ["f"], "limit": 2}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[5], ids[4]) {
		t.Fatal(q)
	}
	// all documents
	q, err = runQuery(`"all"`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[1], ids[2], ids[3], ids[4], ids[5], ids[6], ids[7]) {
		t.Fatal(q)
	}
	// union
	col.Index([]string{"c"})
	q, err = runQuery(`[{"eq": 4, "limit": 1, "in": ["a", "b"]}, {"eq": 1, "limit": 1, "in": ["c"]}]`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[4]) && !ensureMapHasKeys(q, ids[1], ids[4]) {
		t.Fatal(q)
	}
	// intersection
	col.Index([]string{"d"})
	q, err = runQuery(`{"n": [{"eq": 2, "in": ["d"]}, "all"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[1], ids[3]) {
		t.Fatal(q)
	}
	// intersection with incorrect input
	q, err = runQuery(`{"c": null}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingSubQuery {
		t.Fatal(err)
	}
	// complement
	q, err = runQuery(`{"c": [{"eq": 4,  "in": ["c"]}, {"eq": 2, "in": ["d"]}, "all"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[2], ids[6], ids[7]) {
		t.Fatal(q)
	}
	// complement with incorrect input
	q, err = runQuery(`{"c": null}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingSubQuery {
		t.Fatal(err)
	}
	// union of intersection
	q, err = runQuery(`[{"n": [{"eq": 3, "in": ["c"]}]}, {"n": [{"eq": 2, "in": ["c"]}]}]`, col)
	if !ensureMapHasKeys(q, ids[2], ids[3]) {
		t.Fatal(q)
	}
	// union of complement
	q, err = runQuery(`[{"c": [{"eq": 3, "in": ["c"]}]}, {"c": [{"eq": 2, "in": ["c"]}]}]`, col)
	if !ensureMapHasKeys(q, ids[2], ids[3]) {
		t.Fatal(q)
	}
	// union of complement of intersection
	q, err = runQuery(`[{"c": [{"n": [{"eq": 1, "in": ["d"]},{"eq": 1, "in": ["c"]}]},{"eq": 1, "in": ["d"]}]},{"eq": 2, "in": ["c"]}]`, col)
	if !ensureMapHasKeys(q, ids[2], ids[4], ids[5]) {
		t.Fatal(q)
	}
}
<commit_msg>Remove unused variable 'err' (but assigned)<commit_after>package db

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/HouzuoGuo/tiedot/dberr"
)

func ensureMapHasKeys(m map[int]struct{}, keys ...int) bool {
	if len(m) != len(keys) {
		return false
	}
	for _, v := range keys {
		if _, ok := m[v]; !ok {
			return false
		}
	}
	return true
}

func runQuery(query string, col *Col) (map[int]struct{}, error) {
	result := make(map[int]struct{})
	var jq interface{}
	if err := json.Unmarshal([]byte(query), &jq); err != nil {
		fmt.Println(err)
	}
	return result, EvalQuery(jq, col, &result)
}

func TestQuery(t *testing.T) {
	os.RemoveAll(TEST_DATA_DIR)
	defer os.RemoveAll(TEST_DATA_DIR)
	if err := os.MkdirAll(TEST_DATA_DIR, 0700); err != nil {
		t.Fatal(err)
	}
	if err := ioutil.WriteFile(TEST_DATA_DIR+"/number_of_partitions", []byte("2"), 0600); err != nil {
		t.Fatal(err)
	}
	db, err := OpenDB(TEST_DATA_DIR)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	// Prepare collection and index
	if err = db.Create("col"); err != nil {
		t.Fatal(err)
	}
	col := db.Use("col")
	docs := []string{
		`{"a": {"b": [1]}, "c": 1, "d": 1, "f": 1, "g": 1, "special": {"thing": null}, "h": 1}`,
		`{"a": {"b": 1}, "c": [1], "d": 2, "f": 2, "g": 2}`,
		`{"a": [{"b": [2]}], "c": 2, "d": 1, "f": 3, "g": 3, "h": 3}`,
		`{"a": {"b": 3}, "c": [3], "d": 2, "f": 4, "g": 4}`,
		`{"a": {"b": [4]}, "c": 4, "d": 1, "f": 5, "g": 5}`,
		`{"a": [{"b": 5}, {"b": 6}], "c": 4, "d": 1, "f": 5, "g": 5, "h": 2}`,
		`{"a": [{"b": "val1"}, {"b": "val2"}]}`,
		`{"a": [{"b": "val3"}, {"b": ["val4", "val5"]}]}`}
	ids := make([]int, len(docs))
	for i, doc := range docs {
		var jsonDoc map[string]interface{}
		if err := json.Unmarshal([]byte(doc), &jsonDoc); err != nil {
			panic(err)
		}
		if ids[i], err = col.Insert(jsonDoc); err != nil {
			t.Fatal(err)
			return
		}
	}
	q, err := runQuery(`["all"]`, col)
	if err != nil {
		t.Fatal(err)
	}
	col.Index([]string{"a", "b"})
	col.Index([]string{"f"})
	col.Index([]string{"h"})
	col.Index([]string{"special"})
	col.Index([]string{"e"})
	// expand numbers
	q, err = runQuery(`["1", "2", ["3", "4"], "5"]`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, 1, 2, 3, 4, 5) {
		t.Fatal(q)
	}
	// hash scan
	q, err = runQuery(`{"eq": 1, "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[1]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"eq": 5, "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[5]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"eq": 6, "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[5]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"eq": 1, "limit": 1, "in": ["a", "b"]}`, col)
	if err != nil {
		fmt.Println(err)
	}
	if !ensureMapHasKeys(q, ids[1]) && !ensureMapHasKeys(q, ids[0]) {
		t.Fatal(q, ids[1], ids[0])
	}
	// collection scan
	q, err = runQuery(`{"eq": 1, "in": ["c"]}`, col)
	if dberr.Type(err) != dberr.ErrorNeedIndex {
		t.Fatal("Collection scan should not happen")
	}
	// lookup on "special" (null)
	q, err = runQuery(`{"eq": {"thing": null},  "in": ["special"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0]) {
		t.Fatal(q)
	}
	// lookup in list
	q, err = runQuery(`{"eq": "val1",  "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[6]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"eq": "val5",  "in": ["a", "b"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[7]) {
		t.Fatal(q)
	}
	// "e" should not exist
	q, err = runQuery(`{"has": ["e"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q) {
		t.Fatal(q)
	}
	// existence test, hash scan, with limit
	q, err = runQuery(`{"has": ["h"], "limit": 2}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[2]) && !ensureMapHasKeys(q, ids[2], ids[5]) && !ensureMapHasKeys(q, ids[5], ids[0]) {
		t.Fatal(q, ids[0], ids[1], ids[2])
	}
	// existence test with incorrect input
	q, err = runQuery(`{"has": ["c"], "limit": "a"}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingInt {
		t.Fatal(err)
	}
	// existence test, collection scan & PK
	q, err = runQuery(`{"has": ["c"], "limit": 2}`, col)
	if dberr.Type(err) != dberr.ErrorNeedIndex {
		t.Fatal("Existence test should return error")
	}
	q, err = runQuery(`{"has": ["@id"], "limit": 2}`, col)
	if dberr.Type(err) != dberr.ErrorNeedIndex {
		t.Fatal("Existence test should return error")
	}
	// int range scan with incorrect input
	q, err = runQuery(`{"int-from": "a", "int-to": 4, "in": ["f"], "limit": 1}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingInt {
		t.Fatal(err)
	}
	q, err = runQuery(`{"int-from": 1, "int-to": "a", "in": ["f"], "limit": 1}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingInt {
		t.Fatal(err)
	}
	q, err = runQuery(`{"int-from": 1, "int-to": 2, "in": ["f"], "limit": "a"}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingInt {
		t.Fatal(err)
	}
	// int range scan
	q, err = runQuery(`{"int-from": 2, "int-to": 4, "in": ["f"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[1], ids[2], ids[3]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"int-from": 2, "int-to": 4, "in": ["f"], "limit": 2}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[1], ids[2]) {
		t.Fatal(q, ids[1], ids[2])
	}
	// int hash scan using reversed range and limit
	q, err = runQuery(`{"int-from": 10, "int-to": 0, "in": ["f"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[5], ids[4], ids[3], ids[2], ids[1], ids[0]) {
		t.Fatal(q)
	}
	q, err = runQuery(`{"int-from": 10, "int-to": 0, "in": ["f"], "limit": 2}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[5], ids[4]) {
		t.Fatal(q)
	}
	// all documents
	q, err = runQuery(`"all"`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[1], ids[2], ids[3], ids[4], ids[5], ids[6], ids[7]) {
		t.Fatal(q)
	}
	// union
	col.Index([]string{"c"})
	q, err = runQuery(`[{"eq": 4, "limit": 1, "in": ["a", "b"]}, {"eq": 1, "limit": 1, "in": ["c"]}]`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[4]) && !ensureMapHasKeys(q, ids[1], ids[4]) {
		t.Fatal(q)
	}
	// intersection
	col.Index([]string{"d"})
	q, err = runQuery(`{"n": [{"eq": 2, "in": ["d"]}, "all"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[1], ids[3]) {
		t.Fatal(q)
	}
	// intersection with incorrect input
	q, err = runQuery(`{"c": null}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingSubQuery {
		t.Fatal(err)
	}
	// complement
	q, err = runQuery(`{"c": [{"eq": 4,  "in": ["c"]}, {"eq": 2, "in": ["d"]}, "all"]}`, col)
	if err != nil {
		t.Fatal(err)
	}
	if !ensureMapHasKeys(q, ids[0], ids[2], ids[6], ids[7]) {
		t.Fatal(q)
	}
	// complement with incorrect input
	q, err = runQuery(`{"c": null}`, col)
	if dberr.Type(err) != dberr.ErrorExpectingSubQuery {
		t.Fatal(err)
	}
	// union of intersection
	q, _ = runQuery(`[{"n": [{"eq": 3, "in": ["c"]}]}, {"n": [{"eq": 2, "in": ["c"]}]}]`, col)
	if !ensureMapHasKeys(q, ids[2], ids[3]) {
		t.Fatal(q)
	}
	// union of complement
	q, _ = runQuery(`[{"c": [{"eq": 3, "in": ["c"]}]}, {"c": [{"eq": 2, "in": ["c"]}]}]`, col)
	if !ensureMapHasKeys(q, ids[2], ids[3]) {
		t.Fatal(q)
	}
	// union of complement of intersection
	q, _ = runQuery(`[{"c": [{"n": [{"eq": 1, "in": ["d"]},{"eq": 1, "in": ["c"]}]},{"eq": 1, "in": ["d"]}]},{"eq": 2, "in": ["c"]}]`, col)
	if !ensureMapHasKeys(q, ids[2], ids[4], ids[5]) {
		t.Fatal(q)
	}
}
<|endoftext|>


<commit_before>package bender

import (
	"math"
	"sync"
	"time"
	"fmt"
)

type Request struct {
	Rid int64
	Request interface{}
}

type RequestExecutor func(int64, *Request) error

type StartMsg struct {
	Start int64
}

func (m StartMsg) String() string {
	return fmt.Sprintf("StartMsg{start=%d}", m.Start)
}

type EndMsg struct {
	Start, End int64
}

func (m EndMsg) String() string {
	return fmt.Sprintf("EndMsg{start=%d,end=%d}", m.Start, m.End)
}

type WaitMsg struct {
	Wait, Overage int64
}

func (m WaitMsg) String() string {
	return fmt.Sprintf("WaitMsg{wait=%d,overage=%d}", m.Wait, m.Overage)
}

type StartRequestMsg struct {
	Start, Rid int64
}

func (m StartRequestMsg) String() string {
	return fmt.Sprintf("StartRequestMsg{start=%d,requestId=%d}", m.Start, m.Rid)
}

type EndRequestMsg struct {
	Start, End, Rid int64
	Err       error
}

func (m EndRequestMsg) String() string {
	return fmt.Sprintf("EndRequestMsg{start=%d,end=%d,requestId=%d,error=%s}", m.Start, m.End, m.Rid, m.Err)
}

func LoadTestThroughput(intervals chan int64, requests chan *Request, requestExec RequestExecutor) chan interface{} {
	reporter := make(chan interface{})

	go func() {
		start := time.Now().UnixNano()
		reporter <- StartMsg{start}

		var wg sync.WaitGroup
		var overage int64
		for request := range requests {
			overageStart := time.Now().UnixNano()

			wait, more := <-intervals
			if !more {
				break
			}

			adjust := int64(math.Min(float64(wait), float64(overage)))
			wait -= adjust
			overage -= adjust
			reporter <- WaitMsg{wait, overage}
			time.Sleep(time.Duration(wait))

			wg.Add(1)
			go func(req *Request) {
				defer wg.Done()
				reqStart := time.Now().UnixNano()
				reporter <- StartRequestMsg{start, request.Rid}
				err := requestExec(time.Now().UnixNano(), req)
				reporter <- EndRequestMsg{reqStart, time.Now().UnixNano(), request.Rid, err}
			}(request)

			overage += time.Now().UnixNano() - overageStart - wait
		}
		wg.Wait()
		reporter <- EndMsg{start, time.Now().UnixNano()}
		close(reporter)
	}()

	return reporter
}

func LoadTestConcurrency(starts chan int64, requests chan *Request, requestExec RequestExecutor) chan interface{} {
	reporter := make(chan interface{})
	permits := make(chan bool)

	go func() {
		for n := range starts {
			for i := 0; i < n; i++ {
				permits <- true
			}
		}
		close(permits)
	}()

	go func() {
		start := time.Now().UnixNano()
		reporter <- StartMsg{start}

		var wg sync.WaitGroup
		for request := range requests {
			<-permits

			wg.Add(1)
			go func(req *Request) {
				defer func() {
					wg.Done()
					permits <- true
				}()

				reqStart := time.Now().UnixNano()
				reporter <- StartRequestMsg{start, request.Rid}
				err := requestExec(time.Now().UnixNano(), req)
				reporter <- EndRequestMsg{reqStart, time.Now().UnixNano(), request.Rid, err}
			}(request)
		}

		wg.Wait()
		reporter <- EndMsg{start, time.Now().UnixNano()}
		close(reporter)
	}()

	return reporter
}
<commit_msg>Bug fix<commit_after>package bender

import (
	"math"
	"sync"
	"time"
	"fmt"
)

type Request struct {
	Rid int64
	Request interface{}
}

type RequestExecutor func(int64, *Request) error

type StartMsg struct {
	Start int64
}

func (m StartMsg) String() string {
	return fmt.Sprintf("StartMsg{start=%d}", m.Start)
}

type EndMsg struct {
	Start, End int64
}

func (m EndMsg) String() string {
	return fmt.Sprintf("EndMsg{start=%d,end=%d}", m.Start, m.End)
}

type WaitMsg struct {
	Wait, Overage int64
}

func (m WaitMsg) String() string {
	return fmt.Sprintf("WaitMsg{wait=%d,overage=%d}", m.Wait, m.Overage)
}

type StartRequestMsg struct {
	Start, Rid int64
}

func (m StartRequestMsg) String() string {
	return fmt.Sprintf("StartRequestMsg{start=%d,requestId=%d}", m.Start, m.Rid)
}

type EndRequestMsg struct {
	Start, End, Rid int64
	Err       error
}

func (m EndRequestMsg) String() string {
	return fmt.Sprintf("EndRequestMsg{start=%d,end=%d,requestId=%d,error=%s}", m.Start, m.End, m.Rid, m.Err)
}

func LoadTestThroughput(intervals chan int64, requests chan *Request, requestExec RequestExecutor) chan interface{} {
	reporter := make(chan interface{})

	go func() {
		start := time.Now().UnixNano()
		reporter <- StartMsg{start}

		var wg sync.WaitGroup
		var overage int64
		for request := range requests {
			overageStart := time.Now().UnixNano()

			wait, more := <-intervals
			if !more {
				break
			}

			adjust := int64(math.Min(float64(wait), float64(overage)))
			wait -= adjust
			overage -= adjust
			reporter <- WaitMsg{wait, overage}
			time.Sleep(time.Duration(wait))

			wg.Add(1)
			go func(req *Request) {
				defer wg.Done()
				reqStart := time.Now().UnixNano()
				reporter <- StartRequestMsg{start, request.Rid}
				err := requestExec(time.Now().UnixNano(), req)
				reporter <- EndRequestMsg{reqStart, time.Now().UnixNano(), request.Rid, err}
			}(request)

			overage += time.Now().UnixNano() - overageStart - wait
		}
		wg.Wait()
		reporter <- EndMsg{start, time.Now().UnixNano()}
		close(reporter)
	}()

	return reporter
}

func LoadTestConcurrency(starts chan int64, requests chan *Request, requestExec RequestExecutor) chan interface{} {
	reporter := make(chan interface{})
	permits := make(chan bool)

	go func() {
		for n := range starts {
			for i := int64(0); i < n; i++ {
				permits <- true
			}
		}
		close(permits)
	}()

	go func() {
		start := time.Now().UnixNano()
		reporter <- StartMsg{start}

		var wg sync.WaitGroup
		for request := range requests {
			<-permits

			wg.Add(1)
			go func(req *Request) {
				defer func() {
					wg.Done()
					permits <- true
				}()

				reqStart := time.Now().UnixNano()
				reporter <- StartRequestMsg{start, request.Rid}
				err := requestExec(time.Now().UnixNano(), req)
				reporter <- EndRequestMsg{reqStart, time.Now().UnixNano(), request.Rid, err}
			}(request)
		}

		wg.Wait()
		reporter <- EndMsg{start, time.Now().UnixNano()}
		close(reporter)
	}()

	return reporter
}
<|endoftext|>


<commit_before>// Package resolvconf provides utility code to query and update DNS configuration in /etc/resolv.conf
package resolvconf

import (
	"bytes"
	"io/ioutil"
	"regexp"
	"strings"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/pkg/ioutils"
	"github.com/docker/libnetwork/resolvconf/dns"
	"github.com/docker/libnetwork/types"
)

var (
	// Note: the default IPv4 & IPv6 resolvers are set to Google's Public DNS
	defaultIPv4Dns = []string{"nameserver 8.8.8.8", "nameserver 8.8.4.4"}
	defaultIPv6Dns = []string{"nameserver 2001:4860:4860::8888", "nameserver 2001:4860:4860::8844"}
	ipv4NumBlock   = `(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)`
	ipv4Address    = `(` + ipv4NumBlock + `\.){3}` + ipv4NumBlock
	// This is not an IPv6 address verifier as it will accept a super-set of IPv6, and also
	// will *not match* IPv4-Embedded IPv6 Addresses (RFC6052), but that and other variants
	// -- e.g. other link-local types -- either won't work in containers or are unnecessary.
	// For readability and sufficiency for Docker purposes this seemed more reasonable than a
	// 1000+ character regexp with exact and complete IPv6 validation
	ipv6Address = `([0-9A-Fa-f]{0,4}:){2,7}([0-9A-Fa-f]{0,4})(%\w+)?`

	localhostNSRegexp = regexp.MustCompile(`(?m)^nameserver\s+` + dns.IPLocalhost + `\s*\n*`)
	nsIPv6Regexp      = regexp.MustCompile(`(?m)^nameserver\s+` + ipv6Address + `\s*\n*`)
	nsRegexp          = regexp.MustCompile(`^\s*nameserver\s*((` + ipv4Address + `)|(` + ipv6Address + `))\s*$`)
	nsIPv6Regexpmatch = regexp.MustCompile(`^\s*nameserver\s*((` + ipv6Address + `))\s*$`)
	nsIPv4Regexpmatch = regexp.MustCompile(`^\s*nameserver\s*((` + ipv4Address + `))\s*$`)
	searchRegexp      = regexp.MustCompile(`^\s*search\s*(([^\s]+\s*)*)$`)
	optionsRegexp     = regexp.MustCompile(`^\s*options\s*(([^\s]+\s*)*)$`)
)

var lastModified struct {
	sync.Mutex
	sha256   string
	contents []byte
}

// File contains the resolv.conf content and its hash
type File struct {
	Content []byte
	Hash    string
}

// Get returns the contents of /etc/resolv.conf and its hash
func Get() (*File, error) {
	resolv, err := ioutil.ReadFile("/etc/resolv.conf")
	if err != nil {
		return nil, err
	}
	hash, err := ioutils.HashData(bytes.NewReader(resolv))
	if err != nil {
		return nil, err
	}
	return &File{Content: resolv, Hash: hash}, nil
}

// GetSpecific returns the contents of the user specified resolv.conf file and its hash
func GetSpecific(path string) (*File, error) {
	resolv, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	hash, err := ioutils.HashData(bytes.NewReader(resolv))
	if err != nil {
		return nil, err
	}
	return &File{Content: resolv, Hash: hash}, nil
}

// GetIfChanged retrieves the host /etc/resolv.conf file, checks against the last hash
// and, if modified since last check, returns the bytes and new hash.
// This feature is used by the resolv.conf updater for containers
func GetIfChanged() (*File, error) {
	lastModified.Lock()
	defer lastModified.Unlock()

	resolv, err := ioutil.ReadFile("/etc/resolv.conf")
	if err != nil {
		return nil, err
	}
	newHash, err := ioutils.HashData(bytes.NewReader(resolv))
	if err != nil {
		return nil, err
	}
	if lastModified.sha256 != newHash {
		lastModified.sha256 = newHash
		lastModified.contents = resolv
		return &File{Content: resolv, Hash: newHash}, nil
	}
	// nothing changed, so return no data
	return nil, nil
}

// GetLastModified retrieves the last used contents and hash of the host resolv.conf.
// Used by containers updating on restart
func GetLastModified() *File {
	lastModified.Lock()
	defer lastModified.Unlock()

	return &File{Content: lastModified.contents, Hash: lastModified.sha256}
}

// FilterResolvDNS cleans up the config in resolvConf.  It has two main jobs:
// 1. It looks for localhost (127.*|::1) entries in the provided
//    resolv.conf, removing local nameserver entries, and, if the resulting
//    cleaned config has no defined nameservers left, adds default DNS entries
// 2. Given the caller provides the enable/disable state of IPv6, the filter
//    code will remove all IPv6 nameservers if it is not enabled for containers
//
func FilterResolvDNS(resolvConf []byte, ipv6Enabled bool) (*File, error) {
	cleanedResolvConf := localhostNSRegexp.ReplaceAll(resolvConf, []byte{})
	// if IPv6 is not enabled, also clean out any IPv6 address nameserver
	if !ipv6Enabled {
		cleanedResolvConf = nsIPv6Regexp.ReplaceAll(cleanedResolvConf, []byte{})
	}
	// if the resulting resolvConf has no more nameservers defined, add appropriate
	// default DNS servers for IPv4 and (optionally) IPv6
	if len(GetNameservers(cleanedResolvConf, types.IP)) == 0 {
		logrus.Infof("No non-localhost DNS nameservers are left in resolv.conf. Using default external servers: %v", defaultIPv4Dns)
		dns := defaultIPv4Dns
		if ipv6Enabled {
			logrus.Infof("IPv6 enabled; Adding default IPv6 external servers: %v", defaultIPv6Dns)
			dns = append(dns, defaultIPv6Dns...)
		}
		cleanedResolvConf = append(cleanedResolvConf, []byte("\n"+strings.Join(dns, "\n"))...)
	}
	hash, err := ioutils.HashData(bytes.NewReader(cleanedResolvConf))
	if err != nil {
		return nil, err
	}
	return &File{Content: cleanedResolvConf, Hash: hash}, nil
}

// getLines parses input into lines and strips away comments.
func getLines(input []byte, commentMarker []byte) [][]byte {
	lines := bytes.Split(input, []byte("\n"))
	var output [][]byte
	for _, currentLine := range lines {
		var commentIndex = bytes.Index(currentLine, commentMarker)
		if commentIndex == -1 {
			output = append(output, currentLine)
		} else {
			output = append(output, currentLine[:commentIndex])
		}
	}
	return output
}

// GetNameservers returns nameservers (if any) listed in /etc/resolv.conf
func GetNameservers(resolvConf []byte, kind int) []string {
	nameservers := []string{}
	for _, line := range getLines(resolvConf, []byte("#")) {
		var ns [][]byte
		if kind == types.IP {
			ns = nsRegexp.FindSubmatch(line)
		} else if kind == types.IPv4 {
			ns = nsIPv4Regexpmatch.FindSubmatch(line)
		} else if kind == types.IPv6 {
			ns = nsIPv6Regexpmatch.FindSubmatch(line)
		}
		if len(ns) > 0 {
			nameservers = append(nameservers, string(ns[1]))
		}
	}
	return nameservers
}

// GetNameserversAsCIDR returns nameservers (if any) listed in
// /etc/resolv.conf as CIDR blocks (e.g., "1.2.3.4/32")
// This function's output is intended for net.ParseCIDR
func GetNameserversAsCIDR(resolvConf []byte) []string {
	nameservers := []string{}
	for _, nameserver := range GetNameservers(resolvConf, types.IP) {
		nameservers = append(nameservers, nameserver+"/32")
	}
	return nameservers
}

// GetSearchDomains returns search domains (if any) listed in /etc/resolv.conf
// If more than one search line is encountered, only the contents of the last
// one is returned.
func GetSearchDomains(resolvConf []byte) []string {
	domains := []string{}
	for _, line := range getLines(resolvConf, []byte("#")) {
		match := searchRegexp.FindSubmatch(line)
		if match == nil {
			continue
		}
		domains = strings.Fields(string(match[1]))
	}
	return domains
}

// GetOptions returns options (if any) listed in /etc/resolv.conf
// If more than one options line is encountered, only the contents of the last
// one is returned.
func GetOptions(resolvConf []byte) []string {
	options := []string{}
	for _, line := range getLines(resolvConf, []byte("#")) {
		match := optionsRegexp.FindSubmatch(line)
		if match == nil {
			continue
		}
		options = strings.Fields(string(match[1]))
	}
	return options
}

// Build writes a configuration file to path containing a "nameserver" entry
// for every element in dns, a "search" entry for every element in
// dnsSearch, and an "options" entry for every element in dnsOptions.
func Build(path string, dns, dnsSearch, dnsOptions []string) (*File, error) {
	content := bytes.NewBuffer(nil)
	if len(dnsSearch) > 0 {
		if searchString := strings.Join(dnsSearch, " "); strings.Trim(searchString, " ") != "." {
			if _, err := content.WriteString("search " + searchString + "\n"); err != nil {
				return nil, err
			}
		}
	}
	for _, dns := range dns {
		if _, err := content.WriteString("nameserver " + dns + "\n"); err != nil {
			return nil, err
		}
	}
	if len(dnsOptions) > 0 {
		if optsString := strings.Join(dnsOptions, " "); strings.Trim(optsString, " ") != "" {
			if _, err := content.WriteString("options " + optsString + "\n"); err != nil {
				return nil, err
			}
		}
	}

	hash, err := ioutils.HashData(bytes.NewReader(content.Bytes()))
	if err != nil {
		return nil, err
	}

	return &File{Content: content.Bytes(), Hash: hash}, ioutil.WriteFile(path, content.Bytes(), 0644)
}
<commit_msg>Fix GetNameServerAsCIDR<commit_after>// Package resolvconf provides utility code to query and update DNS configuration in /etc/resolv.conf
package resolvconf

import (
	"bytes"
	"io/ioutil"
	"regexp"
	"strings"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/pkg/ioutils"
	"github.com/docker/libnetwork/resolvconf/dns"
	"github.com/docker/libnetwork/types"
)

var (
	// Note: the default IPv4 & IPv6 resolvers are set to Google's Public DNS
	defaultIPv4Dns = []string{"nameserver 8.8.8.8", "nameserver 8.8.4.4"}
	defaultIPv6Dns = []string{"nameserver 2001:4860:4860::8888", "nameserver 2001:4860:4860::8844"}
	ipv4NumBlock   = `(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)`
	ipv4Address    = `(` + ipv4NumBlock + `\.){3}` + ipv4NumBlock
	// This is not an IPv6 address verifier as it will accept a super-set of IPv6, and also
	// will *not match* IPv4-Embedded IPv6 Addresses (RFC6052), but that and other variants
	// -- e.g. other link-local types -- either won't work in containers or are unnecessary.
	// For readability and sufficiency for Docker purposes this seemed more reasonable than a
	// 1000+ character regexp with exact and complete IPv6 validation
	ipv6Address = `([0-9A-Fa-f]{0,4}:){2,7}([0-9A-Fa-f]{0,4})(%\w+)?`

	localhostNSRegexp = regexp.MustCompile(`(?m)^nameserver\s+` + dns.IPLocalhost + `\s*\n*`)
	nsIPv6Regexp      = regexp.MustCompile(`(?m)^nameserver\s+` + ipv6Address + `\s*\n*`)
	nsRegexp          = regexp.MustCompile(`^\s*nameserver\s*((` + ipv4Address + `)|(` + ipv6Address + `))\s*$`)
	nsIPv6Regexpmatch = regexp.MustCompile(`^\s*nameserver\s*((` + ipv6Address + `))\s*$`)
	nsIPv4Regexpmatch = regexp.MustCompile(`^\s*nameserver\s*((` + ipv4Address + `))\s*$`)
	searchRegexp      = regexp.MustCompile(`^\s*search\s*(([^\s]+\s*)*)$`)
	optionsRegexp     = regexp.MustCompile(`^\s*options\s*(([^\s]+\s*)*)$`)
)

var lastModified struct {
	sync.Mutex
	sha256   string
	contents []byte
}

// File contains the resolv.conf content and its hash
type File struct {
	Content []byte
	Hash    string
}

// Get returns the contents of /etc/resolv.conf and its hash
func Get() (*File, error) {
	resolv, err := ioutil.ReadFile("/etc/resolv.conf")
	if err != nil {
		return nil, err
	}
	hash, err := ioutils.HashData(bytes.NewReader(resolv))
	if err != nil {
		return nil, err
	}
	return &File{Content: resolv, Hash: hash}, nil
}

// GetSpecific returns the contents of the user specified resolv.conf file and its hash
func GetSpecific(path string) (*File, error) {
	resolv, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	hash, err := ioutils.HashData(bytes.NewReader(resolv))
	if err != nil {
		return nil, err
	}
	return &File{Content: resolv, Hash: hash}, nil
}

// GetIfChanged retrieves the host /etc/resolv.conf file, checks against the last hash
// and, if modified since last check, returns the bytes and new hash.
// This feature is used by the resolv.conf updater for containers
func GetIfChanged() (*File, error) {
	lastModified.Lock()
	defer lastModified.Unlock()

	resolv, err := ioutil.ReadFile("/etc/resolv.conf")
	if err != nil {
		return nil, err
	}
	newHash, err := ioutils.HashData(bytes.NewReader(resolv))
	if err != nil {
		return nil, err
	}
	if lastModified.sha256 != newHash {
		lastModified.sha256 = newHash
		lastModified.contents = resolv
		return &File{Content: resolv, Hash: newHash}, nil
	}
	// nothing changed, so return no data
	return nil, nil
}

// GetLastModified retrieves the last used contents and hash of the host resolv.conf.
// Used by containers updating on restart
func GetLastModified() *File {
	lastModified.Lock()
	defer lastModified.Unlock()

	return &File{Content: lastModified.contents, Hash: lastModified.sha256}
}

// FilterResolvDNS cleans up the config in resolvConf.  It has two main jobs:
// 1. It looks for localhost (127.*|::1) entries in the provided
//    resolv.conf, removing local nameserver entries, and, if the resulting
//    cleaned config has no defined nameservers left, adds default DNS entries
// 2. Given the caller provides the enable/disable state of IPv6, the filter
//    code will remove all IPv6 nameservers if it is not enabled for containers
//
func FilterResolvDNS(resolvConf []byte, ipv6Enabled bool) (*File, error) {
	cleanedResolvConf := localhostNSRegexp.ReplaceAll(resolvConf, []byte{})
	// if IPv6 is not enabled, also clean out any IPv6 address nameserver
	if !ipv6Enabled {
		cleanedResolvConf = nsIPv6Regexp.ReplaceAll(cleanedResolvConf, []byte{})
	}
	// if the resulting resolvConf has no more nameservers defined, add appropriate
	// default DNS servers for IPv4 and (optionally) IPv6
	if len(GetNameservers(cleanedResolvConf, types.IP)) == 0 {
		logrus.Infof("No non-localhost DNS nameservers are left in resolv.conf. Using default external servers: %v", defaultIPv4Dns)
		dns := defaultIPv4Dns
		if ipv6Enabled {
			logrus.Infof("IPv6 enabled; Adding default IPv6 external servers: %v", defaultIPv6Dns)
			dns = append(dns, defaultIPv6Dns...)
		}
		cleanedResolvConf = append(cleanedResolvConf, []byte("\n"+strings.Join(dns, "\n"))...)
	}
	hash, err := ioutils.HashData(bytes.NewReader(cleanedResolvConf))
	if err != nil {
		return nil, err
	}
	return &File{Content: cleanedResolvConf, Hash: hash}, nil
}

// getLines parses input into lines and strips away comments.
func getLines(input []byte, commentMarker []byte) [][]byte {
	lines := bytes.Split(input, []byte("\n"))
	var output [][]byte
	for _, currentLine := range lines {
		var commentIndex = bytes.Index(currentLine, commentMarker)
		if commentIndex == -1 {
			output = append(output, currentLine)
		} else {
			output = append(output, currentLine[:commentIndex])
		}
	}
	return output
}

// GetNameservers returns nameservers (if any) listed in /etc/resolv.conf
func GetNameservers(resolvConf []byte, kind int) []string {
	nameservers := []string{}
	for _, line := range getLines(resolvConf, []byte("#")) {
		var ns [][]byte
		if kind == types.IP {
			ns = nsRegexp.FindSubmatch(line)
		} else if kind == types.IPv4 {
			ns = nsIPv4Regexpmatch.FindSubmatch(line)
		} else if kind == types.IPv6 {
			ns = nsIPv6Regexpmatch.FindSubmatch(line)
		}
		if len(ns) > 0 {
			nameservers = append(nameservers, string(ns[1]))
		}
	}
	return nameservers
}

// GetNameserversAsCIDR returns nameservers (if any) listed in
// /etc/resolv.conf as CIDR blocks (e.g., "1.2.3.4/32")
// This function's output is intended for net.ParseCIDR
func GetNameserversAsCIDR(resolvConf []byte) []string {
	nameservers := []string{}
	for _, nameserver := range GetNameservers(resolvConf, types.IP) {
		var address string
		// If IPv6, strip zone if present
		if strings.Contains(nameserver, ":") {
			address = strings.Split(nameserver, "%")[0] + "/128"
		} else {
			address = nameserver + "/32"
		}
		nameservers = append(nameservers, address)
	}
	return nameservers
}

// GetSearchDomains returns search domains (if any) listed in /etc/resolv.conf
// If more than one search line is encountered, only the contents of the last
// one is returned.
func GetSearchDomains(resolvConf []byte) []string {
	domains := []string{}
	for _, line := range getLines(resolvConf, []byte("#")) {
		match := searchRegexp.FindSubmatch(line)
		if match == nil {
			continue
		}
		domains = strings.Fields(string(match[1]))
	}
	return domains
}

// GetOptions returns options (if any) listed in /etc/resolv.conf
// If more than one options line is encountered, only the contents of the last
// one is returned.
func GetOptions(resolvConf []byte) []string {
	options := []string{}
	for _, line := range getLines(resolvConf, []byte("#")) {
		match := optionsRegexp.FindSubmatch(line)
		if match == nil {
			continue
		}
		options = strings.Fields(string(match[1]))
	}
	return options
}

// Build writes a configuration file to path containing a "nameserver" entry
// for every element in dns, a "search" entry for every element in
// dnsSearch, and an "options" entry for every element in dnsOptions.
func Build(path string, dns, dnsSearch, dnsOptions []string) (*File, error) {
	content := bytes.NewBuffer(nil)
	if len(dnsSearch) > 0 {
		if searchString := strings.Join(dnsSearch, " "); strings.Trim(searchString, " ") != "." {
			if _, err := content.WriteString("search " + searchString + "\n"); err != nil {
				return nil, err
			}
		}
	}
	for _, dns := range dns {
		if _, err := content.WriteString("nameserver " + dns + "\n"); err != nil {
			return nil, err
		}
	}
	if len(dnsOptions) > 0 {
		if optsString := strings.Join(dnsOptions, " "); strings.Trim(optsString, " ") != "" {
			if _, err := content.WriteString("options " + optsString + "\n"); err != nil {
				return nil, err
			}
		}
	}

	hash, err := ioutils.HashData(bytes.NewReader(content.Bytes()))
	if err != nil {
		return nil, err
	}

	return &File{Content: content.Bytes(), Hash: hash}, ioutil.WriteFile(path, content.Bytes(), 0644)
}
<|endoftext|>


<commit_before>package filer2

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"google.golang.org/grpc"

	"github.com/karlseguin/ccache"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
)

const PaginationSize = 1024 * 256

var (
	OS_UID = uint32(os.Getuid())
	OS_GID = uint32(os.Getgid())
)

type Filer struct {
	store              *FilerStoreWrapper
	directoryCache     *ccache.Cache
	MasterClient       *wdclient.MasterClient
	fileIdDeletionChan chan string
	GrpcDialOption     grpc.DialOption
}

func NewFiler(masters []string, grpcDialOption grpc.DialOption) *Filer {
	f := &Filer{
		directoryCache:     ccache.New(ccache.Configure().MaxSize(1000).ItemsToPrune(100)),
		MasterClient:       wdclient.NewMasterClient(context.Background(), grpcDialOption, "filer", masters),
		fileIdDeletionChan: make(chan string, PaginationSize),
		GrpcDialOption:     grpcDialOption,
	}

	go f.loopProcessingDeletion()

	return f
}

func (f *Filer) SetStore(store FilerStore) {
	f.store = NewFilerStoreWrapper(store)
}

func (f *Filer) DisableDirectoryCache() {
	f.directoryCache = nil
}

func (fs *Filer) GetMaster() string {
	return fs.MasterClient.GetMaster()
}

func (fs *Filer) KeepConnectedToMaster() {
	fs.MasterClient.KeepConnectedToMaster()
}

func (f *Filer) BeginTransaction(ctx context.Context) (context.Context, error) {
	return f.store.BeginTransaction(ctx)
}

func (f *Filer) CommitTransaction(ctx context.Context) error {
	return f.store.CommitTransaction(ctx)
}

func (f *Filer) RollbackTransaction(ctx context.Context) error {
	return f.store.RollbackTransaction(ctx)
}

func (f *Filer) CreateEntry(ctx context.Context, entry *Entry, o_excl bool) error {

	glog.V(4).Infof("CreateEntry create %s", entry.FullPath)

	if string(entry.FullPath) == "/" {
		return nil
	}

	dirParts := strings.Split(string(entry.FullPath), "/")

	// fmt.Printf("directory parts: %+v\n", dirParts)

	var lastDirectoryEntry *Entry

	for i := 1; i < len(dirParts); i++ {
		dirPath := "/" + filepath.ToSlash(filepath.Join(dirParts[:i]...))
		// fmt.Printf("%d directory: %+v\n", i, dirPath)

		// first check local cache
		dirEntry := f.cacheGetDirectory(dirPath)

		// not found, check the store directly
		if dirEntry == nil {
			glog.V(4).Infof("find uncached directory: %s", dirPath)
			dirEntry, _ = f.FindEntry(ctx, FullPath(dirPath))
		} else {
			// glog.V(4).Infof("found cached directory: %s", dirPath)
		}

		// no such existing directory
		if dirEntry == nil {

			// create the directory
			now := time.Now()

			dirEntry = &Entry{
				FullPath: FullPath(dirPath),
				Attr: Attr{
					Mtime:  now,
					Crtime: now,
					Mode:   os.ModeDir | 0770,
					Uid:    entry.Uid,
					Gid:    entry.Gid,
				},
			}

			glog.V(2).Infof("create directory: %s %v", dirPath, dirEntry.Mode)
			mkdirErr := f.store.InsertEntry(ctx, dirEntry)
			if mkdirErr != nil {
				if _, err := f.FindEntry(ctx, FullPath(dirPath)); err == ErrNotFound {
					glog.V(3).Infof("mkdir %s: %v", dirPath, mkdirErr)
					return fmt.Errorf("mkdir %s: %v", dirPath, mkdirErr)
				}
			} else {
				f.NotifyUpdateEvent(nil, dirEntry, false)
			}

		} else if !dirEntry.IsDirectory() {
			glog.Errorf("CreateEntry %s: %s should be a directory", entry.FullPath, dirPath)
			return fmt.Errorf("%s is a file", dirPath)
		}

		// cache the directory entry
		f.cacheSetDirectory(dirPath, dirEntry, i)

		// remember the direct parent directory entry
		if i == len(dirParts)-1 {
			lastDirectoryEntry = dirEntry
		}

	}

	if lastDirectoryEntry == nil {
		glog.Errorf("CreateEntry %s: lastDirectoryEntry is nil", entry.FullPath)
		return fmt.Errorf("parent folder not found: %v", entry.FullPath)
	}

	/*
		if !hasWritePermission(lastDirectoryEntry, entry) {
			glog.V(0).Infof("directory %s: %v, entry: uid=%d gid=%d",
				lastDirectoryEntry.FullPath, lastDirectoryEntry.Attr, entry.Uid, entry.Gid)
			return fmt.Errorf("no write permission in folder %v", lastDirectoryEntry.FullPath)
		}
	*/

	oldEntry, _ := f.FindEntry(ctx, entry.FullPath)

	if oldEntry == nil {
		if err := f.store.InsertEntry(ctx, entry); err != nil {
			glog.Errorf("insert entry %s: %v", entry.FullPath, err)
			return fmt.Errorf("insert entry %s: %v", entry.FullPath, err)
		}
	} else {
		if o_excl {
			glog.V(3).Infof("EEXIST: entry %s already exists", entry.FullPath)
			return fmt.Errorf("EEXIST: entry %s already exists", entry.FullPath)
		}
		if err := f.UpdateEntry(ctx, oldEntry, entry); err != nil {
			glog.Errorf("update entry %s: %v", entry.FullPath, err)
			return fmt.Errorf("update entry %s: %v", entry.FullPath, err)
		}
	}

	glog.V(4).Infof("CreateEntry %s: created", entry.FullPath)

	f.NotifyUpdateEvent(oldEntry, entry, true)

	f.deleteChunksIfNotNew(oldEntry, entry)

	return nil
}

func (f *Filer) UpdateEntry(ctx context.Context, oldEntry, entry *Entry) (err error) {
	if oldEntry != nil {
		if oldEntry.IsDirectory() && !entry.IsDirectory() {
			glog.Errorf("existing %s is a directory", entry.FullPath)
			return fmt.Errorf("existing %s is a directory", entry.FullPath)
		}
		if !oldEntry.IsDirectory() && entry.IsDirectory() {
			glog.Errorf("existing %s is a file", entry.FullPath)
			return fmt.Errorf("existing %s is a file", entry.FullPath)
		}
	}
	return f.store.UpdateEntry(ctx, entry)
}

func (f *Filer) FindEntry(ctx context.Context, p FullPath) (entry *Entry, err error) {

	now := time.Now()

	if string(p) == "/" {
		return &Entry{
			FullPath: p,
			Attr: Attr{
				Mtime:  now,
				Crtime: now,
				Mode:   os.ModeDir | 0755,
				Uid:    OS_UID,
				Gid:    OS_GID,
			},
		}, nil
	}
	return f.store.FindEntry(ctx, p)
}

func (f *Filer) ListDirectoryEntries(ctx context.Context, p FullPath, startFileName string, inclusive bool, limit int) ([]*Entry, error) {
	if strings.HasSuffix(string(p), "/") && len(p) > 1 {
		p = p[0 : len(p)-1]
	}
	return f.store.ListDirectoryEntries(ctx, p, startFileName, inclusive, limit)
}

func (f *Filer) cacheDelDirectory(dirpath string) {

	if dirpath == "/" {
		return
	}

	if f.directoryCache == nil {
		return
	}
	f.directoryCache.Delete(dirpath)
	return
}

func (f *Filer) cacheGetDirectory(dirpath string) *Entry {

	if f.directoryCache == nil {
		return nil
	}
	item := f.directoryCache.Get(dirpath)
	if item == nil {
		return nil
	}
	return item.Value().(*Entry)
}

func (f *Filer) cacheSetDirectory(dirpath string, dirEntry *Entry, level int) {

	if f.directoryCache == nil {
		return
	}

	minutes := 60
	if level < 10 {
		minutes -= level * 6
	}

	f.directoryCache.Set(dirpath, dirEntry, time.Duration(minutes)*time.Minute)
}
<commit_msg>logging<commit_after>package filer2

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"google.golang.org/grpc"

	"github.com/karlseguin/ccache"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
)

const PaginationSize = 1024 * 256

var (
	OS_UID = uint32(os.Getuid())
	OS_GID = uint32(os.Getgid())
)

type Filer struct {
	store              *FilerStoreWrapper
	directoryCache     *ccache.Cache
	MasterClient       *wdclient.MasterClient
	fileIdDeletionChan chan string
	GrpcDialOption     grpc.DialOption
}

func NewFiler(masters []string, grpcDialOption grpc.DialOption) *Filer {
	f := &Filer{
		directoryCache:     ccache.New(ccache.Configure().MaxSize(1000).ItemsToPrune(100)),
		MasterClient:       wdclient.NewMasterClient(context.Background(), grpcDialOption, "filer", masters),
		fileIdDeletionChan: make(chan string, PaginationSize),
		GrpcDialOption:     grpcDialOption,
	}

	go f.loopProcessingDeletion()

	return f
}

func (f *Filer) SetStore(store FilerStore) {
	f.store = NewFilerStoreWrapper(store)
}

func (f *Filer) DisableDirectoryCache() {
	f.directoryCache = nil
}

func (fs *Filer) GetMaster() string {
	return fs.MasterClient.GetMaster()
}

func (fs *Filer) KeepConnectedToMaster() {
	fs.MasterClient.KeepConnectedToMaster()
}

func (f *Filer) BeginTransaction(ctx context.Context) (context.Context, error) {
	return f.store.BeginTransaction(ctx)
}

func (f *Filer) CommitTransaction(ctx context.Context) error {
	return f.store.CommitTransaction(ctx)
}

func (f *Filer) RollbackTransaction(ctx context.Context) error {
	return f.store.RollbackTransaction(ctx)
}

func (f *Filer) CreateEntry(ctx context.Context, entry *Entry, o_excl bool) error {

	glog.V(4).Infof("CreateEntry create %s", entry.FullPath)

	if string(entry.FullPath) == "/" {
		return nil
	}

	dirParts := strings.Split(string(entry.FullPath), "/")

	// fmt.Printf("directory parts: %+v\n", dirParts)

	var lastDirectoryEntry *Entry

	for i := 1; i < len(dirParts); i++ {
		dirPath := "/" + filepath.ToSlash(filepath.Join(dirParts[:i]...))
		// fmt.Printf("%d directory: %+v\n", i, dirPath)

		// first check local cache
		dirEntry := f.cacheGetDirectory(dirPath)

		// not found, check the store directly
		if dirEntry == nil {
			glog.V(4).Infof("find uncached directory: %s", dirPath)
			dirEntry, _ = f.FindEntry(ctx, FullPath(dirPath))
		} else {
			// glog.V(4).Infof("found cached directory: %s", dirPath)
		}

		// no such existing directory
		if dirEntry == nil {

			// create the directory
			now := time.Now()

			dirEntry = &Entry{
				FullPath: FullPath(dirPath),
				Attr: Attr{
					Mtime:  now,
					Crtime: now,
					Mode:   os.ModeDir | 0770,
					Uid:    entry.Uid,
					Gid:    entry.Gid,
				},
			}

			glog.V(2).Infof("create directory: %s %v", dirPath, dirEntry.Mode)
			mkdirErr := f.store.InsertEntry(ctx, dirEntry)
			if mkdirErr != nil {
				if _, err := f.FindEntry(ctx, FullPath(dirPath)); err == ErrNotFound {
					glog.V(3).Infof("mkdir %s: %v", dirPath, mkdirErr)
					return fmt.Errorf("mkdir %s: %v", dirPath, mkdirErr)
				}
			} else {
				f.NotifyUpdateEvent(nil, dirEntry, false)
			}

		} else if !dirEntry.IsDirectory() {
			glog.Errorf("CreateEntry %s: %s should be a directory", entry.FullPath, dirPath)
			return fmt.Errorf("%s is a file", dirPath)
		}

		// cache the directory entry
		f.cacheSetDirectory(dirPath, dirEntry, i)

		// remember the direct parent directory entry
		if i == len(dirParts)-1 {
			lastDirectoryEntry = dirEntry
		}

	}

	if lastDirectoryEntry == nil {
		glog.Errorf("CreateEntry %s: lastDirectoryEntry is nil", entry.FullPath)
		return fmt.Errorf("parent folder not found: %v", entry.FullPath)
	}

	/*
		if !hasWritePermission(lastDirectoryEntry, entry) {
			glog.V(0).Infof("directory %s: %v, entry: uid=%d gid=%d",
				lastDirectoryEntry.FullPath, lastDirectoryEntry.Attr, entry.Uid, entry.Gid)
			return fmt.Errorf("no write permission in folder %v", lastDirectoryEntry.FullPath)
		}
	*/

	oldEntry, _ := f.FindEntry(ctx, entry.FullPath)

	if oldEntry == nil {
		if err := f.store.InsertEntry(ctx, entry); err != nil {
			glog.Errorf("insert entry %s: %v", entry.FullPath, err)
			return fmt.Errorf("insert entry %s: %v", entry.FullPath, err)
		}
	} else {
		if o_excl {
			glog.V(3).Infof("EEXIST: entry %s already exists", entry.FullPath)
			return fmt.Errorf("EEXIST: entry %s already exists", entry.FullPath)
		}
		if err := f.UpdateEntry(ctx, oldEntry, entry); err != nil {
			glog.Errorf("update entry %s: %v", entry.FullPath, err)
			return fmt.Errorf("update entry %s: %v", entry.FullPath, err)
		}
	}

	f.NotifyUpdateEvent(oldEntry, entry, true)

	f.deleteChunksIfNotNew(oldEntry, entry)

	glog.V(4).Infof("CreateEntry %s: created", entry.FullPath)

	return nil
}

func (f *Filer) UpdateEntry(ctx context.Context, oldEntry, entry *Entry) (err error) {
	if oldEntry != nil {
		if oldEntry.IsDirectory() && !entry.IsDirectory() {
			glog.Errorf("existing %s is a directory", entry.FullPath)
			return fmt.Errorf("existing %s is a directory", entry.FullPath)
		}
		if !oldEntry.IsDirectory() && entry.IsDirectory() {
			glog.Errorf("existing %s is a file", entry.FullPath)
			return fmt.Errorf("existing %s is a file", entry.FullPath)
		}
	}
	return f.store.UpdateEntry(ctx, entry)
}

func (f *Filer) FindEntry(ctx context.Context, p FullPath) (entry *Entry, err error) {

	now := time.Now()

	if string(p) == "/" {
		return &Entry{
			FullPath: p,
			Attr: Attr{
				Mtime:  now,
				Crtime: now,
				Mode:   os.ModeDir | 0755,
				Uid:    OS_UID,
				Gid:    OS_GID,
			},
		}, nil
	}
	return f.store.FindEntry(ctx, p)
}

func (f *Filer) ListDirectoryEntries(ctx context.Context, p FullPath, startFileName string, inclusive bool, limit int) ([]*Entry, error) {
	if strings.HasSuffix(string(p), "/") && len(p) > 1 {
		p = p[0 : len(p)-1]
	}
	return f.store.ListDirectoryEntries(ctx, p, startFileName, inclusive, limit)
}

func (f *Filer) cacheDelDirectory(dirpath string) {

	if dirpath == "/" {
		return
	}

	if f.directoryCache == nil {
		return
	}
	f.directoryCache.Delete(dirpath)
	return
}

func (f *Filer) cacheGetDirectory(dirpath string) *Entry {

	if f.directoryCache == nil {
		return nil
	}
	item := f.directoryCache.Get(dirpath)
	if item == nil {
		return nil
	}
	return item.Value().(*Entry)
}

func (f *Filer) cacheSetDirectory(dirpath string, dirEntry *Entry, level int) {

	if f.directoryCache == nil {
		return
	}

	minutes := 60
	if level < 10 {
		minutes -= level * 6
	}

	f.directoryCache.Set(dirpath, dirEntry, time.Duration(minutes)*time.Minute)
}
<|endoftext|>


<commit_before>package project

var (
	description string = "The azure-operator manages Kubernetes clusters on Azure."
	gitSHA             = "n/a"
	name        string = "azure-operator"
	source      string = "https://github.com/giantswarm/azure-operator"
	version            = "6.0.0"
)

func Description() string {
	return description
}

func GitSHA() string {
	return gitSHA
}

func Name() string {
	return name
}

func Source() string {
	return source
}

func Version() string {
	return version
}
<commit_msg>Bump version to 6.0.1-dev (#1682)<commit_after>package project

var (
	description string = "The azure-operator manages Kubernetes clusters on Azure."
	gitSHA             = "n/a"
	name        string = "azure-operator"
	source      string = "https://github.com/giantswarm/azure-operator"
	version            = "6.0.1-dev"
)

func Description() string {
	return description
}

func GitSHA() string {
	return gitSHA
}

func Name() string {
	return name
}

func Source() string {
	return source
}

func Version() string {
	return version
}
<|endoftext|>


<commit_before>// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

import (
	"context"
	"encoding/json"
	"net/url"
	"testing"
)

func TestUpdateViaScript(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Script(NewScript("ctx._source.tags += tag").Params(map[string]interface{}{"tag": "blue"}).Lang("groovy"))
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"script":{"lang":"groovy","params":{"tag":"blue"},"source":"ctx._source.tags += tag"}}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaScriptId(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	scriptParams := map[string]interface{}{
		"pageViewEvent": map[string]interface{}{
			"url":      "foo.com/bar",
			"response": 404,
			"time":     "2014-01-01 12:32",
		},
	}
	script := NewScriptStored("my_web_session_summariser").Params(scriptParams)

	update := client.Update().
		Index("sessions").Type("session").Id("dh3sgudg8gsrgl").
		Script(script).
		ScriptedUpsert(true).
		Upsert(map[string]interface{}{})
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/sessions/session/dh3sgudg8gsrgl/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"script":{"id":"my_web_session_summariser","params":{"pageViewEvent":{"response":404,"time":"2014-01-01 12:32","url":"foo.com/bar"}}},"scripted_upsert":true,"upsert":{}}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaScriptAndUpsert(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Script(NewScript("ctx._source.counter += count").Params(map[string]interface{}{"count": 4})).
		Upsert(map[string]interface{}{"counter": 1})
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"script":{"params":{"count":4},"source":"ctx._source.counter += count"},"upsert":{"counter":1}}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaDoc(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Doc(map[string]interface{}{"name": "new_name"}).
		DetectNoop(true)
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"detect_noop":true,"doc":{"name":"new_name"}}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaDocAndUpsert(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Doc(map[string]interface{}{"name": "new_name"}).
		DocAsUpsert(true).
		Timeout("1s").
		Refresh("true")
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{"refresh": []string{"true"}, "timeout": []string{"1s"}}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"doc":{"name":"new_name"},"doc_as_upsert":true}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaDocAndUpsertAndFetchSource(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Doc(map[string]interface{}{"name": "new_name"}).
		DocAsUpsert(true).
		Timeout("1s").
		Refresh("true").
		FetchSource(true)
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{
		"refresh": []string{"true"},
		"timeout": []string{"1s"},
	}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"_source":true,"doc":{"name":"new_name"},"doc_as_upsert":true}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateAndFetchSource(t *testing.T) {
	client := setupTestClientAndCreateIndexAndAddDocs(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	res, err := client.Update().
		Index(testIndexName).Id("1").
		Doc(map[string]interface{}{"user": "sandrae"}).
		DetectNoop(true).
		FetchSource(true).
		Do(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if res == nil {
		t.Fatal("expected response != nil")
	}
	if res.GetResult == nil {
		t.Fatal("expected GetResult != nil")
	}
	data, err := json.Marshal(res.GetResult.Source)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"user":"sandrae","message":"Welcome to Golang and Elasticsearch.","retweets":108,"created":"0001-01-01T00:00:00Z","tags":["golang","elasticsearch"]}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateOptimistic(t *testing.T) {
	client := setupTestClientAndCreateIndexAndAddDocs(t) //, SetTraceLog(log.New(os.Stdout, "", 0)))

	doc, err := client.Get().
		Index(testIndexName).Id("1").
		Do(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if doc.SeqNo == nil {
		t.Fatal("expected seq_no != nil")
	}
	if doc.PrimaryTerm == nil {
		t.Fatal("expected primary_term != nil")
	}

	// Update with seqNo != doc.SeqNo and primaryTerm != doc.PrimaryTerm
	_, err = client.Update().
		Index(testIndexName).Id(doc.Id).
		Doc(map[string]interface{}{"user": "sandrae"}).
		IfSeqNo(*doc.SeqNo + 1000).
		IfPrimaryTerm(*doc.PrimaryTerm + 1000).
		Do(context.Background())
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	if !IsConflict(err) {
		t.Fatalf("expected conflict error, got %v (%T)", err, err)
	}

	// Update with seqNo == doc.SeqNo and primaryTerm == doc.PrimaryTerm
	res, err := client.Update().
		Index(testIndexName).Id(doc.Id).
		Doc(map[string]interface{}{"user": "sandrae"}).
		IfSeqNo(*doc.SeqNo).
		IfPrimaryTerm(*doc.PrimaryTerm).
		FetchSource(true).
		Do(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if res == nil {
		t.Fatal("expected response != nil")
	}
	if res.GetResult == nil {
		t.Fatal("expected GetResult != nil")
	}
	data, err := json.Marshal(res.GetResult.Source)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"user":"sandrae","message":"Welcome to Golang and Elasticsearch.","retweets":108,"created":"0001-01-01T00:00:00Z","tags":["golang","elasticsearch"]}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}
<commit_msg>Fix usage of Update API (#1477)<commit_after>// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

import (
	"context"
	"encoding/json"
	"net/url"
	"testing"
)

func TestUpdateViaScript(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Script(NewScript("ctx._source.tags += params.tag").Params(map[string]interface{}{"tag": "blue"}).Lang("groovy"))
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"script":{"lang":"groovy","params":{"tag":"blue"},"source":"ctx._source.tags += params.tag"}}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaScriptId(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	scriptParams := map[string]interface{}{
		"pageViewEvent": map[string]interface{}{
			"url":      "foo.com/bar",
			"response": 404,
			"time":     "2014-01-01 12:32",
		},
	}
	script := NewScriptStored("my_web_session_summariser").Params(scriptParams)

	update := client.Update().
		Index("sessions").Type("session").Id("dh3sgudg8gsrgl").
		Script(script).
		ScriptedUpsert(true).
		Upsert(map[string]interface{}{})
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/sessions/session/dh3sgudg8gsrgl/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"script":{"id":"my_web_session_summariser","params":{"pageViewEvent":{"response":404,"time":"2014-01-01 12:32","url":"foo.com/bar"}}},"scripted_upsert":true,"upsert":{}}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaScriptAndUpsert(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Script(NewScript("ctx._source.counter += params.count").Params(map[string]interface{}{"count": 4})).
		Upsert(map[string]interface{}{"counter": 1})
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"script":{"params":{"count":4},"source":"ctx._source.counter += params.count"},"upsert":{"counter":1}}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaDoc(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Doc(map[string]interface{}{"name": "new_name"}).
		DetectNoop(true)
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"detect_noop":true,"doc":{"name":"new_name"}}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaDocAndUpsert(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Doc(map[string]interface{}{"name": "new_name"}).
		DocAsUpsert(true).
		Timeout("1s").
		Refresh("true")
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{"refresh": []string{"true"}, "timeout": []string{"1s"}}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"doc":{"name":"new_name"},"doc_as_upsert":true}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateViaDocAndUpsertAndFetchSource(t *testing.T) {
	client := setupTestClient(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	update := client.Update().
		Index("test").Type("type1").Id("1").
		Doc(map[string]interface{}{"name": "new_name"}).
		DocAsUpsert(true).
		Timeout("1s").
		Refresh("true").
		FetchSource(true)
	path, params, err := update.url()
	if err != nil {
		t.Fatalf("expected to return URL, got: %v", err)
	}
	expectedPath := `/test/type1/1/_update`
	if expectedPath != path {
		t.Errorf("expected URL path\n%s\ngot:\n%s", expectedPath, path)
	}
	expectedParams := url.Values{
		"refresh": []string{"true"},
		"timeout": []string{"1s"},
	}
	if expectedParams.Encode() != params.Encode() {
		t.Errorf("expected URL parameters\n%s\ngot:\n%s", expectedParams.Encode(), params.Encode())
	}
	body, err := update.body()
	if err != nil {
		t.Fatalf("expected to return body, got: %v", err)
	}
	data, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"_source":true,"doc":{"name":"new_name"},"doc_as_upsert":true}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateAndFetchSource(t *testing.T) {
	client := setupTestClientAndCreateIndexAndAddDocs(t) // , SetTraceLog(log.New(os.Stdout, "", 0)))

	res, err := client.Update().
		Index(testIndexName).Id("1").
		Doc(map[string]interface{}{"user": "sandrae"}).
		DetectNoop(true).
		FetchSource(true).
		Do(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if res == nil {
		t.Fatal("expected response != nil")
	}
	if res.GetResult == nil {
		t.Fatal("expected GetResult != nil")
	}
	data, err := json.Marshal(res.GetResult.Source)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"user":"sandrae","message":"Welcome to Golang and Elasticsearch.","retweets":108,"created":"0001-01-01T00:00:00Z","tags":["golang","elasticsearch"]}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}

func TestUpdateOptimistic(t *testing.T) {
	client := setupTestClientAndCreateIndexAndAddDocs(t) //, SetTraceLog(log.New(os.Stdout, "", 0)))

	doc, err := client.Get().
		Index(testIndexName).Id("1").
		Do(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if doc.SeqNo == nil {
		t.Fatal("expected seq_no != nil")
	}
	if doc.PrimaryTerm == nil {
		t.Fatal("expected primary_term != nil")
	}

	// Update with seqNo != doc.SeqNo and primaryTerm != doc.PrimaryTerm
	_, err = client.Update().
		Index(testIndexName).Id(doc.Id).
		Doc(map[string]interface{}{"user": "sandrae"}).
		IfSeqNo(*doc.SeqNo + 1000).
		IfPrimaryTerm(*doc.PrimaryTerm + 1000).
		Do(context.Background())
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	if !IsConflict(err) {
		t.Fatalf("expected conflict error, got %v (%T)", err, err)
	}

	// Update with seqNo == doc.SeqNo and primaryTerm == doc.PrimaryTerm
	res, err := client.Update().
		Index(testIndexName).Id(doc.Id).
		Doc(map[string]interface{}{"user": "sandrae"}).
		IfSeqNo(*doc.SeqNo).
		IfPrimaryTerm(*doc.PrimaryTerm).
		FetchSource(true).
		Do(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if res == nil {
		t.Fatal("expected response != nil")
	}
	if res.GetResult == nil {
		t.Fatal("expected GetResult != nil")
	}
	data, err := json.Marshal(res.GetResult.Source)
	if err != nil {
		t.Fatalf("expected to marshal body as JSON, got: %v", err)
	}
	got := string(data)
	expected := `{"user":"sandrae","message":"Welcome to Golang and Elasticsearch.","retweets":108,"created":"0001-01-01T00:00:00Z","tags":["golang","elasticsearch"]}`
	if got != expected {
		t.Errorf("expected\n%s\ngot:\n%s", expected, got)
	}
}
<|endoftext|>


<commit_before>/*
Copyright The Helm 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 provenance

import (
	"bytes"
	"crypto"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/ghodss/yaml"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/clearsign"
	"golang.org/x/crypto/openpgp/packet"

	"k8s.io/helm/pkg/chartutil"
	hapi "k8s.io/helm/pkg/proto/hapi/chart"
)

var defaultPGPConfig = packet.Config{
	DefaultHash: crypto.SHA512,
}

// SumCollection represents a collection of file and image checksums.
//
// Files are of the form:
//	FILENAME: "sha256:SUM"
// Images are of the form:
//	"IMAGE:TAG": "sha256:SUM"
// Docker optionally supports sha512, and if this is the case, the hash marker
// will be 'sha512' instead of 'sha256'.
type SumCollection struct {
	Files  map[string]string `json:"files"`
	Images map[string]string `json:"images,omitempty"`
}

// Verification contains information about a verification operation.
type Verification struct {
	// SignedBy contains the entity that signed a chart.
	SignedBy *openpgp.Entity
	// FileHash is the hash, prepended with the scheme, for the file that was verified.
	FileHash string
	// FileName is the name of the file that FileHash verifies.
	FileName string
}

// Signatory signs things.
//
// Signatories can be constructed from a PGP private key file using NewFromFiles
// or they can be constructed manually by setting the Entity to a valid
// PGP entity.
//
// The same Signatory can be used to sign or validate multiple charts.
type Signatory struct {
	// The signatory for this instance of Helm. This is used for signing.
	Entity *openpgp.Entity
	// The keyring for this instance of Helm. This is used for verification.
	KeyRing openpgp.EntityList
}

// NewFromFiles constructs a new Signatory from the PGP key in the given filename.
//
// This will emit an error if it cannot find a valid GPG keyfile (entity) at the
// given location.
//
// Note that the keyfile may have just a public key, just a private key, or
// both. The Signatory methods may have different requirements of the keys. For
// example, ClearSign must have a valid `openpgp.Entity.PrivateKey` before it
// can sign something.
func NewFromFiles(keyfile, keyringfile string) (*Signatory, error) {
	e, err := loadKey(keyfile)
	if err != nil {
		return nil, err
	}

	ring, err := loadKeyRing(keyringfile)
	if err != nil {
		return nil, err
	}

	return &Signatory{
		Entity:  e,
		KeyRing: ring,
	}, nil
}

// NewFromKeyring reads a keyring file and creates a Signatory.
//
// If id is not the empty string, this will also try to find an Entity in the
// keyring whose name matches, and set that as the signing entity. It will return
// an error if the id is not empty and also not found.
func NewFromKeyring(keyringfile, id string) (*Signatory, error) {
	ring, err := loadKeyRing(keyringfile)
	if err != nil {
		return nil, err
	}

	s := &Signatory{KeyRing: ring}

	// If the ID is empty, we can return now.
	if id == "" {
		return s, nil
	}

	// We're gonna go all GnuPG on this and look for a string that _contains_. If
	// two or more keys contain the string and none are a direct match, we error
	// out.
	var candidate *openpgp.Entity
	vague := false
	for _, e := range ring {
		for n := range e.Identities {
			if n == id {
				s.Entity = e
				return s, nil
			}
			if strings.Contains(n, id) {
				if candidate != nil {
					vague = true
				}
				candidate = e
			}
		}
	}
	if vague {
		return s, fmt.Errorf("more than one key contain the id %q", id)
	}

	s.Entity = candidate
	return s, nil
}

// PassphraseFetcher returns a passphrase for decrypting keys.
//
// This is used as a callback to read a passphrase from some other location. The
// given name is the Name field on the key, typically of the form:
//
//	USER_NAME (COMMENT) <EMAIL>
type PassphraseFetcher func(name string) ([]byte, error)

// DecryptKey decrypts a private key in the Signatory.
//
// If the key is not encrypted, this will return without error.
//
// If the key does not exist, this will return an error.
//
// If the key exists, but cannot be unlocked with the passphrase returned by
// the PassphraseFetcher, this will return an error.
//
// If the key is successfully unlocked, it will return nil.
func (s *Signatory) DecryptKey(fn PassphraseFetcher) error {
	if s.Entity == nil {
		return errors.New("private key not found")
	} else if s.Entity.PrivateKey == nil {
		return errors.New("provided key is not a private key")
	}

	// Nothing else to do if key is not encrypted.
	if !s.Entity.PrivateKey.Encrypted {
		return nil
	}

	fname := "Unknown"
	for i := range s.Entity.Identities {
		if i != "" {
			fname = i
			break
		}
	}

	p, err := fn(fname)
	if err != nil {
		return err
	}

	return s.Entity.PrivateKey.Decrypt(p)
}

// ClearSign signs a chart with the given key.
//
// This takes the path to a chart archive file and a key, and it returns a clear signature.
//
// The Signatory must have a valid Entity.PrivateKey for this to work. If it does
// not, an error will be returned.
func (s *Signatory) ClearSign(chartpath string) (string, error) {
	if s.Entity == nil {
		return "", errors.New("private key not found")
	} else if s.Entity.PrivateKey == nil {
		return "", errors.New("provided key is not a private key")
	}

	if fi, err := os.Stat(chartpath); err != nil {
		return "", err
	} else if fi.IsDir() {
		return "", errors.New("cannot sign a directory")
	}

	out := bytes.NewBuffer(nil)

	b, err := messageBlock(chartpath)
	if err != nil {
		return "", nil
	}

	// Sign the buffer
	w, err := clearsign.Encode(out, s.Entity.PrivateKey, &defaultPGPConfig)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(w, b)
	w.Close()
	return out.String(), err
}

// Verify checks a signature and verifies that it is legit for a chart.
func (s *Signatory) Verify(chartpath, sigpath string) (*Verification, error) {
	ver := &Verification{}
	for _, fname := range []string{chartpath, sigpath} {
		if fi, err := os.Stat(fname); err != nil {
			return ver, err
		} else if fi.IsDir() {
			return ver, fmt.Errorf("%s cannot be a directory", fname)
		}
	}

	// First verify the signature
	sig, err := s.decodeSignature(sigpath)
	if err != nil {
		return ver, fmt.Errorf("failed to decode signature: %s", err)
	}

	by, err := s.verifySignature(sig)
	if err != nil {
		return ver, err
	}
	ver.SignedBy = by

	// Second, verify the hash of the tarball.
	sum, err := DigestFile(chartpath)
	if err != nil {
		return ver, err
	}
	_, sums, err := parseMessageBlock(sig.Plaintext)
	if err != nil {
		return ver, err
	}

	sum = "sha256:" + sum
	basename := filepath.Base(chartpath)
	if sha, ok := sums.Files[basename]; !ok {
		return ver, fmt.Errorf("provenance does not contain a SHA for a file named %q", basename)
	} else if sha != sum {
		return ver, fmt.Errorf("sha256 sum does not match for %s: %q != %q", basename, sha, sum)
	}
	ver.FileHash = sum
	ver.FileName = basename

	// TODO: when image signing is added, verify that here.

	return ver, nil
}

func (s *Signatory) decodeSignature(filename string) (*clearsign.Block, error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	block, _ := clearsign.Decode(data)
	if block == nil {
		// There was no sig in the file.
		return nil, errors.New("signature block not found")
	}

	return block, nil
}

// verifySignature verifies that the given block is validly signed, and returns the signer.
func (s *Signatory) verifySignature(block *clearsign.Block) (*openpgp.Entity, error) {
	return openpgp.CheckDetachedSignature(
		s.KeyRing,
		bytes.NewBuffer(block.Bytes),
		block.ArmoredSignature.Body,
	)
}

func messageBlock(chartpath string) (*bytes.Buffer, error) {
	var b *bytes.Buffer
	// Checksum the archive
	chash, err := DigestFile(chartpath)
	if err != nil {
		return b, err
	}

	base := filepath.Base(chartpath)
	sums := &SumCollection{
		Files: map[string]string{
			base: "sha256:" + chash,
		},
	}

	// Load the archive into memory.
	chart, err := chartutil.LoadFile(chartpath)
	if err != nil {
		return b, err
	}

	// Buffer a hash + checksums YAML file
	data, err := yaml.Marshal(chart.Metadata)
	if err != nil {
		return b, err
	}

	// FIXME: YAML uses ---\n as a file start indicator, but this is not legal in a PGP
	// clearsign block. So we use ...\n, which is the YAML document end marker.
	// http://yaml.org/spec/1.2/spec.html#id2800168
	b = bytes.NewBuffer(data)
	b.WriteString("\n...\n")

	data, err = yaml.Marshal(sums)
	if err != nil {
		return b, err
	}
	b.Write(data)

	return b, nil
}

// parseMessageBlock
func parseMessageBlock(data []byte) (*hapi.Metadata, *SumCollection, error) {
	// This sucks.
	parts := bytes.Split(data, []byte("\n...\n"))
	if len(parts) < 2 {
		return nil, nil, errors.New("message block must have at least two parts")
	}

	md := &hapi.Metadata{}
	sc := &SumCollection{}

	if err := yaml.Unmarshal(parts[0], md); err != nil {
		return md, sc, err
	}
	err := yaml.Unmarshal(parts[1], sc)
	return md, sc, err
}

// loadKey loads a GPG key found at a particular path.
func loadKey(keypath string) (*openpgp.Entity, error) {
	f, err := os.Open(keypath)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	pr := packet.NewReader(f)
	return openpgp.ReadEntity(pr)
}

func loadKeyRing(ringpath string) (openpgp.EntityList, error) {
	f, err := os.Open(ringpath)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return openpgp.ReadKeyRing(f)
}

// DigestFile calculates a SHA256 hash (like Docker) for a given file.
//
// It takes the path to the archive file, and returns a string representation of
// the SHA256 sum.
//
// The intended use of this function is to generate a sum of a chart TGZ file.
func DigestFile(filename string) (string, error) {
	f, err := os.Open(filename)
	if err != nil {
		return "", err
	}
	defer f.Close()
	return Digest(f)
}

// Digest hashes a reader and returns a SHA256 digest.
//
// Helm uses SHA256 as its default hash for all non-cryptographic applications.
func Digest(in io.Reader) (string, error) {
	hash := crypto.SHA256.New()
	io.Copy(hash, in)
	return hex.EncodeToString(hash.Sum(nil)), nil
}
<commit_msg>add errcheck for Digest method in sign.go<commit_after>/*
Copyright The Helm 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 provenance

import (
	"bytes"
	"crypto"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/ghodss/yaml"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/clearsign"
	"golang.org/x/crypto/openpgp/packet"

	"k8s.io/helm/pkg/chartutil"
	hapi "k8s.io/helm/pkg/proto/hapi/chart"
)

var defaultPGPConfig = packet.Config{
	DefaultHash: crypto.SHA512,
}

// SumCollection represents a collection of file and image checksums.
//
// Files are of the form:
//	FILENAME: "sha256:SUM"
// Images are of the form:
//	"IMAGE:TAG": "sha256:SUM"
// Docker optionally supports sha512, and if this is the case, the hash marker
// will be 'sha512' instead of 'sha256'.
type SumCollection struct {
	Files  map[string]string `json:"files"`
	Images map[string]string `json:"images,omitempty"`
}

// Verification contains information about a verification operation.
type Verification struct {
	// SignedBy contains the entity that signed a chart.
	SignedBy *openpgp.Entity
	// FileHash is the hash, prepended with the scheme, for the file that was verified.
	FileHash string
	// FileName is the name of the file that FileHash verifies.
	FileName string
}

// Signatory signs things.
//
// Signatories can be constructed from a PGP private key file using NewFromFiles
// or they can be constructed manually by setting the Entity to a valid
// PGP entity.
//
// The same Signatory can be used to sign or validate multiple charts.
type Signatory struct {
	// The signatory for this instance of Helm. This is used for signing.
	Entity *openpgp.Entity
	// The keyring for this instance of Helm. This is used for verification.
	KeyRing openpgp.EntityList
}

// NewFromFiles constructs a new Signatory from the PGP key in the given filename.
//
// This will emit an error if it cannot find a valid GPG keyfile (entity) at the
// given location.
//
// Note that the keyfile may have just a public key, just a private key, or
// both. The Signatory methods may have different requirements of the keys. For
// example, ClearSign must have a valid `openpgp.Entity.PrivateKey` before it
// can sign something.
func NewFromFiles(keyfile, keyringfile string) (*Signatory, error) {
	e, err := loadKey(keyfile)
	if err != nil {
		return nil, err
	}

	ring, err := loadKeyRing(keyringfile)
	if err != nil {
		return nil, err
	}

	return &Signatory{
		Entity:  e,
		KeyRing: ring,
	}, nil
}

// NewFromKeyring reads a keyring file and creates a Signatory.
//
// If id is not the empty string, this will also try to find an Entity in the
// keyring whose name matches, and set that as the signing entity. It will return
// an error if the id is not empty and also not found.
func NewFromKeyring(keyringfile, id string) (*Signatory, error) {
	ring, err := loadKeyRing(keyringfile)
	if err != nil {
		return nil, err
	}

	s := &Signatory{KeyRing: ring}

	// If the ID is empty, we can return now.
	if id == "" {
		return s, nil
	}

	// We're gonna go all GnuPG on this and look for a string that _contains_. If
	// two or more keys contain the string and none are a direct match, we error
	// out.
	var candidate *openpgp.Entity
	vague := false
	for _, e := range ring {
		for n := range e.Identities {
			if n == id {
				s.Entity = e
				return s, nil
			}
			if strings.Contains(n, id) {
				if candidate != nil {
					vague = true
				}
				candidate = e
			}
		}
	}
	if vague {
		return s, fmt.Errorf("more than one key contain the id %q", id)
	}

	s.Entity = candidate
	return s, nil
}

// PassphraseFetcher returns a passphrase for decrypting keys.
//
// This is used as a callback to read a passphrase from some other location. The
// given name is the Name field on the key, typically of the form:
//
//	USER_NAME (COMMENT) <EMAIL>
type PassphraseFetcher func(name string) ([]byte, error)

// DecryptKey decrypts a private key in the Signatory.
//
// If the key is not encrypted, this will return without error.
//
// If the key does not exist, this will return an error.
//
// If the key exists, but cannot be unlocked with the passphrase returned by
// the PassphraseFetcher, this will return an error.
//
// If the key is successfully unlocked, it will return nil.
func (s *Signatory) DecryptKey(fn PassphraseFetcher) error {
	if s.Entity == nil {
		return errors.New("private key not found")
	} else if s.Entity.PrivateKey == nil {
		return errors.New("provided key is not a private key")
	}

	// Nothing else to do if key is not encrypted.
	if !s.Entity.PrivateKey.Encrypted {
		return nil
	}

	fname := "Unknown"
	for i := range s.Entity.Identities {
		if i != "" {
			fname = i
			break
		}
	}

	p, err := fn(fname)
	if err != nil {
		return err
	}

	return s.Entity.PrivateKey.Decrypt(p)
}

// ClearSign signs a chart with the given key.
//
// This takes the path to a chart archive file and a key, and it returns a clear signature.
//
// The Signatory must have a valid Entity.PrivateKey for this to work. If it does
// not, an error will be returned.
func (s *Signatory) ClearSign(chartpath string) (string, error) {
	if s.Entity == nil {
		return "", errors.New("private key not found")
	} else if s.Entity.PrivateKey == nil {
		return "", errors.New("provided key is not a private key")
	}

	if fi, err := os.Stat(chartpath); err != nil {
		return "", err
	} else if fi.IsDir() {
		return "", errors.New("cannot sign a directory")
	}

	out := bytes.NewBuffer(nil)

	b, err := messageBlock(chartpath)
	if err != nil {
		return "", nil
	}

	// Sign the buffer
	w, err := clearsign.Encode(out, s.Entity.PrivateKey, &defaultPGPConfig)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(w, b)
	w.Close()
	return out.String(), err
}

// Verify checks a signature and verifies that it is legit for a chart.
func (s *Signatory) Verify(chartpath, sigpath string) (*Verification, error) {
	ver := &Verification{}
	for _, fname := range []string{chartpath, sigpath} {
		if fi, err := os.Stat(fname); err != nil {
			return ver, err
		} else if fi.IsDir() {
			return ver, fmt.Errorf("%s cannot be a directory", fname)
		}
	}

	// First verify the signature
	sig, err := s.decodeSignature(sigpath)
	if err != nil {
		return ver, fmt.Errorf("failed to decode signature: %s", err)
	}

	by, err := s.verifySignature(sig)
	if err != nil {
		return ver, err
	}
	ver.SignedBy = by

	// Second, verify the hash of the tarball.
	sum, err := DigestFile(chartpath)
	if err != nil {
		return ver, err
	}
	_, sums, err := parseMessageBlock(sig.Plaintext)
	if err != nil {
		return ver, err
	}

	sum = "sha256:" + sum
	basename := filepath.Base(chartpath)
	if sha, ok := sums.Files[basename]; !ok {
		return ver, fmt.Errorf("provenance does not contain a SHA for a file named %q", basename)
	} else if sha != sum {
		return ver, fmt.Errorf("sha256 sum does not match for %s: %q != %q", basename, sha, sum)
	}
	ver.FileHash = sum
	ver.FileName = basename

	// TODO: when image signing is added, verify that here.

	return ver, nil
}

func (s *Signatory) decodeSignature(filename string) (*clearsign.Block, error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	block, _ := clearsign.Decode(data)
	if block == nil {
		// There was no sig in the file.
		return nil, errors.New("signature block not found")
	}

	return block, nil
}

// verifySignature verifies that the given block is validly signed, and returns the signer.
func (s *Signatory) verifySignature(block *clearsign.Block) (*openpgp.Entity, error) {
	return openpgp.CheckDetachedSignature(
		s.KeyRing,
		bytes.NewBuffer(block.Bytes),
		block.ArmoredSignature.Body,
	)
}

func messageBlock(chartpath string) (*bytes.Buffer, error) {
	var b *bytes.Buffer
	// Checksum the archive
	chash, err := DigestFile(chartpath)
	if err != nil {
		return b, err
	}

	base := filepath.Base(chartpath)
	sums := &SumCollection{
		Files: map[string]string{
			base: "sha256:" + chash,
		},
	}

	// Load the archive into memory.
	chart, err := chartutil.LoadFile(chartpath)
	if err != nil {
		return b, err
	}

	// Buffer a hash + checksums YAML file
	data, err := yaml.Marshal(chart.Metadata)
	if err != nil {
		return b, err
	}

	// FIXME: YAML uses ---\n as a file start indicator, but this is not legal in a PGP
	// clearsign block. So we use ...\n, which is the YAML document end marker.
	// http://yaml.org/spec/1.2/spec.html#id2800168
	b = bytes.NewBuffer(data)
	b.WriteString("\n...\n")

	data, err = yaml.Marshal(sums)
	if err != nil {
		return b, err
	}
	b.Write(data)

	return b, nil
}

// parseMessageBlock
func parseMessageBlock(data []byte) (*hapi.Metadata, *SumCollection, error) {
	// This sucks.
	parts := bytes.Split(data, []byte("\n...\n"))
	if len(parts) < 2 {
		return nil, nil, errors.New("message block must have at least two parts")
	}

	md := &hapi.Metadata{}
	sc := &SumCollection{}

	if err := yaml.Unmarshal(parts[0], md); err != nil {
		return md, sc, err
	}
	err := yaml.Unmarshal(parts[1], sc)
	return md, sc, err
}

// loadKey loads a GPG key found at a particular path.
func loadKey(keypath string) (*openpgp.Entity, error) {
	f, err := os.Open(keypath)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	pr := packet.NewReader(f)
	return openpgp.ReadEntity(pr)
}

func loadKeyRing(ringpath string) (openpgp.EntityList, error) {
	f, err := os.Open(ringpath)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return openpgp.ReadKeyRing(f)
}

// DigestFile calculates a SHA256 hash (like Docker) for a given file.
//
// It takes the path to the archive file, and returns a string representation of
// the SHA256 sum.
//
// The intended use of this function is to generate a sum of a chart TGZ file.
func DigestFile(filename string) (string, error) {
	f, err := os.Open(filename)
	if err != nil {
		return "", err
	}
	defer f.Close()
	return Digest(f)
}

// Digest hashes a reader and returns a SHA256 digest.
//
// Helm uses SHA256 as its default hash for all non-cryptographic applications.
func Digest(in io.Reader) (string, error) {
	hash := crypto.SHA256.New()
	if _, err := io.Copy(hash, in); err != nil {
		return "", nil
	}
	return hex.EncodeToString(hash.Sum(nil)), nil
}
<|endoftext|>