<commit_before>/*
 * Copyright (c) 2013-2014, Jeremy Bingham (<jbingham@gmail.com>)
 *
 * 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 search

import (
	"fmt"
	"github.com/ctdk/goas/v2/logger"
	"github.com/ctdk/goiardi/indexer"
	//"github.com/ctdk/goiardi/util"
	"regexp"
	"strings"
)

type PostgresSearch struct {
}

type PgQuery struct {
	queryChain Queryable
	paths []string
	queryStrs []string
	arguments []string
}

type gClause struct {
	clause string
	op Op
}

func (p *PostgresSearch) Search(idx string, q string, rows int, sortOrder string, start int, partialData map[string]interface{}) ([]map[string]interface{}, error) {
	// keep up with the ersatz solr.
	qq := &Tokenizer{Buffer: q}
	qq.Init()
	if err := qq.Parse(); err != nil {
		return nil, err
	}
	qq.Execute()
	qchain := qq.Evaluate()

	pgQ := &PgQuery{ queryChain: qchain }

	logger.Debugf("what on earth is the chain? %q", qchain)
	err := pgQ.execute()
	if err != nil {
		return nil, err
	}

	// dummy
	dres := make([]map[string]interface{}, 0)
	return dres, nil
}

func (p *PostgresSearch) GetEndpoints() []string {
	// TODO: deal with possible errors
	endpoints, err := indexer.Endpoints()
	return endpoints
	if err != nil {
		panic(err)
	}
	return endpoints
}

func (pq *PgQuery) execute(startTableID ...*int) error {
	p := pq.queryChain
	curOp := OpNotAnOp
	opMap := map[Op]string{
		OpNotAnOp: "(not an op)",
		OpUnaryNot: "not",
		OpUnaryReq: "req",
		OpUnaryPro: "pro",
		OpBinAnd: "and",
		OpBinOr: "or",
		OpBoost: "boost",
		OpFuzzy: "fuzzy",
		OpStartGroup: "start group",
		OpEndGroup: "end group",
		OpStartIncl: "start inc",
		OpEndIncl: "end inc",
		OpStartExcl: "start exc",
		OpEndExcl: "end exc",
	}
	var t *int
	if len(startTableID) == 0 {
		z := 0
		t = &z
	} else {
		t = startTableID[0]
	}
	for p != nil {
		switch c := p.(type) {
		case *BasicQuery:
			pq.paths = append(pq.paths, string(c.field))
			logger.Debugf("basic t%d: field: %s op: %s term: %+v complete %v", *t, c.field, opMap[c.op], c.term, c.complete)
			args, qstr := buildBasicQuery(c.field, c.term, t, curOp)
			logger.Debugf("qstr: %s", qstr)
			pq.arguments = append(pq.arguments, args...)
			pq.queryStrs = append(pq.queryStrs, qstr)
			*t++
		case *GroupedQuery:
			pq.paths = append(pq.paths, string(c.field))
			logger.Debugf("grouped t%d: field: %s op: %s terms: %+v complete %v", *t, c.field, opMap[c.op], c.terms, c.complete)
			args, qstr := buildGroupedQuery(c.field, c.terms, t, curOp)
			logger.Debugf("qstr: %s", qstr)
			pq.arguments = append(pq.arguments, args...)
			pq.queryStrs = append(pq.queryStrs, qstr)
			*t++
		case *RangeQuery:
			pq.paths = append(pq.paths, string(c.field))
			logger.Debugf("range t%d: field %s op %s start %s end %s inclusive %v complete %v", *t, c.field, opMap[c.op], c.start, c.end, c.inclusive, c.complete)
			args, qstr := buildRangeQuery(c.field, c.start, c.end, c.inclusive, t, curOp)
			logger.Debugf("qstr: %s", qstr)
			pq.arguments = append(pq.arguments, args...)
			pq.queryStrs = append(pq.queryStrs, qstr)
			*t++
		case *SubQuery:
			logger.Debugf("STARTING SUBQUERY: op %s complete %v", opMap[c.op], c.complete)
			newq, nend, nerr := extractSubQuery(c)
			if nerr != nil {
				return nerr
			}
			p = nend
			logger.Debugf("OP NOW: %s", opMap[p.Op()])
			np := &PgQuery{ queryChain: newq }
			err := np.execute(t)
			if err != nil {
				return err
			}
			logger.Debugf("subquery paths: %v", np.paths)
			logger.Debugf("subquery args: %v", np.arguments)
			logger.Debugf("subquery qstrs: %v", np.queryStrs)
			pq.paths = append(pq.paths, np.paths...)
			pq.arguments = append(pq.arguments, np.arguments...)
			pq.queryStrs = append(pq.queryStrs, fmt.Sprintf("%s(%s)", binOp(curOp), strings.Join(np.queryStrs, " ")))
			logger.Debugf("ENDING SUBQUERY")
		default:
			err := fmt.Errorf("Unknown type %T for query", c)
			return err
		}
		curOp = p.Op()
		p = p.Next()
	}
	logger.Debugf("paths: %v", pq.paths)
	logger.Debugf("arguments: %v", pq.arguments)
	logger.Debugf("query strings: %v", pq.queryStrs)
	logger.Debugf("number of tables: %d", *t)
	return nil
}

func buildBasicQuery(field Field, term QueryTerm, tNum *int, op Op) ([]string, string) {
	opStr := binOp(op)
	cop := matchOp(term.mod, term)

	var q string
	args := []string{ string(field) }
	if term.term == "*" {
		q = fmt.Sprintf("%s(f%d.path ~ _ARG_)", opStr, *tNum)
	} else {
		q = fmt.Sprintf("%s(f%d.path ~ _ARG_ AND f%d.value %s _ARG_)", opStr, *tNum, *tNum, cop)
		args = append(args, string(term.term))
	}

	return args, q
}

func buildGroupedQuery(field Field, terms []QueryTerm, tNum *int, op Op) ([]string, string) {
	opStr := binOp(op)

	var q string
	args := []string{ string(field) }
	var grouped []*gClause

	for _, v := range terms {
		cop := matchOp(op, v)

		clause := fmt.Sprintf("f%d.value %s _ARG_", *tNum, cop)
		g := &gClause{ clause, v.mod }
		grouped = append(grouped, g)
		args = append(args, string(v.term))
	}
	var clauseArr []string
	for i, g := range grouped {
		var j string
		if i != 0 {
			if g.op == OpUnaryPro || g.op == OpUnaryReq || g.op == OpUnaryNot {
				j = " AND "
			} else {
				j = " OR "
			}
		}
		clauseArr = append(clauseArr, fmt.Sprintf("%s%s", j, g.clause))
	}
	clauses := strings.Join(clauseArr, " ")
	q = fmt.Sprintf("%s(f%d.path ~ _ARG_ AND (%s))", opStr, *tNum, clauses)
	return args, q
}

func buildRangeQuery(field Field, start RangeTerm, end RangeTerm, inclusive bool, tNum *int, op Op) ([]string, string) {
	if start > end {
		start, end = end, start
	}

	var q string
	args := []string{ string(field) }

	opStr := binOp(op)
	var equals string
	if inclusive {
		equals = "="
	}
	var ranges []string
	if string(start) != "*" {
		s := fmt.Sprintf("f%d.value >%s _ARG_", *tNum, equals)
		ranges = append(ranges, s)
		args = append(args, string(start))
	}
	if string(end) != "*" {
		e := fmt.Sprintf("f%d.value <%s _ARG_", *tNum, equals)
		ranges = append(ranges, e)
		args = append(args, string(end))
	}
	var rangeStr string
	if len(ranges) != 0 {
		rangeStr = fmt.Sprintf(" AND (%s)", strings.Join(ranges, " AND "))
	}
	q = fmt.Sprintf("%s(f%d.path ~ _ARG_%s)", opStr, *tNum, rangeStr)
	return args, q
}

func matchOp(op Op, term QueryTerm) string {
	r := regexp.MustCompile(`\*|\?`)
	var cop string
	if r.MatchString(string(term.term)) {
		if term.mod == OpUnaryNot || term.mod == OpUnaryPro {
			cop = "NOT LIKE"
		} else {
			cop = "LIKE"
		}
	} else {
		if term.mod == OpUnaryNot || term.mod == OpUnaryPro {
			cop = "<>"
		} else {
			cop = "="
		}
	}
	return cop
}

func binOp(op Op) string {
	var opStr string
	if op != OpNotAnOp {
		if op == OpBinAnd {
			opStr = " AND "
		} else {
			opStr = " OR "
		}
	}
	return opStr
}
<commit_msg>queries build now<commit_after>/*
 * Copyright (c) 2013-2014, Jeremy Bingham (<jbingham@gmail.com>)
 *
 * 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 search

import (
	"fmt"
	"github.com/ctdk/goas/v2/logger"
	"github.com/ctdk/goiardi/indexer"
	//"github.com/ctdk/goiardi/util"
	"regexp"
	"strings"
)

type PostgresSearch struct {
}

type PgQuery struct {
	idx string
	queryChain Queryable
	paths []string
	queryStrs []string
	arguments []string
}

type gClause struct {
	clause string
	op Op
}

func (p *PostgresSearch) Search(idx string, q string, rows int, sortOrder string, start int, partialData map[string]interface{}) ([]map[string]interface{}, error) {
	// keep up with the ersatz solr.
	qq := &Tokenizer{Buffer: q}
	qq.Init()
	if err := qq.Parse(); err != nil {
		return nil, err
	}
	qq.Execute()
	qchain := qq.Evaluate()

	pgQ := &PgQuery{ idx: idx, queryChain: qchain }

	logger.Debugf("what on earth is the chain? %q", qchain)
	err := pgQ.execute()
	if err != nil {
		return nil, err
	}

	// dummy
	dres := make([]map[string]interface{}, 0)
	return dres, nil
}

func (p *PostgresSearch) GetEndpoints() []string {
	// TODO: deal with possible errors
	endpoints, err := indexer.Endpoints()
	return endpoints
	if err != nil {
		panic(err)
	}
	return endpoints
}

func (pq *PgQuery) execute(startTableID ...*int) error {
	p := pq.queryChain
	curOp := OpNotAnOp
	opMap := map[Op]string{
		OpNotAnOp: "(not an op)",
		OpUnaryNot: "not",
		OpUnaryReq: "req",
		OpUnaryPro: "pro",
		OpBinAnd: "and",
		OpBinOr: "or",
		OpBoost: "boost",
		OpFuzzy: "fuzzy",
		OpStartGroup: "start group",
		OpEndGroup: "end group",
		OpStartIncl: "start inc",
		OpEndIncl: "end inc",
		OpStartExcl: "start exc",
		OpEndExcl: "end exc",
	}
	var t *int
	if len(startTableID) == 0 {
		z := 0
		t = &z
	} else {
		t = startTableID[0]
	}
	for p != nil {
		switch c := p.(type) {
		case *BasicQuery:
			pq.paths = append(pq.paths, string(c.field))
			logger.Debugf("basic t%d: field: %s op: %s term: %+v complete %v", *t, c.field, opMap[c.op], c.term, c.complete)
			args, qstr := buildBasicQuery(c.field, c.term, t, curOp)
			logger.Debugf("qstr: %s", qstr)
			pq.arguments = append(pq.arguments, args...)
			pq.queryStrs = append(pq.queryStrs, qstr)
			*t++
		case *GroupedQuery:
			pq.paths = append(pq.paths, string(c.field))
			logger.Debugf("grouped t%d: field: %s op: %s terms: %+v complete %v", *t, c.field, opMap[c.op], c.terms, c.complete)
			args, qstr := buildGroupedQuery(c.field, c.terms, t, curOp)
			logger.Debugf("qstr: %s", qstr)
			pq.arguments = append(pq.arguments, args...)
			pq.queryStrs = append(pq.queryStrs, qstr)
			*t++
		case *RangeQuery:
			pq.paths = append(pq.paths, string(c.field))
			logger.Debugf("range t%d: field %s op %s start %s end %s inclusive %v complete %v", *t, c.field, opMap[c.op], c.start, c.end, c.inclusive, c.complete)
			args, qstr := buildRangeQuery(c.field, c.start, c.end, c.inclusive, t, curOp)
			logger.Debugf("qstr: %s", qstr)
			pq.arguments = append(pq.arguments, args...)
			pq.queryStrs = append(pq.queryStrs, qstr)
			*t++
		case *SubQuery:
			logger.Debugf("STARTING SUBQUERY: op %s complete %v", opMap[c.op], c.complete)
			newq, nend, nerr := extractSubQuery(c)
			if nerr != nil {
				return nerr
			}
			p = nend
			logger.Debugf("OP NOW: %s", opMap[p.Op()])
			np := &PgQuery{ queryChain: newq }
			err := np.execute(t)
			if err != nil {
				return err
			}
			logger.Debugf("subquery paths: %v", np.paths)
			logger.Debugf("subquery args: %v", np.arguments)
			logger.Debugf("subquery qstrs: %v", np.queryStrs)
			pq.paths = append(pq.paths, np.paths...)
			pq.arguments = append(pq.arguments, np.arguments...)
			pq.queryStrs = append(pq.queryStrs, fmt.Sprintf("%s(%s)", binOp(curOp), strings.Join(np.queryStrs, " ")))
			logger.Debugf("ENDING SUBQUERY")
		default:
			err := fmt.Errorf("Unknown type %T for query", c)
			return err
		}
		curOp = p.Op()
		p = p.Next()
	}
	logger.Debugf("paths: %v", pq.paths)
	logger.Debugf("arguments: %v", pq.arguments)
	logger.Debugf("query strings: %v", pq.queryStrs)
	logger.Debugf("number of tables: %d", *t)
	fullQ, allArgs := craftFullQuery(1, pq.idx, pq.paths, pq.arguments, pq.queryStrs, t)
	logger.Debugf("full query: %s", fullQ)
	logger.Debugf("all %d args: %v", len(allArgs), allArgs)
	return nil
}

func buildBasicQuery(field Field, term QueryTerm, tNum *int, op Op) ([]string, string) {
	opStr := binOp(op)
	cop := matchOp(term.mod, term)

	var q string
	args := []string{ string(field) }
	if term.term == "*" {
		q = fmt.Sprintf("%s(f%d.path ~ _ARG_)", opStr, *tNum)
	} else {
		q = fmt.Sprintf("%s(f%d.path ~ _ARG_ AND f%d.value %s _ARG_)", opStr, *tNum, *tNum, cop)
		args = append(args, string(term.term))
	}

	return args, q
}

func buildGroupedQuery(field Field, terms []QueryTerm, tNum *int, op Op) ([]string, string) {
	opStr := binOp(op)

	var q string
	args := []string{ string(field) }
	var grouped []*gClause

	for _, v := range terms {
		cop := matchOp(op, v)

		clause := fmt.Sprintf("f%d.value %s _ARG_", *tNum, cop)
		g := &gClause{ clause, v.mod }
		grouped = append(grouped, g)
		args = append(args, string(v.term))
	}
	var clauseArr []string
	for i, g := range grouped {
		var j string
		if i != 0 {
			if g.op == OpUnaryPro || g.op == OpUnaryReq || g.op == OpUnaryNot {
				j = " AND "
			} else {
				j = " OR "
			}
		}
		clauseArr = append(clauseArr, fmt.Sprintf("%s%s", j, g.clause))
	}
	clauses := strings.Join(clauseArr, " ")
	q = fmt.Sprintf("%s(f%d.path ~ _ARG_ AND (%s))", opStr, *tNum, clauses)
	return args, q
}

func buildRangeQuery(field Field, start RangeTerm, end RangeTerm, inclusive bool, tNum *int, op Op) ([]string, string) {
	if start > end {
		start, end = end, start
	}

	var q string
	args := []string{ string(field) }

	opStr := binOp(op)
	var equals string
	if inclusive {
		equals = "="
	}
	var ranges []string
	if string(start) != "*" {
		s := fmt.Sprintf("f%d.value >%s _ARG_", *tNum, equals)
		ranges = append(ranges, s)
		args = append(args, string(start))
	}
	if string(end) != "*" {
		e := fmt.Sprintf("f%d.value <%s _ARG_", *tNum, equals)
		ranges = append(ranges, e)
		args = append(args, string(end))
	}
	var rangeStr string
	if len(ranges) != 0 {
		rangeStr = fmt.Sprintf(" AND (%s)", strings.Join(ranges, " AND "))
	}
	q = fmt.Sprintf("%s(f%d.path ~ _ARG_%s)", opStr, *tNum, rangeStr)
	return args, q
}

func matchOp(op Op, term QueryTerm) string {
	r := regexp.MustCompile(`\*|\?`)
	var cop string
	if r.MatchString(string(term.term)) {
		if term.mod == OpUnaryNot || term.mod == OpUnaryPro {
			cop = "NOT LIKE"
		} else {
			cop = "LIKE"
		}
	} else {
		if term.mod == OpUnaryNot || term.mod == OpUnaryPro {
			cop = "<>"
		} else {
			cop = "="
		}
	}
	return cop
}

func binOp(op Op) string {
	var opStr string
	if op != OpNotAnOp {
		if op == OpBinAnd {
			opStr = " AND "
		} else {
			opStr = " OR "
		}
	}
	return opStr
}

func craftFullQuery(orgID int, idx string, paths []string, arguments []string, queryStrs []string, tNum *int) (string, []interface{}) {
	// TODO: FIX
	allArgs := make([]interface{}, 0, len(paths) + len(arguments) + 2)
	allArgs = append(allArgs, orgID)
	allArgs = append(allArgs, idx)
	for _, v := range paths {
		allArgs = append(allArgs, v)
	}
	for _, v := range arguments {
		allArgs = append(allArgs, v)
	}

	pcount := 3
	params := make([]string, 0, len(paths))
	for range paths {
		params = append(params, fmt.Sprintf("$%d", pcount))
		pcount++
	}
	withStatement := fmt.Sprintf("WITH found_items AS (SELECT item_name, path, value FROM goiardi.search_items si JOIN goiardi.search_collections sc ON si.search_collection_id = sc.id WHERE si.organization_id = $1 AND sc.name = $2 AND path ? ARRAY[%s]::lquery[]), items AS (SELECT DISTINCT item_name FROM found_items)", strings.Join(params, ", "))
	var selectStmt string
	if *tNum == 1 {
		selectStmt = fmt.Sprintf("SELECT DISTINCT item_name FROM found_items f0 WHERE %s", queryStrs[0])
	} else {
		joins := make([]string, 0, *tNum)
		for i := 0; i < *tNum; i++ {
			j := fmt.Sprintf("INNER JOIN found_items AS f%d ON i.item_name = f%d.item_name", i, i)
			joins = append(joins, j)
		}
		selectStmt = fmt.Sprintf("SELECT i.item_name FROM items i %s WHERE %s", strings.Join(joins, " "), strings.Join(queryStrs, " "))
	}
	fullQuery := strings.Join([]string{ withStatement, selectStmt }, " ")
	re := regexp.MustCompile("_ARG_")
	rfunc := func([]byte) []byte {
		r := []byte(fmt.Sprintf("$%d", pcount))
		pcount++
		return r
	}
	fullQuery = string(re.ReplaceAllFunc([]byte(fullQuery), rfunc))

	return fullQuery, allArgs
}
<|endoftext|>


<commit_before>package icsi

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"
)

const notaryDomain = "notary.icsi.berkeley.edu"

type Status int

const (
	Unknown Status = iota
	Seen
	Validated
)

var (
	ipSeen      = net.IP{127, 0, 0, 1}
	ipValidated = net.IP{127, 0, 0, 2}

	errInvalidResponse = errors.New("icsi: invalid response")
	errUnknownVersion  = errors.New("icsi: unknown version")
	errMultipleRecords = errors.New("icsi: multiple records")
)

func dnsname(sha []byte) string {
	return fmt.Sprintf("%x.%s", sha, notaryDomain)
}

func isnxdomain(err error) bool {
	if err, ok := err.(*net.DNSError); ok {
		return err.Err == "no such host"
	}
	return false
}

func QueryStatus(sha []byte) (Status, error) {
	ips, err := net.LookupIP(dnsname(sha))
	if err != nil {
		if isnxdomain(err) {
			return Unknown, nil
		}
		return Unknown, err
	}

	if len(ips) != 1 {
		return Unknown, errMultipleRecords
	}
	if bytes.Equal(ips[0], ipSeen) {
		return Seen, nil
	}
	if bytes.Equal(ips[0], ipValidated) {
		return Validated, nil
	}

	return Unknown, nil
}

type Response struct {
	Version   int
	FirstSeen time.Time
	LastSeen  time.Time
	TimesSeen int
	Validated bool
}

func parseDate(s string) (time.Time, error) {
	i, err := strconv.ParseInt(s, 10, 32)
	if err != nil {
		return time.Time{}, err
	}
	return time.Unix(0, 0).UTC().AddDate(0, 0, int(i)), nil
}

func parseResponse(txt string) (*Response, error) {
	var r Response

	tok := strings.Split(txt, " ")
	for _, t := range tok {
		pair := strings.Split(t, "=")
		if len(pair) != 2 {
			return nil, errInvalidResponse
		}
		switch pair[0] {
		case "version":
			i, err := strconv.ParseInt(pair[1], 10, 32)
			if err != nil {
				return nil, errInvalidResponse
			}
			r.Version = int(i)
		case "first_seen":
			t, err := parseDate(pair[1])
			if err != nil {
				return nil, errInvalidResponse
			}
			r.FirstSeen = t
		case "last_seen":
			t, err := parseDate(pair[1])
			if err != nil {
				return nil, errInvalidResponse
			}
			r.LastSeen = t
		case "times_seen":
			i, err := strconv.ParseInt(pair[1], 10, 32)
			if err != nil {
				return nil, errInvalidResponse
			}
			r.TimesSeen = int(i)
		case "validated":
			i, err := strconv.ParseInt(pair[1], 10, 32)
			if err != nil {
				return nil, errInvalidResponse
			}
			r.Validated = i == 1
		}
	}

	if r.Version != 1 {
		return nil, errUnknownVersion
	}

	return &r, nil

}

func Query(sha []byte) (*Response, error) {
	txts, err := net.LookupTXT(dnsname(sha))
	if err != nil {
		if isnxdomain(err) {
			return nil, nil
		}
		return nil, err
	}

	if len(txts) != 1 {
		return nil, errMultipleRecords
	}

	return parseResponse(txts[0])

}
<commit_msg>add Hash<commit_after>package icsi

import (
	"bytes"
	"crypto/sha1"
	"crypto/x509"
	"errors"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"
)

const notaryDomain = "notary.icsi.berkeley.edu"

type Status int

const (
	Unknown Status = iota
	Seen
	Validated
)

var (
	ipSeen      = net.IP{127, 0, 0, 1}
	ipValidated = net.IP{127, 0, 0, 2}

	errInvalidResponse = errors.New("icsi: invalid response")
	errUnknownVersion  = errors.New("icsi: unknown version")
	errMultipleRecords = errors.New("icsi: multiple records")
)

func dnsname(sha []byte) string {
	return fmt.Sprintf("%x.%s", sha, notaryDomain)
}

func isnxdomain(err error) bool {
	if err, ok := err.(*net.DNSError); ok {
		return err.Err == "no such host"
	}
	return false
}

func QueryStatus(sha []byte) (Status, error) {
	ips, err := net.LookupIP(dnsname(sha))
	if err != nil {
		if isnxdomain(err) {
			return Unknown, nil
		}
		return Unknown, err
	}

	if len(ips) != 1 {
		return Unknown, errMultipleRecords
	}
	if bytes.Equal(ips[0], ipSeen) {
		return Seen, nil
	}
	if bytes.Equal(ips[0], ipValidated) {
		return Validated, nil
	}

	return Unknown, nil
}

type Response struct {
	Version   int
	FirstSeen time.Time
	LastSeen  time.Time
	TimesSeen int
	Validated bool
}

func parseDate(s string) (time.Time, error) {
	i, err := strconv.ParseInt(s, 10, 32)
	if err != nil {
		return time.Time{}, err
	}
	return time.Unix(0, 0).UTC().AddDate(0, 0, int(i)), nil
}

func parseResponse(txt string) (*Response, error) {
	var r Response

	tok := strings.Split(txt, " ")
	for _, t := range tok {
		pair := strings.Split(t, "=")
		if len(pair) != 2 {
			return nil, errInvalidResponse
		}
		switch pair[0] {
		case "version":
			i, err := strconv.ParseInt(pair[1], 10, 32)
			if err != nil {
				return nil, errInvalidResponse
			}
			r.Version = int(i)
		case "first_seen":
			t, err := parseDate(pair[1])
			if err != nil {
				return nil, errInvalidResponse
			}
			r.FirstSeen = t
		case "last_seen":
			t, err := parseDate(pair[1])
			if err != nil {
				return nil, errInvalidResponse
			}
			r.LastSeen = t
		case "times_seen":
			i, err := strconv.ParseInt(pair[1], 10, 32)
			if err != nil {
				return nil, errInvalidResponse
			}
			r.TimesSeen = int(i)
		case "validated":
			i, err := strconv.ParseInt(pair[1], 10, 32)
			if err != nil {
				return nil, errInvalidResponse
			}
			r.Validated = i == 1
		}
	}

	if r.Version != 1 {
		return nil, errUnknownVersion
	}

	return &r, nil

}

func Query(sha []byte) (*Response, error) {
	txts, err := net.LookupTXT(dnsname(sha))
	if err != nil {
		if isnxdomain(err) {
			return nil, nil
		}
		return nil, err
	}

	if len(txts) != 1 {
		return nil, errMultipleRecords
	}

	return parseResponse(txts[0])

}

func Hash(cert *x509.Certificate) []byte {
	h := sha1.New()
	h.Write(cert.Raw)
	return h.Sum(nil)
}
<|endoftext|>


<commit_before>package raft

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"hash"
	"hash/crc64"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

const (
	testPath      = "permTest"
	snapPath      = "snapshots"
	metaFilePath  = "meta.json"
	stateFilePath = "state.bin"
	tmpSuffix     = ".tmp"
)

// FileSnapshotStore implements the SnapshotStore interface and allows
// snapshots to be made on the local disk.
type FileSnapshotStore struct {
	path   string
	retain int
	logger *log.Logger
}

type snapMetaSlice []*fileSnapshotMeta

// FileSnapshotSink implements SnapshotSink with a file.
type FileSnapshotSink struct {
	store  *FileSnapshotStore
	logger *log.Logger
	dir    string
	meta   fileSnapshotMeta

	stateFile *os.File
	stateHash hash.Hash64
	buffered  *bufio.Writer

	closed bool
}

// fileSnapshotMeta is stored on disk. We also put a CRC
// on disk so that we can verify the snapshot
type fileSnapshotMeta struct {
	SnapshotMeta
	CRC []byte
}

// bufferedFile is returned when we open a snapshot. This way
// reads are buffered and the file still gets closed.
type bufferedFile struct {
	bh *bufio.Reader
	fh *os.File
}

func (b *bufferedFile) Read(p []byte) (n int, err error) {
	return b.bh.Read(p)
}

func (b *bufferedFile) Close() error {
	return b.fh.Close()
}

// NewFileSnapshotStore creates a new FileSnapshotStore based
// on a base directory. The `retain` parameter controls how many
// snapshots are retained. Must be at least 1.
func NewFileSnapshotStore(base string, retain int, logOutput io.Writer) (*FileSnapshotStore, error) {
	if retain < 1 {
		return nil, fmt.Errorf("must retain at least one snapshot")
	}
	if logOutput == nil {
		logOutput = os.Stderr
	}

	// Ensure our path exists
	path := filepath.Join(base, snapPath)
	if err := os.Mkdir(path, 0755); err != nil && !os.IsExist(err) {
		return nil, fmt.Errorf("snapshot path not accessible: %v", err)
	}

	// Setup the store
	store := &FileSnapshotStore{
		path:   path,
		retain: retain,
		logger: log.New(logOutput, "", log.LstdFlags),
	}

	// Do a permissions test
	if err := store.testPermissions(); err != nil {
		return nil, fmt.Errorf("permissions test failed: %v", err)
	}
	return store, nil
}

// testPermissions tries to touch a file in our path to see if it works
func (f *FileSnapshotStore) testPermissions() error {
	path := filepath.Join(f.path, testPath)
	fh, err := os.Create(path)
	if err != nil {
		return err
	}
	fh.Close()
	os.Remove(path)
	return nil
}

// snapshotName generate s name for the snapshot
func snapshotName(term, index uint64) string {
	now := time.Now()
	msec := now.UnixNano() / int64(time.Millisecond)
	return fmt.Sprintf("%d-%d-%d", term, index, msec)
}

// Create is used to start a new snapshot
func (f *FileSnapshotStore) Create(index, term uint64, peers []byte) (SnapshotSink, error) {
	// Create a new path
	name := snapshotName(term, index)
	path := filepath.Join(f.path, name+tmpSuffix)
	f.logger.Printf("[INFO] snapshot: Creating new snapshot at %s", path)

	// Make the directory
	if err := os.Mkdir(path, 0755); err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to make snapshot directory: %v", err)
		return nil, err
	}

	// Create the sink
	sink := &FileSnapshotSink{
		store:  f,
		logger: f.logger,
		dir:    path,
		meta: fileSnapshotMeta{
			SnapshotMeta: SnapshotMeta{
				ID:    name,
				Index: index,
				Term:  term,
				Peers: peers,
			},
			CRC: nil,
		},
	}

	// Write out the meta data
	if err := sink.writeMeta(); err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to write metadata: %v", err)
		return nil, err
	}

	// Open the state file
	statePath := filepath.Join(path, stateFilePath)
	fh, err := os.Create(statePath)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to create state file: %v", err)
		return nil, err
	}
	sink.stateFile = fh

	// Create a CRC64 hash
	sink.stateHash = crc64.New(crc64.MakeTable(crc64.ECMA))

	// Wrap both the hash and file in a MultiWriter with buffering
	multi := io.MultiWriter(sink.stateFile, sink.stateHash)
	sink.buffered = bufio.NewWriter(multi)

	// Done
	return sink, nil
}

// List returns available snapshots in the store.
func (f *FileSnapshotStore) List() ([]*SnapshotMeta, error) {
	// Get the eligible snapshots
	snapshots, err := f.getSnapshots()
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get snapshots: %v", err)
		return nil, err
	}

	var snapMeta []*SnapshotMeta
	for _, meta := range snapshots {
		snapMeta = append(snapMeta, &meta.SnapshotMeta)
		if len(snapMeta) == f.retain {
			break
		}
	}
	return snapMeta, nil
}

// getSnapshots returns all the known snapshots
func (f *FileSnapshotStore) getSnapshots() ([]*fileSnapshotMeta, error) {
	// Get the eligible snapshots
	snapshots, err := ioutil.ReadDir(f.path)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to scan snapshot dir: %v", err)
		return nil, err
	}

	// Populate the metadata
	var snapMeta []*fileSnapshotMeta
	for _, snap := range snapshots {
		// Ignore any files
		if !snap.IsDir() {
			continue
		}

		// Ignore any temporary snapshots
		dirName := snap.Name()
		if strings.HasSuffix(dirName, tmpSuffix) {
			f.logger.Printf("[WARN] snapshot: Found temporary snapshot: %v", dirName)
			continue
		}

		// Try to read the meta data
		meta, err := f.readMeta(dirName)
		if err != nil {
			f.logger.Printf("[WARN] snapshot: Failed to read metadata for %v: %v", dirName, err)
			continue
		}

		// Append, but only return up to the retain count
		snapMeta = append(snapMeta, meta)
	}

	// Sort the snapshot, reverse so we get new -> old
	sort.Sort(sort.Reverse(snapMetaSlice(snapMeta)))

	return snapMeta, nil
}

// readMeta is used to read the meta data for a given named backup
func (f *FileSnapshotStore) readMeta(name string) (*fileSnapshotMeta, error) {
	// Open the meta file
	metaPath := filepath.Join(f.path, name, metaFilePath)
	fh, err := os.Open(metaPath)
	if err != nil {
		return nil, err
	}
	defer fh.Close()

	// Buffer the file IO
	buffered := bufio.NewReader(fh)

	// Read in the JSON
	meta := &fileSnapshotMeta{}
	dec := json.NewDecoder(buffered)
	if err := dec.Decode(meta); err != nil {
		return nil, err
	}
	return meta, nil
}

// Open takes a snapshot ID and returns a ReadCloser for that snapshot.
func (f *FileSnapshotStore) Open(id string) (*SnapshotMeta, io.ReadCloser, error) {
	// Get the metadata
	meta, err := f.readMeta(id)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get meta data to open snapshot: %v", err)
		return nil, nil, err
	}

	// Open the state file
	statePath := filepath.Join(f.path, id, stateFilePath)
	fh, err := os.Open(statePath)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to open state file: %v", err)
		return nil, nil, err
	}

	// Create a CRC64 hash
	stateHash := crc64.New(crc64.MakeTable(crc64.ECMA))

	// Compute the hash
	_, err = io.Copy(stateHash, fh)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to read state file: %v", err)
		fh.Close()
		return nil, nil, err
	}

	// Verify the hash
	computed := stateHash.Sum(nil)
	if bytes.Compare(meta.CRC, computed) != 0 {
		f.logger.Printf("[ERR] snapshot: CRC checksum failed (stored: %v computed: %v)",
			meta.CRC, computed)
		fh.Close()
		return nil, nil, fmt.Errorf("CRC mismatch")
	}

	// Seek to the start
	if _, err := fh.Seek(0, 0); err != nil {
		f.logger.Printf("[ERR] snapshot: State file seek failed: %v", err)
		fh.Close()
		return nil, nil, err
	}

	// Return a buffered file
	buffered := &bufferedFile{
		bh: bufio.NewReader(fh),
		fh: fh,
	}

	return &meta.SnapshotMeta, buffered, nil
}

// ReapSnapshots reaps any snapshots beyond the retain count.
func (f *FileSnapshotStore) ReapSnapshots() error {
	snapshots, err := f.getSnapshots()
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get snapshots: %v", err)
		return err
	}

	for i := f.retain; i < len(snapshots); i++ {
		path := filepath.Join(f.path, snapshots[i].ID)
		f.logger.Printf("[INFO] snapshot: reaping snapshot %v", path)
		if err := os.RemoveAll(path); err != nil {
			f.logger.Printf("[ERR] snapshot: Failed to reap snapshot %v: %v", path, err)
			return err
		}
	}
	return nil
}

// ID returns the ID of the snapshot, can be used with Open()
// after the snapshot is finalized.
func (s *FileSnapshotSink) ID() string {
	return s.meta.ID
}

// Write is used to append to the state file. We write to the
// buffered IO object to reduce the amount of context switches
func (s *FileSnapshotSink) Write(b []byte) (int, error) {
	return s.buffered.Write(b)
}

// Close is used to indicate a successful end
func (s *FileSnapshotSink) Close() error {
	// Make sure close is idempotent
	if s.closed {
		return nil
	}
	s.closed = true

	// Close the open handles
	if err := s.finalize(); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to finalize snapshot: %v", err)
		return err
	}

	// Write out the meta data
	if err := s.writeMeta(); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to write metadata: %v", err)
		return err
	}

	// Move the directory into place
	newPath := strings.TrimSuffix(s.dir, tmpSuffix)
	if err := os.Rename(s.dir, newPath); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to move snapshot into place: %v", err)
		return err
	}

	// Reap any old snapshots
	s.store.ReapSnapshots()
	return nil
}

// Cancel is used to indicate an unsuccessful end
func (s *FileSnapshotSink) Cancel() error {
	// Make sure close is idempotent
	if s.closed {
		return nil
	}
	s.closed = true

	// Close the open handles
	if err := s.finalize(); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to finalize snapshot: %v", err)
		return err
	}

	// Attempt to remove all artifacts
	return os.RemoveAll(s.dir)
}

// finalize is used to close all of our resources
func (s *FileSnapshotSink) finalize() error {
	// Flush any remaining data
	if err := s.buffered.Flush(); err != nil {
		return err
	}

	// Get the file size
	stat, statErr := s.stateFile.Stat()

	// Close the file
	if err := s.stateFile.Close(); err != nil {
		return err
	}

	// Set the file size, check after we close
	if statErr != nil {
		return statErr
	}
	s.meta.Size = stat.Size()

	// Set the CRC
	s.meta.CRC = s.stateHash.Sum(nil)
	return nil
}

// writeMeta is used to write out the metadata we have
func (s *FileSnapshotSink) writeMeta() error {
	// Open the meta file
	metaPath := filepath.Join(s.dir, metaFilePath)
	fh, err := os.Create(metaPath)
	if err != nil {
		return err
	}
	defer fh.Close()

	// Buffer the file IO
	buffered := bufio.NewWriter(fh)
	defer buffered.Flush()

	// Write out as JSON
	enc := json.NewEncoder(buffered)
	if err := enc.Encode(&s.meta); err != nil {
		return err
	}
	return nil
}

// Implement the sort interface for []*fileSnapshotMeta
func (s snapMetaSlice) Len() int {
	return len(s)
}

func (s snapMetaSlice) Less(i, j int) bool {
	if s[i].Term != s[j].Term {
		return s[i].Term < s[j].Term
	}
	if s[i].Index != s[j].Index {
		return s[i].Index < s[j].Index
	}
	return s[i].ID < s[j].ID
}

func (s snapMetaSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
<commit_msg>Use MkdirAll instead to be more resilient.<commit_after>package raft

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"hash"
	"hash/crc64"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

const (
	testPath      = "permTest"
	snapPath      = "snapshots"
	metaFilePath  = "meta.json"
	stateFilePath = "state.bin"
	tmpSuffix     = ".tmp"
)

// FileSnapshotStore implements the SnapshotStore interface and allows
// snapshots to be made on the local disk.
type FileSnapshotStore struct {
	path   string
	retain int
	logger *log.Logger
}

type snapMetaSlice []*fileSnapshotMeta

// FileSnapshotSink implements SnapshotSink with a file.
type FileSnapshotSink struct {
	store  *FileSnapshotStore
	logger *log.Logger
	dir    string
	meta   fileSnapshotMeta

	stateFile *os.File
	stateHash hash.Hash64
	buffered  *bufio.Writer

	closed bool
}

// fileSnapshotMeta is stored on disk. We also put a CRC
// on disk so that we can verify the snapshot
type fileSnapshotMeta struct {
	SnapshotMeta
	CRC []byte
}

// bufferedFile is returned when we open a snapshot. This way
// reads are buffered and the file still gets closed.
type bufferedFile struct {
	bh *bufio.Reader
	fh *os.File
}

func (b *bufferedFile) Read(p []byte) (n int, err error) {
	return b.bh.Read(p)
}

func (b *bufferedFile) Close() error {
	return b.fh.Close()
}

// NewFileSnapshotStore creates a new FileSnapshotStore based
// on a base directory. The `retain` parameter controls how many
// snapshots are retained. Must be at least 1.
func NewFileSnapshotStore(base string, retain int, logOutput io.Writer) (*FileSnapshotStore, error) {
	if retain < 1 {
		return nil, fmt.Errorf("must retain at least one snapshot")
	}
	if logOutput == nil {
		logOutput = os.Stderr
	}

	// Ensure our path exists
	path := filepath.Join(base, snapPath)
	if err := os.MkdirAll(path, 0755); err != nil && !os.IsExist(err) {
		return nil, fmt.Errorf("snapshot path not accessible: %v", err)
	}

	// Setup the store
	store := &FileSnapshotStore{
		path:   path,
		retain: retain,
		logger: log.New(logOutput, "", log.LstdFlags),
	}

	// Do a permissions test
	if err := store.testPermissions(); err != nil {
		return nil, fmt.Errorf("permissions test failed: %v", err)
	}
	return store, nil
}

// testPermissions tries to touch a file in our path to see if it works
func (f *FileSnapshotStore) testPermissions() error {
	path := filepath.Join(f.path, testPath)
	fh, err := os.Create(path)
	if err != nil {
		return err
	}
	fh.Close()
	os.Remove(path)
	return nil
}

// snapshotName generate s name for the snapshot
func snapshotName(term, index uint64) string {
	now := time.Now()
	msec := now.UnixNano() / int64(time.Millisecond)
	return fmt.Sprintf("%d-%d-%d", term, index, msec)
}

// Create is used to start a new snapshot
func (f *FileSnapshotStore) Create(index, term uint64, peers []byte) (SnapshotSink, error) {
	// Create a new path
	name := snapshotName(term, index)
	path := filepath.Join(f.path, name+tmpSuffix)
	f.logger.Printf("[INFO] snapshot: Creating new snapshot at %s", path)

	// Make the directory
	if err := os.MkdirAll(path, 0755); err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to make snapshot directory: %v", err)
		return nil, err
	}

	// Create the sink
	sink := &FileSnapshotSink{
		store:  f,
		logger: f.logger,
		dir:    path,
		meta: fileSnapshotMeta{
			SnapshotMeta: SnapshotMeta{
				ID:    name,
				Index: index,
				Term:  term,
				Peers: peers,
			},
			CRC: nil,
		},
	}

	// Write out the meta data
	if err := sink.writeMeta(); err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to write metadata: %v", err)
		return nil, err
	}

	// Open the state file
	statePath := filepath.Join(path, stateFilePath)
	fh, err := os.Create(statePath)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to create state file: %v", err)
		return nil, err
	}
	sink.stateFile = fh

	// Create a CRC64 hash
	sink.stateHash = crc64.New(crc64.MakeTable(crc64.ECMA))

	// Wrap both the hash and file in a MultiWriter with buffering
	multi := io.MultiWriter(sink.stateFile, sink.stateHash)
	sink.buffered = bufio.NewWriter(multi)

	// Done
	return sink, nil
}

// List returns available snapshots in the store.
func (f *FileSnapshotStore) List() ([]*SnapshotMeta, error) {
	// Get the eligible snapshots
	snapshots, err := f.getSnapshots()
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get snapshots: %v", err)
		return nil, err
	}

	var snapMeta []*SnapshotMeta
	for _, meta := range snapshots {
		snapMeta = append(snapMeta, &meta.SnapshotMeta)
		if len(snapMeta) == f.retain {
			break
		}
	}
	return snapMeta, nil
}

// getSnapshots returns all the known snapshots
func (f *FileSnapshotStore) getSnapshots() ([]*fileSnapshotMeta, error) {
	// Get the eligible snapshots
	snapshots, err := ioutil.ReadDir(f.path)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to scan snapshot dir: %v", err)
		return nil, err
	}

	// Populate the metadata
	var snapMeta []*fileSnapshotMeta
	for _, snap := range snapshots {
		// Ignore any files
		if !snap.IsDir() {
			continue
		}

		// Ignore any temporary snapshots
		dirName := snap.Name()
		if strings.HasSuffix(dirName, tmpSuffix) {
			f.logger.Printf("[WARN] snapshot: Found temporary snapshot: %v", dirName)
			continue
		}

		// Try to read the meta data
		meta, err := f.readMeta(dirName)
		if err != nil {
			f.logger.Printf("[WARN] snapshot: Failed to read metadata for %v: %v", dirName, err)
			continue
		}

		// Append, but only return up to the retain count
		snapMeta = append(snapMeta, meta)
	}

	// Sort the snapshot, reverse so we get new -> old
	sort.Sort(sort.Reverse(snapMetaSlice(snapMeta)))

	return snapMeta, nil
}

// readMeta is used to read the meta data for a given named backup
func (f *FileSnapshotStore) readMeta(name string) (*fileSnapshotMeta, error) {
	// Open the meta file
	metaPath := filepath.Join(f.path, name, metaFilePath)
	fh, err := os.Open(metaPath)
	if err != nil {
		return nil, err
	}
	defer fh.Close()

	// Buffer the file IO
	buffered := bufio.NewReader(fh)

	// Read in the JSON
	meta := &fileSnapshotMeta{}
	dec := json.NewDecoder(buffered)
	if err := dec.Decode(meta); err != nil {
		return nil, err
	}
	return meta, nil
}

// Open takes a snapshot ID and returns a ReadCloser for that snapshot.
func (f *FileSnapshotStore) Open(id string) (*SnapshotMeta, io.ReadCloser, error) {
	// Get the metadata
	meta, err := f.readMeta(id)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get meta data to open snapshot: %v", err)
		return nil, nil, err
	}

	// Open the state file
	statePath := filepath.Join(f.path, id, stateFilePath)
	fh, err := os.Open(statePath)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to open state file: %v", err)
		return nil, nil, err
	}

	// Create a CRC64 hash
	stateHash := crc64.New(crc64.MakeTable(crc64.ECMA))

	// Compute the hash
	_, err = io.Copy(stateHash, fh)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to read state file: %v", err)
		fh.Close()
		return nil, nil, err
	}

	// Verify the hash
	computed := stateHash.Sum(nil)
	if bytes.Compare(meta.CRC, computed) != 0 {
		f.logger.Printf("[ERR] snapshot: CRC checksum failed (stored: %v computed: %v)",
			meta.CRC, computed)
		fh.Close()
		return nil, nil, fmt.Errorf("CRC mismatch")
	}

	// Seek to the start
	if _, err := fh.Seek(0, 0); err != nil {
		f.logger.Printf("[ERR] snapshot: State file seek failed: %v", err)
		fh.Close()
		return nil, nil, err
	}

	// Return a buffered file
	buffered := &bufferedFile{
		bh: bufio.NewReader(fh),
		fh: fh,
	}

	return &meta.SnapshotMeta, buffered, nil
}

// ReapSnapshots reaps any snapshots beyond the retain count.
func (f *FileSnapshotStore) ReapSnapshots() error {
	snapshots, err := f.getSnapshots()
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get snapshots: %v", err)
		return err
	}

	for i := f.retain; i < len(snapshots); i++ {
		path := filepath.Join(f.path, snapshots[i].ID)
		f.logger.Printf("[INFO] snapshot: reaping snapshot %v", path)
		if err := os.RemoveAll(path); err != nil {
			f.logger.Printf("[ERR] snapshot: Failed to reap snapshot %v: %v", path, err)
			return err
		}
	}
	return nil
}

// ID returns the ID of the snapshot, can be used with Open()
// after the snapshot is finalized.
func (s *FileSnapshotSink) ID() string {
	return s.meta.ID
}

// Write is used to append to the state file. We write to the
// buffered IO object to reduce the amount of context switches
func (s *FileSnapshotSink) Write(b []byte) (int, error) {
	return s.buffered.Write(b)
}

// Close is used to indicate a successful end
func (s *FileSnapshotSink) Close() error {
	// Make sure close is idempotent
	if s.closed {
		return nil
	}
	s.closed = true

	// Close the open handles
	if err := s.finalize(); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to finalize snapshot: %v", err)
		return err
	}

	// Write out the meta data
	if err := s.writeMeta(); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to write metadata: %v", err)
		return err
	}

	// Move the directory into place
	newPath := strings.TrimSuffix(s.dir, tmpSuffix)
	if err := os.Rename(s.dir, newPath); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to move snapshot into place: %v", err)
		return err
	}

	// Reap any old snapshots
	s.store.ReapSnapshots()
	return nil
}

// Cancel is used to indicate an unsuccessful end
func (s *FileSnapshotSink) Cancel() error {
	// Make sure close is idempotent
	if s.closed {
		return nil
	}
	s.closed = true

	// Close the open handles
	if err := s.finalize(); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to finalize snapshot: %v", err)
		return err
	}

	// Attempt to remove all artifacts
	return os.RemoveAll(s.dir)
}

// finalize is used to close all of our resources
func (s *FileSnapshotSink) finalize() error {
	// Flush any remaining data
	if err := s.buffered.Flush(); err != nil {
		return err
	}

	// Get the file size
	stat, statErr := s.stateFile.Stat()

	// Close the file
	if err := s.stateFile.Close(); err != nil {
		return err
	}

	// Set the file size, check after we close
	if statErr != nil {
		return statErr
	}
	s.meta.Size = stat.Size()

	// Set the CRC
	s.meta.CRC = s.stateHash.Sum(nil)
	return nil
}

// writeMeta is used to write out the metadata we have
func (s *FileSnapshotSink) writeMeta() error {
	// Open the meta file
	metaPath := filepath.Join(s.dir, metaFilePath)
	fh, err := os.Create(metaPath)
	if err != nil {
		return err
	}
	defer fh.Close()

	// Buffer the file IO
	buffered := bufio.NewWriter(fh)
	defer buffered.Flush()

	// Write out as JSON
	enc := json.NewEncoder(buffered)
	if err := enc.Encode(&s.meta); err != nil {
		return err
	}
	return nil
}

// Implement the sort interface for []*fileSnapshotMeta
func (s snapMetaSlice) Len() int {
	return len(s)
}

func (s snapMetaSlice) Less(i, j int) bool {
	if s[i].Term != s[j].Term {
		return s[i].Term < s[j].Term
	}
	if s[i].Index != s[j].Index {
		return s[i].Index < s[j].Index
	}
	return s[i].ID < s[j].ID
}

func (s snapMetaSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
<|endoftext|>


<commit_before>package gitmediafilters

import (
	".."
	"crypto/sha256"
	"encoding/hex"
	"io"
	"os"
	"strconv"
)

type CleanedAsset struct {
	Size          int64
	File          *os.File
	Sha           string
	mediafilepath string
}

func Clean(reader io.Reader) (*CleanedAsset, error) {
	tmp, err := gitmedia.TempFile()
	if err != nil {
		return nil, err
	}

	oidHash := sha256.New()
	writer := io.MultiWriter(oidHash, tmp)
	written, err := io.Copy(writer, reader)
	oidHash.Write([]byte(strconv.FormatInt(written, 10)))

	return &CleanedAsset{written, tmp, hex.EncodeToString(oidHash.Sum(nil)), ""}, err
}

func (a *CleanedAsset) Close() error {
	return os.Remove(a.File.Name())
}
<commit_msg>only hash the file content<commit_after>package gitmediafilters

import (
	".."
	"crypto/sha256"
	"encoding/hex"
	"io"
	"os"
)

type CleanedAsset struct {
	Size          int64
	File          *os.File
	Sha           string
	mediafilepath string
}

func Clean(reader io.Reader) (*CleanedAsset, error) {
	tmp, err := gitmedia.TempFile()
	if err != nil {
		return nil, err
	}

	oidHash := sha256.New()
	writer := io.MultiWriter(oidHash, tmp)
	written, err := io.Copy(writer, reader)

	return &CleanedAsset{written, tmp, hex.EncodeToString(oidHash.Sum(nil)), ""}, err
}

func (a *CleanedAsset) Close() error {
	return os.Remove(a.File.Name())
}
<|endoftext|>


<commit_before>package firebase_test

import (
	"encoding/json"
	"fmt"
	"os"
	"reflect"
	"testing"
	"time"

	"code.google.com/p/gomock/gomock"

	"github.com/JustinTulloss/firebase"
	"github.com/JustinTulloss/firebase/mock_firebase"
)

type Name struct {
	First string `json:",omitempty"`
	Last  string `json:",omitempty"`
}

func nameAlloc() interface{} {
	return &Name{}
}

/*
Set the two variables below and set them to your own
Firebase URL and credentials (optional) if you're forking the code
and want to test your changes.
*/

// enter your firebase credentials for testing.
var (
	testUrl  string = os.Getenv("FIREBASE_TEST_URL")
	testAuth string = os.Getenv("FIREBASE_TEST_AUTH")
)

func TestValue(t *testing.T) {
	client := firebase.NewClient(testUrl+"/tests", testAuth, nil)

	var r map[string]interface{}
	err := client.Value(&r)
	if err != nil {
		t.Error(err)
	}

	if r == nil {
		t.Fatalf("No values returned from the server\n")
	}
}

func TestChild(t *testing.T) {
	client := firebase.NewClient(testUrl+"/tests", testAuth, nil)

	r := client.Child("")

	if r == nil {
		t.Fatalf("No child returned from the server\n")
	}
}

func TestPush(t *testing.T) {
	client := firebase.NewClient(testUrl+"/tests", testAuth, nil)

	name := &Name{First: "FirstName", Last: "LastName"}

	r, err := client.Push(name, nil)

	if err != nil {
		t.Fatalf("%v\n", err)
	}

	if r == nil {
		t.Fatalf("No client returned from the server\n")
	}

	newName := &Name{}
	c2 := firebase.NewClient(r.String(), testAuth, nil)
	c2.Value(newName)
	if !reflect.DeepEqual(name, newName) {
		t.Errorf("Expected %v to equal %v", name, newName)
	}
}

func TestSet(t *testing.T) {
	c1 := firebase.NewClient(testUrl+"/tests/users", testAuth, nil)

	name := &Name{First: "First", Last: "last"}
	c2, _ := c1.Push(name, nil)

	newName := &Name{First: "NewFirst", Last: "NewLast"}
	r, err := c2.Set("", newName, map[string]string{"print": "silent"})

	if err != nil {
		t.Fatalf("%v\n", err)
	}

	if r == nil {
		t.Fatalf("No client returned from the server\n")
	}
}

func TestUpdate(t *testing.T) {
	c1 := firebase.NewClient(testUrl+"/tests/users", testAuth, nil)

	name := &Name{First: "First", Last: "last"}
	c2, _ := c1.Push(name, nil)

	newName := &Name{Last: "NewLast"}
	err := c2.Update("", newName, nil)

	if err != nil {
		t.Fatalf("%v\n", err)
	}
}

func TestRemovet(t *testing.T) {
	c1 := firebase.NewClient(testUrl+"/tests/users", testAuth, nil)

	name := &Name{First: "First", Last: "last"}
	c2, _ := c1.Push(name, nil)

	err := c2.Remove("", nil)
	if err != nil {
		t.Fatalf("%v\n", err)
	}

	var val map[string]interface{}
	c3 := firebase.NewClient(c2.String(), testAuth, nil)
	err = c3.Value(&val)
	if err != nil {
		t.Error(err)
	}

	if len(val) != 0 {
		t.Errorf("Expected %s to be removed, was %v", c2.String(), val)
	}
}

func TestRules(t *testing.T) {
	client := firebase.NewClient(testUrl, testAuth, nil)

	r, err := client.Rules(nil)

	if err != nil {
		t.Fatalf("Error retrieving rules: %v\n", err)
	}

	if r == nil {
		t.Fatalf("No child returned from the server\n")
	}
}

func TestSetRules(t *testing.T) {
	client := firebase.NewClient(testUrl, testAuth, nil)

	rules := &firebase.Rules{
		"rules": map[string]interface{}{
			".read":  "auth.username == 'admin'",
			".write": "auth.username == 'admin'",
			"ordered": map[string]interface{}{
				".indexOn": []string{"First"},
				"kids": map[string]interface{}{
					".indexOn": []string{"Age"},
				},
			},
		},
	}

	err := client.SetRules(rules, nil)

	if err != nil {
		t.Fatalf("Error setting rules: %v\n", err)
	}
}

func TestOrderBy(t *testing.T) {
	client := firebase.NewClient(testUrl, testAuth, nil).Child("ordered")
	defer client.Remove("", nil)

	names := []*Name{
		&Name{First: "BBBB", Last: "YYYY"},
		&Name{First: "AAAA", Last: "ZZZZZ"},
	}

	for _, n := range names {
		_, err := client.Push(n, nil)
		if err != nil {
			t.Fatalf("Couldn't push new name: %s\n", err)
		}
	}

	i := 0
	for n := range client.OrderBy(firebase.KeyProp).Iterator(nameAlloc) {
		if n.Value.(*Name).First != names[i].First {
			t.Fatalf("Key order was not delivered")
		}
		i++
	}
	if i == 0 {
		t.Fatalf("Did not receive names ordered by key")
	}

	expectedOrder := []*Name{names[1], names[0]}
	i = 0
	for n := range client.OrderBy("First").Iterator(nameAlloc) {
		if n.Value.(*Name).First != expectedOrder[i].First {
			t.Fatalf("Child prop order was not delivered")
		}
		i++
	}
	if i == 0 {
		t.Fatalf("Did not receive names ordered by first name")
	}

	kids := map[string]map[string]interface{}{
		"a": map[string]interface{}{"Name": "Bob", "Age": 14},
		"b": map[string]interface{}{"Name": "Alice", "Age": 13},
	}
	_, err := client.Set("kids", kids, nil)
	if err != nil {
		t.Fatalf("Could not set kids: %s\n", err)
	}

	expectedKidsOrder := []map[string]interface{}{kids["b"], kids["a"]}
	i = 0
	for n := range client.Child("kids").OrderBy("Age").Iterator(nil) {
		if (*n.Value.(*map[string]interface{}))["First"] != expectedKidsOrder[i]["First"] {
			t.Fatalf("Child prop order by age was not delivered")
		}
		i++
	}
	if i == 0 {
		t.Fatalf("Did not receive names ordered by age")
	}
}

func TestTimestamp(t *testing.T) {
	ts := firebase.Timestamp(time.Now())
	marshaled, err := json.Marshal(&ts)
	if err != nil {
		t.Fatalf("Could not marshal a timestamp to json: %s\n", err)
	}
	unmarshaledTs := firebase.Timestamp{}
	err = json.Unmarshal(marshaled, &unmarshaledTs)
	if err != nil {
		t.Fatalf("Could not unmarshal a timestamp to json: %s\n", err)
	}
	// Compare unix timestamps as we lose some fidelity in the nanoseconds
	if time.Time(ts).Unix() != time.Time(unmarshaledTs).Unix() {
		t.Fatalf("Unmarshaled time %s not equivalent to marshaled time %s",
			unmarshaledTs,
			ts,
		)
	}
}

func TestServerTimestamp(t *testing.T) {
	b, err := json.Marshal(firebase.ServerTimestamp)
	if err != nil {
		t.Fatalf("Could not marshal server timestamp: %s\n", err)
	}
	if string(b) != `{".sv":"timestamp"}` {
		t.Fatalf("Unexpected timestamp json value: %s\n", b)
	}
}

func TestIterator(t *testing.T) {
	client := firebase.NewClient(testUrl+"/test-iterator", testAuth, nil)
	defer client.Remove("", nil)
	names := []Name{
		{First: "FirstName", Last: "LastName"},
		{First: "Second", Last: "Seconder"},
	}
	for _, name := range names {
		_, err := client.Push(name, nil)
		if err != nil {
			t.Fatalf("%v\n", err)
		}
	}

	var i = 0
	for nameEntry := range client.Iterator(nameAlloc) {
		name := nameEntry.Value.(*Name)
		if !reflect.DeepEqual(&names[i], name) {
			t.Errorf("Expected %v to equal %v", &names[i], name)
		}
		i++
	}
	if i != len(names) {
		t.Fatalf("Did not receive all names, received %d\n", i)
	}
}

func TestShallow(t *testing.T) {
	client := firebase.NewClient(testUrl+"/test-shallow", testAuth, nil)
	defer client.Remove("", nil)
	client.Set("one", 1, nil)
	client.Set("two", 2, nil)
	keys, err := client.Shallow()
	if err != nil {
		t.Errorf("Error when calling shallow: %s\n", err)
	}
	if !reflect.DeepEqual(keys, []string{"one", "two"}) {
		t.Errorf("keys (%v) were not correct\n", keys)
	}
}

func TestKey(t *testing.T) {
	client := firebase.NewClient(testUrl+"/test", testAuth, nil)
	if client.Key() != "test" {
		t.Errorf("Key should have been test, was %s\n", client.Key())
	}

	client = firebase.NewClient(testUrl, testAuth, nil)
	if client.Key() != "" {
		t.Errorf("Key should have been empty, was %s\n", client.Key())
	}

	client = client.Child("/a/b/c/d/e/f/g")
	if client.Key() != "g" {
		t.Errorf("Key should have been 'g', was %s\n", client.Key())
	}
}

func TestMockable(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	mockFire := mock_firebase.NewMockClient(mockCtrl)
	mockFire.EXPECT().Child("test")
	mockFire.Child("test")
}

func TestMain(m *testing.M) {
	if testUrl == "" || testAuth == "" {
		fmt.Printf("You need to set FIREBASE_TEST_URL and FIREBASE_TEST_AUTH\n")
		os.Exit(1)
	}
	os.Exit(m.Run())
}
<commit_msg>Convert unit tests to ginkgo/gomega<commit_after>package firebase

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var (
	testUrl  = ""
	testAuth = ""
)

type Name struct {
	First string `json:",omitempty"`
	Last  string `json:",omitempty"`
}

func nameAlloc() interface{} {
	return &Name{}
}

func fakeServer(handler http.Handler) (*httptest.Server, *client) {
	testServer := httptest.NewServer(handler)

	c := NewClient(testServer.URL, testAuth, nil)
	testClient, isClient := c.(*client)
	Expect(isClient).To(BeTrue())

	return testServer, testClient
}

var _ = Describe("Transforming client urls/queries", func() {
	var (
		c        *client
		isClient bool
		testURL  string = "https://who.cares.com"
	)

	BeforeEach(func() {
		c, isClient = NewClient(testURL, testAuth, nil).(*client)
		Expect(isClient).To(BeTrue())
	})

	It("Adds the child path to the returned client object", func() {
		child, isClient := c.Child("child").(*client)
		Expect(isClient).To(BeTrue())

		Expect(child.url).To(Equal(testURL + "/child"))
	})

	It("Sets a query string param to ask for a shallow object", func() {
		shallow, isClient := c.Shallow().(*client)
		Expect(isClient).To(BeTrue())

		Expect(shallow.params["shallow"]).To(Equal("true"))
	})

	It("Retrieves the key of a client", func() {
		Expect(c.Key()).To(Equal(""))
		Expect(c.Child("test").Key()).To(Equal("test"))
		Expect(c.Child("/a/b/c/d/e/f/g").Key()).To(Equal("g"))
	})
})

var _ = Describe("Manipulating values from firebase", func() {
	var (
		testResource Name
		testServer   *httptest.Server
		testClient   *client
		handler      func(w http.ResponseWriter, r *http.Request)
	)

	BeforeEach(func() {
		testResource = Name{First: "FirstName", Last: "LastName"}
	})

	JustBeforeEach(func() {
		testServer, testClient = fakeServer(http.HandlerFunc(handler))
	})

	Context("Retrieving a value from firebase", func() {
		BeforeEach(func() {
			handler = func(w http.ResponseWriter, r *http.Request) {
				Expect(r.Method).To(Equal("GET"))
				fmt.Fprintln(w, `{"bru": "haha"}`)
			}
		})

		It("Retrieves the expected value from the resource path", func() {
			var r map[string]interface{}
			err := testClient.Child("").Value(&r)
			Expect(err).To(BeNil())

			Expect(len(r)).To(Equal(1))
			Expect(r["bru"]).To(Equal("haha"))
		})
	})

	Context("Pushing a new value to firebase", func() {
		var (
			pushedName string = "baloo"
		)

		BeforeEach(func() {
			handler = func(w http.ResponseWriter, r *http.Request) {
				Expect(r.Method).To(Equal("POST"))

				var pushed Name
				defer r.Body.Close()

				decoder := json.NewDecoder(r.Body)
				err := decoder.Decode(&pushed)
				Expect(err).To(BeNil())
				Expect(pushed).To(Equal(testResource))

				fmt.Fprintf(w, `{"name": "%s"}`, pushedName)
			}
		})

		It("Pushes the new resource and returns a matching client", func() {
			name := &testResource

			response, err := testClient.Child("path").Push(name, nil)
			Expect(err).To(BeNil())

			responseClient, isClient := response.(*client)
			Expect(isClient).To(BeTrue())
			Expect(responseClient.url).To(Equal(testServer.URL + "/path/" +
				pushedName))
		})
	})

	Context("Setting an existing value in firebase", func() {
		var (
			newName Name   = Name{First: "NewFirst", Last: "NewLast"}
			setPath string = "set"
		)

		BeforeEach(func() {
			handler = func(w http.ResponseWriter, r *http.Request) {
				Expect(r.Method).To(Equal("PUT"))

				var setValue Name
				defer r.Body.Close()

				decoder := json.NewDecoder(r.Body)
				err := decoder.Decode(&setValue)
				Expect(err).To(BeNil())
				Expect(setValue).To(Equal(newName))
			}
		})

		It("Overwrites the value of the existing resource", func() {
			response, err := testClient.Set(setPath, &newName, nil)
			Expect(err).To(BeNil())

			responseClient, isClient := response.(*client)
			Expect(isClient).To(BeTrue())
			Expect(responseClient.url).To(Equal(testServer.URL + "/" +
				setPath))
		})
	})

	Context("Update an existing value in firebase", func() {
		var (
			updatedName Name   = Name{Last: "NewLast"}
			updatePath  string = "update"
		)

		BeforeEach(func() {
			handler = func(w http.ResponseWriter, r *http.Request) {
				Expect(r.Method).To(Equal("PATCH"))
				Expect(r.URL.String()).To(Equal("/" + updatePath + ".json"))

				var updateValue Name
				defer r.Body.Close()

				decoder := json.NewDecoder(r.Body)
				err := decoder.Decode(&updateValue)
				Expect(err).To(BeNil())
				Expect(updateValue).To(Equal(updatedName))
			}
		})

		It("Changes the value of the existing resource", func() {
			err := testClient.Update(updatePath, &updatedName, nil)
			Expect(err).To(BeNil())
		})
	})

	Context("Delete an existing value in firebase", func() {
		var (
			rmPath string = "update"
		)

		BeforeEach(func() {
			handler = func(w http.ResponseWriter, r *http.Request) {
				Expect(r.Method).To(Equal("DELETE"))
				Expect(r.URL.String()).To(Equal("/" + rmPath + ".json"))
			}
		})

		It("Deletes the resource", func() {
			err := testClient.Remove(rmPath, nil)
			Expect(err).To(BeNil())
		})
	})

	Context("Reading the security rules", func() {
		var testRules Rules = make(map[string]interface{})

		BeforeEach(func() {
			testRules["rules"] = "anything goes"

			handler = func(w http.ResponseWriter, r *http.Request) {
				Expect(r.Method).To(Equal("GET"))
				Expect(r.URL.String()).To(Equal("/.settings/rules.json"))

				encoder := json.NewEncoder(w)
				err := encoder.Encode(testRules)
				Expect(err).To(BeNil())
			}
		})

		It("Retrieves the firebase's security rules", func() {
			rules, err := testClient.Rules(nil)
			Expect(err).To(BeNil())
			Expect(*rules).To(Equal(testRules))
		})
	})

	Context("Changing the firebase's security rules", func() {
		var newRules Rules

		BeforeEach(func() {
			newRules = Rules{
				"rules": map[string]interface{}{
					".read":  "auth.username == 'admin'",
					".write": "auth.username == 'admin'",
					"ordered": map[string]interface{}{
						".indexOn": []string{"First"},
						"kids": map[string]interface{}{
							".indexOn": []string{"Age"},
						},
					},
				},
			}

			handler = func(w http.ResponseWriter, r *http.Request) {
				var changedRules Rules
				Expect(r.Method).To(Equal("PUT"))
				Expect(r.URL.String()).To(Equal("/.settings/rules.json"))

				defer r.Body.Close()
				decoder := json.NewDecoder(r.Body)
				err := decoder.Decode(&changedRules)
				Expect(err).To(BeNil())
			}
		})

		It("Changes the security rules", func() {
			err := testClient.SetRules(&newRules, nil)
			Expect(err).To(BeNil())
		})
	})
})

var _ = Describe("Firebase timestamps", func() {
	It("Marshals a timestamp into ms since the epoch", func() {
		ts := Timestamp(time.Now())
		marshaled, err := json.Marshal(&ts)
		Expect(err).To(BeNil())

		unmarshaledTs := Timestamp{}
		err = json.Unmarshal(marshaled, &unmarshaledTs)
		Expect(err).To(BeNil())

		// Compare unix timestamps as we lose some fidelity in the nanoseconds
		Expect(time.Time(ts).Unix()).To(Equal(time.Time(unmarshaledTs).Unix()))
	})

	It("Marhsals a server-side timestamp", func() {
		text, err := json.Marshal(ServerTimestamp)
		Expect(err).To(BeNil())
		Expect(string(text)).To(Equal(`{".sv":"timestamp"}`))
	})
})

func TestOrderBy(t *testing.T) {
	t.Skip("needs httptest")
	client := NewClient(testUrl, testAuth, nil).Child("ordered")
	defer client.Remove("", nil)

	names := []*Name{
		&Name{First: "BBBB", Last: "YYYY"},
		&Name{First: "AAAA", Last: "ZZZZZ"},
	}

	for _, n := range names {
		_, err := client.Push(n, nil)
		if err != nil {
			t.Fatalf("Couldn't push new name: %s\n", err)
		}
	}

	i := 0
	for n := range client.OrderBy(KeyProp).Iterator(nameAlloc) {
		if n.Value.(*Name).First != names[i].First {
			t.Fatalf("Key order was not delivered")
		}
		i++
	}
	if i == 0 {
		t.Fatalf("Did not receive names ordered by key")
	}

	expectedOrder := []*Name{names[1], names[0]}
	i = 0
	for n := range client.OrderBy("First").Iterator(nameAlloc) {
		if n.Value.(*Name).First != expectedOrder[i].First {
			t.Fatalf("Child prop order was not delivered")
		}
		i++
	}
	if i == 0 {
		t.Fatalf("Did not receive names ordered by first name")
	}

	kids := map[string]map[string]interface{}{
		"a": map[string]interface{}{"Name": "Bob", "Age": 14},
		"b": map[string]interface{}{"Name": "Alice", "Age": 13},
	}
	_, err := client.Set("kids", kids, nil)
	if err != nil {
		t.Fatalf("Could not set kids: %s\n", err)
	}

	expectedKidsOrder := []map[string]interface{}{kids["b"], kids["a"]}
	i = 0
	for n := range client.Child("kids").OrderBy("Age").Iterator(nil) {
		if (*n.Value.(*map[string]interface{}))["First"] != expectedKidsOrder[i]["First"] {
			t.Fatalf("Child prop order by age was not delivered")
		}
		i++
	}
	if i == 0 {
		t.Fatalf("Did not receive names ordered by age")
	}
}

func TestIterator(t *testing.T) {
	t.Skip("needs httptest")
	client := NewClient(testUrl+"/test-iterator", testAuth, nil)
	defer client.Remove("", nil)
	names := []Name{
		{First: "FirstName", Last: "LastName"},
		{First: "Second", Last: "Seconder"},
	}
	for _, name := range names {
		_, err := client.Push(name, nil)
		if err != nil {
			t.Fatalf("%v\n", err)
		}
	}

	var i = 0
	for nameEntry := range client.Iterator(nameAlloc) {
		name := nameEntry.Value.(*Name)
		if !reflect.DeepEqual(&names[i], name) {
			t.Errorf("Expected %v to equal %v", &names[i], name)
		}
		i++
	}
	if i != len(names) {
		t.Fatalf("Did not receive all names, received %d\n", i)
	}
}

func TestFirebase(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Firebase Suite")
}
<|endoftext|>


<commit_before>package main

import (
	"os/exec"
	"strings"
	"encoding/gob"
	"fmt"
	"bufio"
	"io"
	"log"
	"net"
	"os"

	"github.com/flynn/go-discover/discover"
	"github.com/flynn/lorne/types"
	"github.com/titanous/go-dockerclient"
	"github.com/flynn/sampi/types"
	sc "github.com/flynn/sampi/client"
	lc "github.com/flynn/lorne/client"
)

// WARNING: assumes one host at the moment (firstHost will always be the same)

var sd *discover.Client
var sched *sc.Client
var host *lc.Client
var hostid string

func init() {
	var err error
	sd, err = discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}

	sched, err = sc.New()
	if err != nil {
		log.Fatal(err)
	}

	hostid = findHost()
	host, err = lc.New(hostid)
	if err != nil {
		log.Fatal(err)
	}
}

func main() {
	root := "/var/lib/demo/apps"
	hostname := shell("curl -s icanhazip.com")

	set, _ := sd.Services("shelf")
	addrs := set.OnlineAddrs()
	if len(addrs) < 1 {
		panic("Shelf is not discoverable")
	}
	shelfHost := addrs[0]

	app := os.Args[2]
	os.MkdirAll(root+"/"+app, 0755)

	fmt.Printf("-----> Building %s on %s ...\n", app, hostname)

	scheduleAndAttach(app + ".build", docker.Config{
		Image:        "flynn/slugbuilder",
		Cmd:          []string{"http://" + shelfHost + "/" + app + ".tgz"},
		Tty:          false,
		AttachStdin:  true,
		AttachStdout: true,
		AttachStderr: true,
		OpenStdin:    true,
		StdinOnce:    true,
	})

	fmt.Printf("-----> Deploying %s ...\n", app)

	jobid := app + ".web"

	stopIfExists(jobid)
	scheduleWithTcpPort(jobid, docker.Config{
		Image:        "flynn/slugrunner",
		Cmd:          []string{"start web"},
		Tty:          false,
		AttachStdin:  false,
		AttachStdout: false,
		AttachStderr: false,
		OpenStdin:    false,
		StdinOnce:    false,
		Env: []string{
			"SLUG_URL=http://" + shelfHost + "/" + app + ".tgz",
		},
	})

	fmt.Printf("=====> Application deployed!\n")
	//fmt.Printf("       %s\n", getUrl(jobid))
	fmt.Println("")

}

func shell(cmdline string) string {
        out, err := exec.Command("bash", "-c", cmdline).Output()
        if err != nil {
                panic(err)
        }
        return strings.Trim(string(out), " \n")
}

func stopIfExists(jobid string) {
	_, err := host.GetJob(jobid)
	if err != nil {
		return
	}
	if err := host.StopJob(jobid); err != nil {
		return
	}
}

func scheduleWithTcpPort(jobid string, config docker.Config) {

	schedReq := &sampi.ScheduleReq{
		Incremental: true,
		HostJobs:    map[string][]*sampi.Job{hostid: {{ID: jobid, Config: &config, TCPPorts: 1}}},
	}
	if _, err := sched.Schedule(schedReq); err != nil {
		log.Fatal(err)
	}
}

/*func getUrl(jobid string) string {
	if job, err := host.GetJob(jobid); err != nil {
		log.Fatal(err)
	}
}*/

func findHost() string {
	state, err := sched.State()
	if err != nil {
		log.Fatal(err)
	}

	var firstHost string
	for k := range state {
		firstHost = k
		break
	}
	if firstHost == "" {
		log.Fatal("no hosts")
	}
	return firstHost
}

func scheduleAndAttach(jobid string, config docker.Config) {

	services, err := sd.Services("flynn-lorne-attach." + hostid)
	if err != nil {
		log.Fatal(err)
	}
	conn, err := net.Dial("tcp", services.OnlineAddrs()[0])
	if err != nil {
		log.Fatal(err)
	}
	err = gob.NewEncoder(conn).Encode(&lorne.AttachReq{
		JobID: jobid,
		Flags: lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagStdin | lorne.AttachFlagStream,
	})
	if err != nil {
		log.Fatal(err)
	}
	attachState := make([]byte, 1)
	if _, err := conn.Read(attachState); err != nil {
		log.Fatal(err)
	}
	switch attachState[0] {
	case lorne.AttachError:
		log.Fatal("attach error")
	}

	schedReq := &sampi.ScheduleReq{
		Incremental: true,
		HostJobs:    map[string][]*sampi.Job{hostid: {{ID: jobid, Config: &config}}},
	}
	if _, err := sched.Schedule(schedReq); err != nil {
		log.Fatal(err)
	}

	if _, err := conn.Read(attachState); err != nil {
		log.Fatal(err)
	}

	go func() {
		io.Copy(conn, os.Stdin)
		conn.(*net.TCPConn).CloseWrite()
	}()
	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		fmt.Fprintln(os.Stdout, scanner.Text()[8:])
	}
	conn.Close()
}
<commit_msg>added lorne and sampi to project<commit_after>package main

import (
	"os/exec"
	"strings"
	"encoding/gob"
	"fmt"
	"bufio"
	"io"
	"log"
	"net"
	"os"

	"github.com/flynn/go-discover/discover"
	"github.com/flynn/lorne/types"
	"github.com/titanous/go-dockerclient"
	"github.com/flynn/sampi/types"
	sc "github.com/flynn/sampi/client"
	lc "github.com/flynn/lorne/client"
)

// WARNING: assumes one host at the moment (firstHost will always be the same)

var sd *discover.Client
var sched *sc.Client
var host *lc.Client
var hostid string

func init() {
	var err error
	sd, err = discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}

	sched, err = sc.New()
	if err != nil {
		log.Fatal(err)
	}

	hostid = findHost()
	host, err = lc.New(hostid)
	if err != nil {
		log.Fatal(err)
	}
}

func main() {
	root := "/var/lib/demo/apps"
	hostname := shell("curl -s icanhazip.com")

	set, _ := sd.Services("shelf")
	addrs := set.OnlineAddrs()
	if len(addrs) < 1 {
		panic("Shelf is not discoverable")
	}
	shelfHost := addrs[0]

	app := os.Args[2]
	os.MkdirAll(root+"/"+app, 0755)

	fmt.Printf("-----> Building %s on %s ...\n", app, hostname)

	scheduleAndAttach(app + ".build", docker.Config{
		Image:        "flynn/slugbuilder",
		Cmd:          []string{"http://" + shelfHost + "/" + app + ".tgz"},
		Tty:          false,
		AttachStdin:  true,
		AttachStdout: true,
		AttachStderr: true,
		OpenStdin:    true,
		StdinOnce:    true,
	})

	fmt.Printf("-----> Deploying %s ...\n", app)

	jobid := app + ".web"

	stopIfExists(jobid)
	scheduleWithTcpPort(jobid, docker.Config{
		Image:        "flynn/slugrunner",
		Cmd:          []string{"start", "web"},
		Tty:          false,
		AttachStdin:  false,
		AttachStdout: false,
		AttachStderr: false,
		OpenStdin:    false,
		StdinOnce:    false,
		Env: []string{
			"SLUG_URL=http://" + shelfHost + "/" + app + ".tgz",
		},
	})

	fmt.Printf("=====> Application deployed!\n")
	//fmt.Printf("       %s\n", getUrl(jobid))
	fmt.Println("")

}

func shell(cmdline string) string {
        out, err := exec.Command("bash", "-c", cmdline).Output()
        if err != nil {
                panic(err)
        }
        return strings.Trim(string(out), " \n")
}

func stopIfExists(jobid string) {
	_, err := host.GetJob(jobid)
	if err != nil {
		return
	}
	if err := host.StopJob(jobid); err != nil {
		return
	}
}

func scheduleWithTcpPort(jobid string, config docker.Config) {

	schedReq := &sampi.ScheduleReq{
		Incremental: true,
		HostJobs:    map[string][]*sampi.Job{hostid: {{ID: jobid, Config: &config, TCPPorts: 1}}},
	}
	if _, err := sched.Schedule(schedReq); err != nil {
		log.Fatal(err)
	}
}

/*func getUrl(jobid string) string {
	if job, err := host.GetJob(jobid); err != nil {
		log.Fatal(err)
	}
}*/

func findHost() string {
	state, err := sched.State()
	if err != nil {
		log.Fatal(err)
	}

	var firstHost string
	for k := range state {
		firstHost = k
		break
	}
	if firstHost == "" {
		log.Fatal("no hosts")
	}
	return firstHost
}

func scheduleAndAttach(jobid string, config docker.Config) {

	services, err := sd.Services("flynn-lorne-attach." + hostid)
	if err != nil {
		log.Fatal(err)
	}
	conn, err := net.Dial("tcp", services.OnlineAddrs()[0])
	if err != nil {
		log.Fatal(err)
	}
	err = gob.NewEncoder(conn).Encode(&lorne.AttachReq{
		JobID: jobid,
		Flags: lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagStdin | lorne.AttachFlagStream,
	})
	if err != nil {
		log.Fatal(err)
	}
	attachState := make([]byte, 1)
	if _, err := conn.Read(attachState); err != nil {
		log.Fatal(err)
	}
	switch attachState[0] {
	case lorne.AttachError:
		log.Fatal("attach error")
	}

	schedReq := &sampi.ScheduleReq{
		Incremental: true,
		HostJobs:    map[string][]*sampi.Job{hostid: {{ID: jobid, Config: &config}}},
	}
	if _, err := sched.Schedule(schedReq); err != nil {
		log.Fatal(err)
	}

	if _, err := conn.Read(attachState); err != nil {
		log.Fatal(err)
	}

	go func() {
		io.Copy(conn, os.Stdin)
		conn.(*net.TCPConn).CloseWrite()
	}()
	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		fmt.Fprintln(os.Stdout, scanner.Text()[8:])
	}
	conn.Close()
}
<|endoftext|>


<commit_before>package dependencies

import (
	"bytes"
	"fmt"
	"github.com/jfrog/jfrog-cli-go/jfrog-cli/artifactory/utils"
	golangutil "github.com/jfrog/jfrog-cli-go/jfrog-cli/artifactory/utils/golang"
	"github.com/jfrog/jfrog-cli-go/jfrog-cli/utils/config"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/artifactory/buildinfo"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/artifactory/services/vgo"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/utils/errorutils"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/utils/io/fileutils"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/utils/io/fileutils/checksum"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/utils/log"
	"io/ioutil"
	"path/filepath"
	"regexp"
	"strings"
)

func Load() ([]Dependency, error) {
	goPath, err := getGOPATH()
	if err != nil {
		return nil, errorutils.CheckError(err)
	}
	cachePath := filepath.Join(goPath, "src", "v", "cache")
	return getDependencies(cachePath)
}

// Represent vgo dependency project.
// Includes publishing capabilities and build info dependencies.
type Dependency struct {
	buildInfoDependencies []buildinfo.Dependency
	id                    string
	modContent            []byte
	zipPath               string
	version               string
}

func (dependency *Dependency) GetId() string {
	return dependency.id
}

func (dependency *Dependency) Publish(targetRepo string, details *config.ArtifactoryDetails) error {
	log.Info("Publishing:", dependency.id, "to", targetRepo)
	servicesManager, err := utils.CreateServiceManager(details, false)
	if err != nil {
		return err
	}
	params := &vgo.VgoParamsImpl{}
	params.ZipPath = dependency.zipPath
	params.ModContent = dependency.modContent
	params.Version = dependency.version
	params.TargetRepo = targetRepo

	return servicesManager.PublishVgoProject(params)
}

func (dependency *Dependency) Dependencies() []buildinfo.Dependency {
	return dependency.buildInfoDependencies
}

func getDependencies(cachePath string) ([]Dependency, error) {
	vgoCmd, err := golangutil.NewCmd()
	if err != nil {
		return nil, err
	}
	vgoCmd.Command = []string{"list"}
	vgoCmd.CommandFlags = []string{"-m"}
	output, err := utils.RunCmdOutput(vgoCmd)
	if err != nil {
		return nil, errorutils.CheckError(err)
	}

	nameVersionMap, err := parseListOutput(output)
	if err != nil {
		return nil, nil
	}

	deps := []Dependency{}
	for name, ver := range nameVersionMap {
		dep, err := createDependency(cachePath, name, ver)
		if err != nil {
			return nil, err
		}
		if dep != nil {
			deps = append(deps, *dep)
		}
	}
	return deps, nil
}

// Creates a vgo dependency.
// Returns a nil value in case the dependency does not include a zip in the cache.
func createDependency(cachePath, dependencyName, version string) (*Dependency, error) {
	// We first check if the this dependency has a zip binary in the local vgo cache.
	// Core golang dependencies do not have a binary, and we should therefore skip them.
	zipPath := filepath.Join(cachePath, dependencyName, "@v", version+".zip")
	fileExists, err := fileutils.IsFileExists(zipPath)
	if err != nil {
		return nil, err
	}
	// Zip binary does not exist, so we skip it by returning a nil dependency.
	if !fileExists {
		return nil, nil
	}

	dep := Dependency{}

	dep.id = strings.Join([]string{dependencyName, version}, ":")
	dep.version = version
	dep.zipPath = zipPath
	dep.modContent, err = ioutil.ReadFile(filepath.Join(cachePath, dependencyName, "@v", version+".mod"))
	if err != nil {
		return &dep, errorutils.CheckError(err)
	}

	// Mod file dependency
	modDependency := buildinfo.Dependency{Id: dep.id}
	checksums, err := checksum.Calc(bytes.NewBuffer(dep.modContent))
	if err != nil {
		return &dep, err
	}
	modDependency.Checksum = &buildinfo.Checksum{Sha1: checksums[checksum.SHA1], Md5: checksums[checksum.MD5]}

	// Zip file dependency
	zipDependency := buildinfo.Dependency{Id: dep.id}
	fileDetails, err := fileutils.GetFileDetails(dep.zipPath)
	if err != nil {
		return &dep, err
	}
	zipDependency.Checksum = &buildinfo.Checksum{Sha1: fileDetails.Checksum.Sha1, Md5: fileDetails.Checksum.Md5}

	dep.buildInfoDependencies = append(dep.buildInfoDependencies, modDependency, zipDependency)
	return &dep, nil
}

func parseListOutput(content []byte) (map[string]string, error) {
	depRegexp, err := regexp.Compile("(\\S+)\\s+(\\S+)")
	if err != nil {
		return nil, errorutils.CheckError(err)
	}

	depMap := map[string]string{}
	lines := bytes.Split(content, []byte("\n"))
	for i := 2; i < len(lines); i++ {
		dependency := depRegexp.FindStringSubmatch(string(lines[i]))
		if len(dependency) == 3 {
			depMap[dependency[1]] = dependency[2]
		}
	}
	return depMap, nil
}

func getGOPATH() (string, error) {
	vgoCmd, err := golangutil.NewCmd()
	if err != nil {
		return "", err
	}
	vgoCmd.Command = []string{"env", "GOPATH"}
	output, err := utils.RunCmdOutput(vgoCmd)
	if err != nil {
		return "", fmt.Errorf("Could not find GOPATH env: %s", err.Error())
	}
	return strings.TrimSpace(string(output)), nil
}
<commit_msg>"jfrog rt go-publish" fix - Skip publishing dependencies which do not have a zip binary.<commit_after>package dependencies

import (
	"bytes"
	"fmt"
	"github.com/jfrog/jfrog-cli-go/jfrog-cli/artifactory/utils"
	golangutil "github.com/jfrog/jfrog-cli-go/jfrog-cli/artifactory/utils/golang"
	"github.com/jfrog/jfrog-cli-go/jfrog-cli/utils/config"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/artifactory/buildinfo"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/artifactory/services/vgo"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/utils/errorutils"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/utils/io/fileutils"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/utils/io/fileutils/checksum"
	"github.com/jfrog/jfrog-cli-go/jfrog-client/utils/log"
	"io/ioutil"
	"path/filepath"
	"regexp"
	"strings"
)

func Load() ([]Dependency, error) {
	goPath, err := getGOPATH()
	if err != nil {
		return nil, errorutils.CheckError(err)
	}
	cachePath := filepath.Join(goPath, "src", "v", "cache")
	return getDependencies(cachePath)
}

// Represent vgo dependency project.
// Includes publishing capabilities and build info dependencies.
type Dependency struct {
	buildInfoDependencies []buildinfo.Dependency
	id                    string
	modContent            []byte
	zipPath               string
	version               string
}

func (dependency *Dependency) GetId() string {
	return dependency.id
}

func (dependency *Dependency) Publish(targetRepo string, details *config.ArtifactoryDetails) error {
	log.Info("Publishing:", dependency.id, "to", targetRepo)
	servicesManager, err := utils.CreateServiceManager(details, false)
	if err != nil {
		return err
	}
	params := &vgo.VgoParamsImpl{}
	params.ZipPath = dependency.zipPath
	params.ModContent = dependency.modContent
	params.Version = dependency.version
	params.TargetRepo = targetRepo

	return servicesManager.PublishVgoProject(params)
}

func (dependency *Dependency) Dependencies() []buildinfo.Dependency {
	return dependency.buildInfoDependencies
}

func getDependencies(cachePath string) ([]Dependency, error) {
	vgoCmd, err := golangutil.NewCmd()
	if err != nil {
		return nil, err
	}
	vgoCmd.Command = []string{"list"}
	vgoCmd.CommandFlags = []string{"-m"}
	output, err := utils.RunCmdOutput(vgoCmd)
	if err != nil {
		return nil, errorutils.CheckError(err)
	}

	nameVersionMap, err := parseListOutput(output)
	if err != nil {
		return nil, nil
	}

	deps := []Dependency{}
	for name, ver := range nameVersionMap {
		dep, err := createDependency(cachePath, name, ver)
		if err != nil {
			return nil, err
		}
		if dep != nil {
			deps = append(deps, *dep)
		}
	}
	return deps, nil
}

// Creates a vgo dependency.
// Returns a nil value in case the dependency does not include a zip in the cache.
func createDependency(cachePath, dependencyName, version string) (*Dependency, error) {
	// We first check if the this dependency has a zip binary in the local vgo cache.
	// If it does not, nil is returned. This seems to be a bug in vgo.
	zipPath := filepath.Join(cachePath, dependencyName, "@v", version+".zip")
	fileExists, err := fileutils.IsFileExists(zipPath)
	if err != nil {
		log.Warn(fmt.Sprintf("Could not find zip binary for dependency '%s' at %s.", dependencyName, zipPath))
		return nil, err
	}
	// Zip binary does not exist, so we skip it by returning a nil dependency.
	if !fileExists {
		return nil, nil
	}

	dep := Dependency{}

	dep.id = strings.Join([]string{dependencyName, version}, ":")
	dep.version = version
	dep.zipPath = zipPath
	dep.modContent, err = ioutil.ReadFile(filepath.Join(cachePath, dependencyName, "@v", version+".mod"))
	if err != nil {
		return &dep, errorutils.CheckError(err)
	}

	// Mod file dependency
	modDependency := buildinfo.Dependency{Id: dep.id}
	checksums, err := checksum.Calc(bytes.NewBuffer(dep.modContent))
	if err != nil {
		return &dep, err
	}
	modDependency.Checksum = &buildinfo.Checksum{Sha1: checksums[checksum.SHA1], Md5: checksums[checksum.MD5]}

	// Zip file dependency
	zipDependency := buildinfo.Dependency{Id: dep.id}
	fileDetails, err := fileutils.GetFileDetails(dep.zipPath)
	if err != nil {
		return &dep, err
	}
	zipDependency.Checksum = &buildinfo.Checksum{Sha1: fileDetails.Checksum.Sha1, Md5: fileDetails.Checksum.Md5}

	dep.buildInfoDependencies = append(dep.buildInfoDependencies, modDependency, zipDependency)
	return &dep, nil
}

func parseListOutput(content []byte) (map[string]string, error) {
	depRegexp, err := regexp.Compile("(\\S+)\\s+(\\S+)")
	if err != nil {
		return nil, errorutils.CheckError(err)
	}

	depMap := map[string]string{}
	lines := bytes.Split(content, []byte("\n"))
	for i := 2; i < len(lines); i++ {
		dependency := depRegexp.FindStringSubmatch(string(lines[i]))
		if len(dependency) == 3 {
			depMap[dependency[1]] = dependency[2]
		}
	}
	return depMap, nil
}

func getGOPATH() (string, error) {
	vgoCmd, err := golangutil.NewCmd()
	if err != nil {
		return "", err
	}
	vgoCmd.Command = []string{"env", "GOPATH"}
	output, err := utils.RunCmdOutput(vgoCmd)
	if err != nil {
		return "", fmt.Errorf("Could not find GOPATH env: %s", err.Error())
	}
	return strings.TrimSpace(string(output)), nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"net"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"koding/kite-handler/command"
	"koding/kite-handler/fs"
	"koding/kite-handler/terminal"
	"koding/kites/klient/collaboration"
	"koding/kites/klient/protocol"
	"koding/kites/klient/usage"

	"github.com/koding/kite"
	"github.com/koding/kite/config"
)

var (
	flagIP          = flag.String("ip", "", "Change public ip")
	flagPort        = flag.Int("port", 56789, "Change running port")
	flagVersion     = flag.Bool("version", false, "Show version and exit")
	flagEnvironment = flag.String("env", protocol.Environment, "Change environment")
	flagRegion      = flag.String("region", protocol.Region, "Change region")
	flagRegisterURL = flag.String("register-url", "", "Change register URL to kontrol")
	flagDebug       = flag.Bool("debug", false, "Debug mode")

	// update paramters
	flagUpdateInterval = flag.Duration("update-interval", time.Minute*5,
		"Change interval for checking for new updates")
	flagUpdateURL = flag.String("update-url",
		"https://s3.amazonaws.com/koding-klient/"+protocol.Environment+"/latest-version.txt",
		"Change update endpoint for latest version")

	VERSION = protocol.Version
	NAME    = protocol.Name

	// this is our main reference to count and measure metrics for the klient
	// we count only those methods, please add/remove methods here that will
	// reset the timer of a klient.
	usg = usage.NewUsage(map[string]bool{
		"fs.readDirectory":    true,
		"fs.glob":             true,
		"fs.readFile":         true,
		"fs.writeFile":        true,
		"fs.uniquePath":       true,
		"fs.getInfo":          true,
		"fs.setPermissions":   true,
		"fs.remove":           true,
		"fs.rename":           true,
		"fs.createDirectory":  true,
		"fs.move":             true,
		"fs.copy":             true,
		"webterm.getSessions": true,
		"webterm.connect":     true,
		"webterm.killSession": true,
		"exec":                true,
		"klient.share":        true,
		"klient.unshare":      true,
		"klient.shared":       true,
	})

	// this is used to allow other users to call any klient method.
	collab = collaboration.New()

	// we also could use an atomic boolean this is simple for now.
	updating   = false
	updatingMu sync.Mutex // protects updating
)

func main() {
	flag.Parse()
	if *flagVersion {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	k := newKite()

	// Close the klient.db in any case. Corrupt db would be catastrophic
	defer collab.Close()

	k.Log.Info("Running as version %s", VERSION)
	k.Run()
}

func newKite() *kite.Kite {
	k := kite.New(NAME, VERSION)

	if *flagDebug {
		k.SetLogLevel(kite.DEBUG)
	}

	conf := config.MustGet()
	k.Config = conf
	k.Config.Port = *flagPort
	k.Config.Environment = *flagEnvironment
	k.Config.Region = *flagRegion
	k.Id = conf.Id // always boot up with the same id in the kite.key

	// FIXME: It's ugly I know. It's a fix for Koding local development and is
	// needed
	if !strings.Contains(k.Config.KontrolURL, "ngrok") {
		// override current kontrolURL so it talks to port 3000, this is needed
		// because ELB can forward requests based on ports. The port 80 and 443 are
		// HTTP/HTTPS only so our kite can't connect it (we use websocket). However
		// We have a TCP proxy at 3000 which allows us to connect via WebSocket.
		u, _ := url.Parse(k.Config.KontrolURL)

		host := u.Host
		if HasPort(u.Host) {
			host, _, _ = net.SplitHostPort(u.Host)
		}

		u.Host = AddPort(host, "3000")
		u.Scheme = "http"
		k.Config.KontrolURL = u.String()
	}

	if *flagUpdateInterval < time.Minute {
		k.Log.Warning("Update interval can't be less than one minute. Setting to one minute.")
		*flagUpdateInterval = time.Minute
	}

	updater := &Updater{
		Endpoint: *flagUpdateURL,
		Interval: *flagUpdateInterval,
		Log:      k.Log,
	}

	// before we register check for latest update and re-update itself before
	// we continue
	k.Log.Info("Checking for new updates")
	if err := updater.checkAndUpdate(); err != nil {
		k.Log.Warning("Self-update: %s", err)
	}

	go updater.Run()

	userIn := func(user string, users ...string) bool {
		for _, u := range users {
			if u == user {
				return true
			}
		}
		return false
	}

	// don't pass any request if the caller is outside of our scope.
	// don't allow anyone to call a method if we are during an update.
	k.PreHandleFunc(func(r *kite.Request) (interface{}, error) {
		// only authenticated methods have correct username. For example
		// kite.ping has authentication disabled so username can be empty.
		if r.Auth != nil {
			k.Log.Info("Kite '%s/%s/%s' called method: '%s'",
				r.Username, r.Client.Environment, r.Client.Name, r.Method)

			// Allow these users by default
			allowedUsers := []string{k.Config.Username, "koding"}

			// Allow collaboration users as well
			sharedUsers, err := collab.GetAll()
			if err != nil {
				return nil, fmt.Errorf("Can't read shared users from the storage. Err: %v", err)
			}
			allowedUsers = append(allowedUsers, sharedUsers...)

			if !userIn(r.Username, allowedUsers...) {
				return nil, fmt.Errorf("User '%s' is not allowed to make a call to us.", r.Username)
			}
		}

		updatingMu.Lock()
		defer updatingMu.Unlock()

		if updating {
			return nil, errors.New("Updating klient. Can't accept any method.")
		}

		return true, nil
	})

	// Unshare collab users if the klient owner disconnects
	k.OnDisconnect(func(c *kite.Client) {
		k.Log.Info("Kite '%s/%s/%s' is disconnected", c.Username, c.Environment, c.Name)
		if c.Username == k.Config.Username {
			sharedUsers, err := collab.GetAll()
			if err != nil {
				k.Log.Warning("Couldn't unshare users: '%s'", err)
				return
			}

			k.Log.Info("Unsharing users '%s'", sharedUsers)
			for _, user := range sharedUsers {
				if err := collab.Delete(user); err != nil {
					k.Log.Warning("Couldn't delete user from storage: '%s'", err)
				}
			}
		}
	})

	// Metrics, is used by Kloud to get usage so Kloud can stop free VMs
	k.PreHandleFunc(usg.Counter) // we measure every incoming request
	k.HandleFunc("klient.usage", usg.Current)

	// Collaboration, is used by our Koding.com browser client
	k.HandleFunc("klient.share", collab.Share)
	k.HandleFunc("klient.unshare", collab.Unshare)
	k.HandleFunc("klient.shared", collab.Shared)

	// Filesystem
	k.HandleFunc("fs.readDirectory", fs.ReadDirectory)
	k.HandleFunc("fs.glob", fs.Glob)
	k.HandleFunc("fs.readFile", fs.ReadFile)
	k.HandleFunc("fs.writeFile", fs.WriteFile)
	k.HandleFunc("fs.uniquePath", fs.UniquePath)
	k.HandleFunc("fs.getInfo", fs.GetInfo)
	k.HandleFunc("fs.setPermissions", fs.SetPermissions)
	k.HandleFunc("fs.remove", fs.Remove)
	k.HandleFunc("fs.rename", fs.Rename)
	k.HandleFunc("fs.createDirectory", fs.CreateDirectory)
	k.HandleFunc("fs.move", fs.Move)
	k.HandleFunc("fs.copy", fs.Copy)

	// Terminal
	terminal.ResetFunc = usg.Reset
	k.HandleFunc("webterm.getSessions", terminal.GetSessions)
	k.HandleFunc("webterm.connect", terminal.Connect)
	k.HandleFunc("webterm.killSession", terminal.KillSession)

	// Execution
	k.HandleFunc("exec", command.Exec)

	if err := register(k); err != nil {
		panic(err)
	}

	return k
}

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func HasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func AddPort(host, port string) string {
	if ok := HasPort(host); ok {
		return host
	}

	return host + ":" + port
}
<commit_msg>kite-handler/terminal: make Terminal Handler a struct to modify certain things<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"net"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"koding/kite-handler/command"
	"koding/kite-handler/fs"
	"koding/kite-handler/terminal"
	"koding/kites/klient/collaboration"
	"koding/kites/klient/protocol"
	"koding/kites/klient/usage"

	"github.com/koding/kite"
	"github.com/koding/kite/config"
)

var (
	flagIP          = flag.String("ip", "", "Change public ip")
	flagPort        = flag.Int("port", 56789, "Change running port")
	flagVersion     = flag.Bool("version", false, "Show version and exit")
	flagEnvironment = flag.String("env", protocol.Environment, "Change environment")
	flagRegion      = flag.String("region", protocol.Region, "Change region")
	flagRegisterURL = flag.String("register-url", "", "Change register URL to kontrol")
	flagDebug       = flag.Bool("debug", false, "Debug mode")

	// update paramters
	flagUpdateInterval = flag.Duration("update-interval", time.Minute*5,
		"Change interval for checking for new updates")
	flagUpdateURL = flag.String("update-url",
		"https://s3.amazonaws.com/koding-klient/"+protocol.Environment+"/latest-version.txt",
		"Change update endpoint for latest version")

	VERSION = protocol.Version
	NAME    = protocol.Name

	// this is our main reference to count and measure metrics for the klient
	// we count only those methods, please add/remove methods here that will
	// reset the timer of a klient.
	usg = usage.NewUsage(map[string]bool{
		"fs.readDirectory":    true,
		"fs.glob":             true,
		"fs.readFile":         true,
		"fs.writeFile":        true,
		"fs.uniquePath":       true,
		"fs.getInfo":          true,
		"fs.setPermissions":   true,
		"fs.remove":           true,
		"fs.rename":           true,
		"fs.createDirectory":  true,
		"fs.move":             true,
		"fs.copy":             true,
		"webterm.getSessions": true,
		"webterm.connect":     true,
		"webterm.killSession": true,
		"exec":                true,
		"klient.share":        true,
		"klient.unshare":      true,
		"klient.shared":       true,
	})

	// this is used to allow other users to call any klient method.
	collab = collaboration.New()

	// we also could use an atomic boolean this is simple for now.
	updating   = false
	updatingMu sync.Mutex // protects updating
)

func main() {
	flag.Parse()
	if *flagVersion {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	k := newKite()

	// Close the klient.db in any case. Corrupt db would be catastrophic
	defer collab.Close()

	k.Log.Info("Running as version %s", VERSION)
	k.Run()
}

func newKite() *kite.Kite {
	k := kite.New(NAME, VERSION)

	if *flagDebug {
		k.SetLogLevel(kite.DEBUG)
	}

	conf := config.MustGet()
	k.Config = conf
	k.Config.Port = *flagPort
	k.Config.Environment = *flagEnvironment
	k.Config.Region = *flagRegion
	k.Id = conf.Id // always boot up with the same id in the kite.key

	// FIXME: It's ugly I know. It's a fix for Koding local development and is
	// needed
	if !strings.Contains(k.Config.KontrolURL, "ngrok") {
		// override current kontrolURL so it talks to port 3000, this is needed
		// because ELB can forward requests based on ports. The port 80 and 443 are
		// HTTP/HTTPS only so our kite can't connect it (we use websocket). However
		// We have a TCP proxy at 3000 which allows us to connect via WebSocket.
		u, _ := url.Parse(k.Config.KontrolURL)

		host := u.Host
		if HasPort(u.Host) {
			host, _, _ = net.SplitHostPort(u.Host)
		}

		u.Host = AddPort(host, "3000")
		u.Scheme = "http"
		k.Config.KontrolURL = u.String()
	}

	if *flagUpdateInterval < time.Minute {
		k.Log.Warning("Update interval can't be less than one minute. Setting to one minute.")
		*flagUpdateInterval = time.Minute
	}

	updater := &Updater{
		Endpoint: *flagUpdateURL,
		Interval: *flagUpdateInterval,
		Log:      k.Log,
	}

	// before we register check for latest update and re-update itself before
	// we continue
	k.Log.Info("Checking for new updates")
	if err := updater.checkAndUpdate(); err != nil {
		k.Log.Warning("Self-update: %s", err)
	}

	go updater.Run()

	userIn := func(user string, users ...string) bool {
		for _, u := range users {
			if u == user {
				return true
			}
		}
		return false
	}

	// don't pass any request if the caller is outside of our scope.
	// don't allow anyone to call a method if we are during an update.
	k.PreHandleFunc(func(r *kite.Request) (interface{}, error) {
		// only authenticated methods have correct username. For example
		// kite.ping has authentication disabled so username can be empty.
		if r.Auth != nil {
			k.Log.Info("Kite '%s/%s/%s' called method: '%s'",
				r.Username, r.Client.Environment, r.Client.Name, r.Method)

			// Allow these users by default
			allowedUsers := []string{k.Config.Username, "koding"}

			// Allow collaboration users as well
			sharedUsers, err := collab.GetAll()
			if err != nil {
				return nil, fmt.Errorf("Can't read shared users from the storage. Err: %v", err)
			}
			allowedUsers = append(allowedUsers, sharedUsers...)

			if !userIn(r.Username, allowedUsers...) {
				return nil, fmt.Errorf("User '%s' is not allowed to make a call to us.", r.Username)
			}
		}

		updatingMu.Lock()
		defer updatingMu.Unlock()

		if updating {
			return nil, errors.New("Updating klient. Can't accept any method.")
		}

		return true, nil
	})

	// Unshare collab users if the klient owner disconnects
	k.OnDisconnect(func(c *kite.Client) {
		k.Log.Info("Kite '%s/%s/%s' is disconnected", c.Username, c.Environment, c.Name)
		if c.Username == k.Config.Username {
			sharedUsers, err := collab.GetAll()
			if err != nil {
				k.Log.Warning("Couldn't unshare users: '%s'", err)
				return
			}

			k.Log.Info("Unsharing users '%s'", sharedUsers)
			for _, user := range sharedUsers {
				if err := collab.Delete(user); err != nil {
					k.Log.Warning("Couldn't delete user from storage: '%s'", err)
				}
			}
		}
	})

	// Metrics, is used by Kloud to get usage so Kloud can stop free VMs
	k.PreHandleFunc(usg.Counter) // we measure every incoming request
	k.HandleFunc("klient.usage", usg.Current)

	// Collaboration, is used by our Koding.com browser client
	k.HandleFunc("klient.share", collab.Share)
	k.HandleFunc("klient.unshare", collab.Unshare)
	k.HandleFunc("klient.shared", collab.Shared)

	// Filesystem
	k.HandleFunc("fs.readDirectory", fs.ReadDirectory)
	k.HandleFunc("fs.glob", fs.Glob)
	k.HandleFunc("fs.readFile", fs.ReadFile)
	k.HandleFunc("fs.writeFile", fs.WriteFile)
	k.HandleFunc("fs.uniquePath", fs.UniquePath)
	k.HandleFunc("fs.getInfo", fs.GetInfo)
	k.HandleFunc("fs.setPermissions", fs.SetPermissions)
	k.HandleFunc("fs.remove", fs.Remove)
	k.HandleFunc("fs.rename", fs.Rename)
	k.HandleFunc("fs.createDirectory", fs.CreateDirectory)
	k.HandleFunc("fs.move", fs.Move)
	k.HandleFunc("fs.copy", fs.Copy)

	// Terminal
	term := &terminal.Terminal{
		InputHook: usg.Reset,
	}

	k.HandleFunc("webterm.getSessions", term.GetSessions)
	k.HandleFunc("webterm.connect", term.Connect)
	k.HandleFunc("webterm.killSession", term.KillSession)

	// Execution
	k.HandleFunc("exec", command.Exec)

	if err := register(k); err != nil {
		panic(err)
	}

	return k
}

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func HasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func AddPort(host, port string) string {
	if ok := HasPort(host); ok {
		return host
	}

	return host + ":" + port
}
<|endoftext|>