<commit_before>package php

import "testing"

var testFile = `<?php

$str = "I'm a string!";
$num = -230.098;

session_start();

function foo(barType $bar, $foobar) {
  if (buzz()) {
    fizzbuzz();
  }
}

class MyClass {
  public myMethod() {
  }
}

$var = $this->myMethod();

?>

<? echo something(); ?>
`

func assertNext(t *testing.T, l *lexer, typ ItemType) Item {
	i := l.nextItem()
	if i.typ != typ {
		t.Fatal("Incorrect lexing. Expected:", typ, "Found:", i)
	}
	return i
}

func assertItem(t *testing.T, i Item, expected string) {
	if i.val != expected {
		t.Fatal("Did not correctly parse item", i)
	}
}

func TestPHPLexer(t *testing.T) {
	l := newLexer(testFile)

	var i Item
	i = assertNext(t, l, itemPHPBegin)

	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemAssignmentOperator)
	i = assertNext(t, l, itemStringLiteral)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemAssignmentOperator)
	i = assertNext(t, l, itemSubtractionOperator)
	i = assertNext(t, l, itemNumberLiteral)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemFunction)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemComma)
	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemBlockBegin)

	i = assertNext(t, l, itemIf)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemCloseParen)

	i = assertNext(t, l, itemBlockBegin)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemStatementEnd)
	i = assertNext(t, l, itemBlockEnd)

	i = assertNext(t, l, itemBlockEnd)

	i = assertNext(t, l, itemClass)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemBlockBegin)
	i = assertNext(t, l, itemPublic)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemBlockBegin)
	i = assertNext(t, l, itemBlockEnd)
	i = assertNext(t, l, itemBlockEnd)

	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemAssignmentOperator)
	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemObjectOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemPHPEnd)
	i = assertNext(t, l, itemHTML)
	assertItem(t, i, "\n\n")

	i = assertNext(t, l, itemPHPBegin)
	i = assertNext(t, l, itemEcho)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemPHPEnd)
	i = assertNext(t, l, itemHTML)
	assertItem(t, i, "\n")

	i = assertNext(t, l, itemEOF)
}
<commit_msg>Tweaking lexer test to not fail when variable is named<commit_after>package php

import "testing"

var testFile = `<?php

$str = "I'm a string!";
$num = -230.098;

session_start();

function foo(barType $bar, $foobar) {
  if (buzz()) {
    fizzbuzz();
  }
}

class MyClass {
  public myMethod() {
  }
}

$myvar = $this->myMethod();

?>

<? echo something(); ?>
`

func assertNext(t *testing.T, l *lexer, typ ItemType) Item {
	i := l.nextItem()
	if i.typ != typ {
		t.Fatal("Incorrect lexing. Expected:", typ, "Found:", i)
	}
	return i
}

func assertItem(t *testing.T, i Item, expected string) {
	if i.val != expected {
		t.Fatal("Did not correctly parse item", i)
	}
}

func TestPHPLexer(t *testing.T) {
	l := newLexer(testFile)

	var i Item
	i = assertNext(t, l, itemPHPBegin)

	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemAssignmentOperator)
	i = assertNext(t, l, itemStringLiteral)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemAssignmentOperator)
	i = assertNext(t, l, itemSubtractionOperator)
	i = assertNext(t, l, itemNumberLiteral)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemFunction)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemComma)
	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemBlockBegin)

	i = assertNext(t, l, itemIf)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemCloseParen)

	i = assertNext(t, l, itemBlockBegin)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemStatementEnd)
	i = assertNext(t, l, itemBlockEnd)

	i = assertNext(t, l, itemBlockEnd)

	i = assertNext(t, l, itemClass)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemBlockBegin)
	i = assertNext(t, l, itemPublic)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemBlockBegin)
	i = assertNext(t, l, itemBlockEnd)
	i = assertNext(t, l, itemBlockEnd)

	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemAssignmentOperator)
	i = assertNext(t, l, itemVariableOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemObjectOperator)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemPHPEnd)
	i = assertNext(t, l, itemHTML)
	assertItem(t, i, "\n\n")

	i = assertNext(t, l, itemPHPBegin)
	i = assertNext(t, l, itemEcho)
	i = assertNext(t, l, itemIdentifier)
	i = assertNext(t, l, itemOpenParen)
	i = assertNext(t, l, itemCloseParen)
	i = assertNext(t, l, itemStatementEnd)

	i = assertNext(t, l, itemPHPEnd)
	i = assertNext(t, l, itemHTML)
	assertItem(t, i, "\n")

	i = assertNext(t, l, itemEOF)
}
<|endoftext|>


<commit_before>package lfs

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/cheggaaa/pb"
	"github.com/rubyist/tracerx"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
)

const (
	mediaType = "application/vnd.git-lfs+json; charset-utf-8"
)

var (
	lfsMediaTypeRE             = regexp.MustCompile(`\Aapplication/vnd\.git\-lfs\+json(;|\z)`)
	jsonMediaTypeRE            = regexp.MustCompile(`\Aapplication/json(;|\z)`)
	objectRelationDoesNotExist = errors.New("relation does not exist")
	hiddenHeaders              = map[string]bool{
		"Authorization": true,
	}

	// 401 and 403 print the same default error message
	defaultErrors = map[int]string{
		400: "Client error: %s",
		401: "Authorization error: %s\nCheck that you have proper access to the repository",
		404: "Repository or object not found: %s\nCheck that it exists and that you have proper access to it",
		500: "Server error: %s",
	}
)

type objectResource struct {
	Oid   string                   `json:"oid,omitempty"`
	Size  int64                    `json:"size,omitempty"`
	Links map[string]*linkRelation `json:"_links,omitempty"`
}

func (o *objectResource) NewRequest(relation, method string) (*http.Request, Creds, error) {
	rel, ok := o.Rel(relation)
	if !ok {
		return nil, nil, objectRelationDoesNotExist
	}

	req, creds, err := newClientRequest(method, rel.Href)
	if err != nil {
		return nil, nil, err
	}

	for h, v := range rel.Header {
		req.Header.Set(h, v)
	}

	return req, creds, nil
}

func (o *objectResource) Rel(name string) (*linkRelation, bool) {
	if o.Links == nil {
		return nil, false
	}

	rel, ok := o.Links[name]
	return rel, ok
}

type linkRelation struct {
	Href   string            `json:"href"`
	Header map[string]string `json:"header,omitempty"`
}

type ClientError struct {
	Message          string `json:"message"`
	DocumentationUrl string `json:"documentation_url,omitempty"`
	RequestId        string `json:"request_id,omitempty"`
}

func (e *ClientError) Error() string {
	msg := e.Message
	if len(e.DocumentationUrl) > 0 {
		msg += "\nDocs: " + e.DocumentationUrl
	}
	if len(e.RequestId) > 0 {
		msg += "\nRequest ID: " + e.RequestId
	}
	return msg
}

func Download(oid string) (io.ReadCloser, int64, *WrappedError) {
	req, creds, err := newApiRequest("GET", oid)
	if err != nil {
		return nil, 0, Error(err)
	}

	res, obj, wErr := doApiRequest(req, creds)
	if wErr != nil {
		return nil, 0, wErr
	}

	req, creds, err = obj.NewRequest("download", "GET")
	if err != nil {
		return nil, 0, Error(err)
	}

	res, wErr = doHttpRequest(req, creds)
	if wErr != nil {
		return nil, 0, wErr
	}

	return res.Body, res.ContentLength, nil
}

type byteCloser struct {
	*bytes.Reader
}

func (b *byteCloser) Close() error {
	return nil
}

func Upload(oidPath, filename string, cb CopyCallback) *WrappedError {
	oid := filepath.Base(oidPath)
	file, err := os.Open(oidPath)
	if err != nil {
		return Error(err)
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return Error(err)
	}

	reqObj := &objectResource{
		Oid:  oid,
		Size: stat.Size(),
	}

	by, err := json.Marshal(reqObj)
	if err != nil {
		return Error(err)
	}

	req, creds, err := newApiRequest("POST", "")
	if err != nil {
		return Error(err)
	}

	req.Header.Set("Content-Type", mediaType)
	req.Header.Set("Content-Length", strconv.Itoa(len(by)))
	req.ContentLength = int64(len(by))
	req.Body = &byteCloser{bytes.NewReader(by)}

	tracerx.Printf("api: uploading %s (%s)", filename, oid)
	res, obj, wErr := doApiRequest(req, creds)
	if wErr != nil {
		return wErr
	}

	if res.StatusCode == 200 {
		return nil
	}

	req, creds, err = obj.NewRequest("upload", "PUT")
	if err != nil {
		return Error(err)
	}

	if len(req.Header.Get("Content-Type")) == 0 {
		req.Header.Set("Content-Type", "application/octet-stream")
	}
	req.Header.Set("Content-Length", strconv.FormatInt(reqObj.Size, 10))
	req.ContentLength = reqObj.Size

	reader := &CallbackReader{
		C:         cb,
		TotalSize: reqObj.Size,
		Reader:    file,
	}

	bar := pb.New64(reqObj.Size)
	bar.SetUnits(pb.U_BYTES)
	bar.Start()

	req.Body = ioutil.NopCloser(bar.NewProxyReader(reader))

	res, wErr = doHttpRequest(req, creds)
	if wErr != nil {
		return wErr
	}

	if res.StatusCode > 299 {
		return Errorf(nil, "Invalid status for %s %s: %d", req.Method, req.URL, res.StatusCode)
	}

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	req, creds, err = obj.NewRequest("verify", "POST")
	if err == objectRelationDoesNotExist {
		return nil
	} else if err != nil {
		return Error(err)
	}

	req.Header.Set("Content-Type", mediaType)
	req.Header.Set("Content-Length", strconv.Itoa(len(by)))
	req.ContentLength = int64(len(by))
	req.Body = ioutil.NopCloser(bytes.NewReader(by))
	res, wErr = doHttpRequest(req, creds)

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	return wErr
}

func doHttpRequest(req *http.Request, creds Creds) (*http.Response, *WrappedError) {
	res, err := DoHTTP(Config, req)

	var wErr *WrappedError

	if err != nil {
		wErr = Errorf(err, "Error for %s %s", res.Request.Method, res.Request.URL)
	} else {
		if creds != nil {
			saveCredentials(creds, res)
		}

		wErr = handleResponse(res)
	}

	if wErr != nil {
		if res != nil {
			setErrorResponseContext(wErr, res)
		} else {
			setErrorRequestContext(wErr, req)
		}
	}

	return res, wErr
}

func doApiRequestWithRedirects(req *http.Request, creds Creds, via []*http.Request) (*http.Response, *WrappedError) {
	res, wErr := doHttpRequest(req, creds)
	if wErr != nil {
		return res, wErr
	}

	if res.StatusCode == 307 {
		redirectedReq, redirectedCreds, err := newClientRequest(req.Method, res.Header.Get("Location"))
		if err != nil {
			return res, Errorf(err, err.Error())
		}

		via = append(via, req)
		if seeker, ok := req.Body.(io.Seeker); ok {
			_, err := seeker.Seek(0, 0)
			if err != nil {
				return res, Error(err)
			}
			redirectedReq.Body = req.Body
			redirectedReq.ContentLength = req.ContentLength
		} else {
			return res, Errorf(nil, "Request body needs to be an io.Seeker to handle redirects.")
		}

		if err = checkRedirect(redirectedReq, via); err != nil {
			return res, Errorf(err, err.Error())
		}

		return doApiRequestWithRedirects(redirectedReq, redirectedCreds, via)
	}

	return res, wErr
}

func doApiRequest(req *http.Request, creds Creds) (*http.Response, *objectResource, *WrappedError) {
	via := make([]*http.Request, 0, 4)
	res, wErr := doApiRequestWithRedirects(req, creds, via)
	if wErr != nil {
		return res, nil, wErr
	}

	obj := &objectResource{}
	wErr = decodeApiResponse(res, obj)

	if wErr != nil {
		setErrorResponseContext(wErr, res)
	}

	return res, obj, wErr
}

func handleResponse(res *http.Response) *WrappedError {
	if res.StatusCode < 400 {
		return nil
	}

	defer func() {
		io.Copy(ioutil.Discard, res.Body)
		res.Body.Close()
	}()

	cliErr := &ClientError{}
	wErr := decodeApiResponse(res, cliErr)
	if wErr == nil {
		if len(cliErr.Message) == 0 {
			wErr = defaultError(res)
		} else {
			wErr = Error(cliErr)
		}
	}

	wErr.Panic = res.StatusCode > 499 && res.StatusCode != 501 && res.StatusCode != 509
	return wErr
}

func decodeApiResponse(res *http.Response, obj interface{}) *WrappedError {
	ctype := res.Header.Get("Content-Type")
	if !(lfsMediaTypeRE.MatchString(ctype) || jsonMediaTypeRE.MatchString(ctype)) {
		return nil
	}

	err := json.NewDecoder(res.Body).Decode(obj)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	if err != nil {
		return Errorf(err, "Unable to parse HTTP response for %s %s", res.Request.Method, res.Request.URL)
	}

	return nil
}

func defaultError(res *http.Response) *WrappedError {
	var msgFmt string

	if f, ok := defaultErrors[res.StatusCode]; ok {
		msgFmt = f
	} else if res.StatusCode < 500 {
		msgFmt = defaultErrors[400] + fmt.Sprintf(" from HTTP %d", res.StatusCode)
	} else {
		msgFmt = defaultErrors[500] + fmt.Sprintf(" from HTTP %d", res.StatusCode)
	}

	return Error(fmt.Errorf(msgFmt, res.Request.URL))
}

func saveCredentials(creds Creds, res *http.Response) {
	if creds == nil {
		return
	}

	if res.StatusCode < 300 {
		execCreds(creds, "approve")
	} else if res.StatusCode == 401 {
		execCreds(creds, "reject")
	}
}

func newApiRequest(method, oid string) (*http.Request, Creds, error) {
	u, err := Config.ObjectUrl(oid)
	if err != nil {
		return nil, nil, err
	}

	req, creds, err := newClientRequest(method, u.String())
	if err == nil {
		req.Header.Set("Accept", mediaType)
	}
	return req, creds, err
}

func newClientRequest(method, rawurl string) (*http.Request, Creds, error) {
	req, err := http.NewRequest(method, rawurl, nil)
	if err != nil {
		return req, nil, err
	}

	req.Header.Set("User-Agent", UserAgent)
	creds, err := getCreds(req)
	return req, creds, err
}

func getCreds(req *http.Request) (Creds, error) {
	if len(req.Header.Get("Authorization")) > 0 {
		return nil, nil
	}

	apiUrl, err := Config.ObjectUrl("")
	if err != nil {
		return nil, err
	}

	if req.URL.Scheme == apiUrl.Scheme &&
		req.URL.Host == apiUrl.Host {
		creds, err := credentials(req.URL)
		if err != nil {
			return nil, err
		}

		token := fmt.Sprintf("%s:%s", creds["username"], creds["password"])
		auth := "Basic " + base64.URLEncoding.EncodeToString([]byte(token))
		req.Header.Set("Authorization", auth)
		return creds, nil
	}

	return nil, nil
}

func setErrorRequestContext(err *WrappedError, req *http.Request) {
	err.Set("Endpoint", Config.Endpoint())
	err.Set("URL", fmt.Sprintf("%s %s", req.Method, req.URL.String()))
	setErrorHeaderContext(err, "Response", req.Header)
}

func setErrorResponseContext(err *WrappedError, res *http.Response) {
	err.Set("Status", res.Status)
	setErrorHeaderContext(err, "Request", res.Header)
	setErrorRequestContext(err, res.Request)
}

func setErrorHeaderContext(err *WrappedError, prefix string, head http.Header) {
	for key, _ := range head {
		contextKey := fmt.Sprintf("%s:%s", prefix, key)
		if _, skip := hiddenHeaders[key]; skip {
			err.Set(contextKey, "--")
		} else {
			err.Set(contextKey, head.Get(key))
		}
	}
}

func init() {
	defaultErrors[403] = defaultErrors[401]
}
<commit_msg>just dup the line<commit_after>package lfs

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/cheggaaa/pb"
	"github.com/rubyist/tracerx"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
)

const (
	mediaType = "application/vnd.git-lfs+json; charset-utf-8"
)

var (
	lfsMediaTypeRE             = regexp.MustCompile(`\Aapplication/vnd\.git\-lfs\+json(;|\z)`)
	jsonMediaTypeRE            = regexp.MustCompile(`\Aapplication/json(;|\z)`)
	objectRelationDoesNotExist = errors.New("relation does not exist")
	hiddenHeaders              = map[string]bool{
		"Authorization": true,
	}

	defaultErrors = map[int]string{
		400: "Client error: %s",
		401: "Authorization error: %s\nCheck that you have proper access to the repository",
		403: "Authorization error: %s\nCheck that you have proper access to the repository",
		404: "Repository or object not found: %s\nCheck that it exists and that you have proper access to it",
		500: "Server error: %s",
	}
)

type objectResource struct {
	Oid   string                   `json:"oid,omitempty"`
	Size  int64                    `json:"size,omitempty"`
	Links map[string]*linkRelation `json:"_links,omitempty"`
}

func (o *objectResource) NewRequest(relation, method string) (*http.Request, Creds, error) {
	rel, ok := o.Rel(relation)
	if !ok {
		return nil, nil, objectRelationDoesNotExist
	}

	req, creds, err := newClientRequest(method, rel.Href)
	if err != nil {
		return nil, nil, err
	}

	for h, v := range rel.Header {
		req.Header.Set(h, v)
	}

	return req, creds, nil
}

func (o *objectResource) Rel(name string) (*linkRelation, bool) {
	if o.Links == nil {
		return nil, false
	}

	rel, ok := o.Links[name]
	return rel, ok
}

type linkRelation struct {
	Href   string            `json:"href"`
	Header map[string]string `json:"header,omitempty"`
}

type ClientError struct {
	Message          string `json:"message"`
	DocumentationUrl string `json:"documentation_url,omitempty"`
	RequestId        string `json:"request_id,omitempty"`
}

func (e *ClientError) Error() string {
	msg := e.Message
	if len(e.DocumentationUrl) > 0 {
		msg += "\nDocs: " + e.DocumentationUrl
	}
	if len(e.RequestId) > 0 {
		msg += "\nRequest ID: " + e.RequestId
	}
	return msg
}

func Download(oid string) (io.ReadCloser, int64, *WrappedError) {
	req, creds, err := newApiRequest("GET", oid)
	if err != nil {
		return nil, 0, Error(err)
	}

	res, obj, wErr := doApiRequest(req, creds)
	if wErr != nil {
		return nil, 0, wErr
	}

	req, creds, err = obj.NewRequest("download", "GET")
	if err != nil {
		return nil, 0, Error(err)
	}

	res, wErr = doHttpRequest(req, creds)
	if wErr != nil {
		return nil, 0, wErr
	}

	return res.Body, res.ContentLength, nil
}

type byteCloser struct {
	*bytes.Reader
}

func (b *byteCloser) Close() error {
	return nil
}

func Upload(oidPath, filename string, cb CopyCallback) *WrappedError {
	oid := filepath.Base(oidPath)
	file, err := os.Open(oidPath)
	if err != nil {
		return Error(err)
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return Error(err)
	}

	reqObj := &objectResource{
		Oid:  oid,
		Size: stat.Size(),
	}

	by, err := json.Marshal(reqObj)
	if err != nil {
		return Error(err)
	}

	req, creds, err := newApiRequest("POST", "")
	if err != nil {
		return Error(err)
	}

	req.Header.Set("Content-Type", mediaType)
	req.Header.Set("Content-Length", strconv.Itoa(len(by)))
	req.ContentLength = int64(len(by))
	req.Body = &byteCloser{bytes.NewReader(by)}

	tracerx.Printf("api: uploading %s (%s)", filename, oid)
	res, obj, wErr := doApiRequest(req, creds)
	if wErr != nil {
		return wErr
	}

	if res.StatusCode == 200 {
		return nil
	}

	req, creds, err = obj.NewRequest("upload", "PUT")
	if err != nil {
		return Error(err)
	}

	if len(req.Header.Get("Content-Type")) == 0 {
		req.Header.Set("Content-Type", "application/octet-stream")
	}
	req.Header.Set("Content-Length", strconv.FormatInt(reqObj.Size, 10))
	req.ContentLength = reqObj.Size

	reader := &CallbackReader{
		C:         cb,
		TotalSize: reqObj.Size,
		Reader:    file,
	}

	bar := pb.New64(reqObj.Size)
	bar.SetUnits(pb.U_BYTES)
	bar.Start()

	req.Body = ioutil.NopCloser(bar.NewProxyReader(reader))

	res, wErr = doHttpRequest(req, creds)
	if wErr != nil {
		return wErr
	}

	if res.StatusCode > 299 {
		return Errorf(nil, "Invalid status for %s %s: %d", req.Method, req.URL, res.StatusCode)
	}

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	req, creds, err = obj.NewRequest("verify", "POST")
	if err == objectRelationDoesNotExist {
		return nil
	} else if err != nil {
		return Error(err)
	}

	req.Header.Set("Content-Type", mediaType)
	req.Header.Set("Content-Length", strconv.Itoa(len(by)))
	req.ContentLength = int64(len(by))
	req.Body = ioutil.NopCloser(bytes.NewReader(by))
	res, wErr = doHttpRequest(req, creds)

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	return wErr
}

func doHttpRequest(req *http.Request, creds Creds) (*http.Response, *WrappedError) {
	res, err := DoHTTP(Config, req)

	var wErr *WrappedError

	if err != nil {
		wErr = Errorf(err, "Error for %s %s", res.Request.Method, res.Request.URL)
	} else {
		if creds != nil {
			saveCredentials(creds, res)
		}

		wErr = handleResponse(res)
	}

	if wErr != nil {
		if res != nil {
			setErrorResponseContext(wErr, res)
		} else {
			setErrorRequestContext(wErr, req)
		}
	}

	return res, wErr
}

func doApiRequestWithRedirects(req *http.Request, creds Creds, via []*http.Request) (*http.Response, *WrappedError) {
	res, wErr := doHttpRequest(req, creds)
	if wErr != nil {
		return res, wErr
	}

	if res.StatusCode == 307 {
		redirectedReq, redirectedCreds, err := newClientRequest(req.Method, res.Header.Get("Location"))
		if err != nil {
			return res, Errorf(err, err.Error())
		}

		via = append(via, req)
		if seeker, ok := req.Body.(io.Seeker); ok {
			_, err := seeker.Seek(0, 0)
			if err != nil {
				return res, Error(err)
			}
			redirectedReq.Body = req.Body
			redirectedReq.ContentLength = req.ContentLength
		} else {
			return res, Errorf(nil, "Request body needs to be an io.Seeker to handle redirects.")
		}

		if err = checkRedirect(redirectedReq, via); err != nil {
			return res, Errorf(err, err.Error())
		}

		return doApiRequestWithRedirects(redirectedReq, redirectedCreds, via)
	}

	return res, wErr
}

func doApiRequest(req *http.Request, creds Creds) (*http.Response, *objectResource, *WrappedError) {
	via := make([]*http.Request, 0, 4)
	res, wErr := doApiRequestWithRedirects(req, creds, via)
	if wErr != nil {
		return res, nil, wErr
	}

	obj := &objectResource{}
	wErr = decodeApiResponse(res, obj)

	if wErr != nil {
		setErrorResponseContext(wErr, res)
	}

	return res, obj, wErr
}

func handleResponse(res *http.Response) *WrappedError {
	if res.StatusCode < 400 {
		return nil
	}

	defer func() {
		io.Copy(ioutil.Discard, res.Body)
		res.Body.Close()
	}()

	cliErr := &ClientError{}
	wErr := decodeApiResponse(res, cliErr)
	if wErr == nil {
		if len(cliErr.Message) == 0 {
			wErr = defaultError(res)
		} else {
			wErr = Error(cliErr)
		}
	}

	wErr.Panic = res.StatusCode > 499 && res.StatusCode != 501 && res.StatusCode != 509
	return wErr
}

func decodeApiResponse(res *http.Response, obj interface{}) *WrappedError {
	ctype := res.Header.Get("Content-Type")
	if !(lfsMediaTypeRE.MatchString(ctype) || jsonMediaTypeRE.MatchString(ctype)) {
		return nil
	}

	err := json.NewDecoder(res.Body).Decode(obj)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	if err != nil {
		return Errorf(err, "Unable to parse HTTP response for %s %s", res.Request.Method, res.Request.URL)
	}

	return nil
}

func defaultError(res *http.Response) *WrappedError {
	var msgFmt string

	if f, ok := defaultErrors[res.StatusCode]; ok {
		msgFmt = f
	} else if res.StatusCode < 500 {
		msgFmt = defaultErrors[400] + fmt.Sprintf(" from HTTP %d", res.StatusCode)
	} else {
		msgFmt = defaultErrors[500] + fmt.Sprintf(" from HTTP %d", res.StatusCode)
	}

	return Error(fmt.Errorf(msgFmt, res.Request.URL))
}

func saveCredentials(creds Creds, res *http.Response) {
	if creds == nil {
		return
	}

	if res.StatusCode < 300 {
		execCreds(creds, "approve")
	} else if res.StatusCode == 401 {
		execCreds(creds, "reject")
	}
}

func newApiRequest(method, oid string) (*http.Request, Creds, error) {
	u, err := Config.ObjectUrl(oid)
	if err != nil {
		return nil, nil, err
	}

	req, creds, err := newClientRequest(method, u.String())
	if err == nil {
		req.Header.Set("Accept", mediaType)
	}
	return req, creds, err
}

func newClientRequest(method, rawurl string) (*http.Request, Creds, error) {
	req, err := http.NewRequest(method, rawurl, nil)
	if err != nil {
		return req, nil, err
	}

	req.Header.Set("User-Agent", UserAgent)
	creds, err := getCreds(req)
	return req, creds, err
}

func getCreds(req *http.Request) (Creds, error) {
	if len(req.Header.Get("Authorization")) > 0 {
		return nil, nil
	}

	apiUrl, err := Config.ObjectUrl("")
	if err != nil {
		return nil, err
	}

	if req.URL.Scheme == apiUrl.Scheme &&
		req.URL.Host == apiUrl.Host {
		creds, err := credentials(req.URL)
		if err != nil {
			return nil, err
		}

		token := fmt.Sprintf("%s:%s", creds["username"], creds["password"])
		auth := "Basic " + base64.URLEncoding.EncodeToString([]byte(token))
		req.Header.Set("Authorization", auth)
		return creds, nil
	}

	return nil, nil
}

func setErrorRequestContext(err *WrappedError, req *http.Request) {
	err.Set("Endpoint", Config.Endpoint())
	err.Set("URL", fmt.Sprintf("%s %s", req.Method, req.URL.String()))
	setErrorHeaderContext(err, "Response", req.Header)
}

func setErrorResponseContext(err *WrappedError, res *http.Response) {
	err.Set("Status", res.Status)
	setErrorHeaderContext(err, "Request", res.Header)
	setErrorRequestContext(err, res.Request)
}

func setErrorHeaderContext(err *WrappedError, prefix string, head http.Header) {
	for key, _ := range head {
		contextKey := fmt.Sprintf("%s:%s", prefix, key)
		if _, skip := hiddenHeaders[key]; skip {
			err.Set(contextKey, "--")
		} else {
			err.Set(contextKey, head.Get(key))
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>chore(gopwt/translator): gofmt -s<commit_after><|endoftext|>


<commit_before>package hbase

import (
	"bytes"
	"context"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/tsuna/gohbase"
	"github.com/tsuna/gohbase/hrpc"
	"io"
)

func init() {
	filer.Stores = append(filer.Stores, &HbaseStore{})
}

type HbaseStore struct {
	Client    gohbase.Client
	table     []byte
	cfKv      string
	cfMetaDir string
	column    string
}

func (store *HbaseStore) GetName() string {
	return "hbase"
}

func (store *HbaseStore) Initialize(configuration util.Configuration, prefix string) (err error) {
	return store.initialize(
		configuration.GetString(prefix+"zkquorum"),
		configuration.GetString(prefix+"table"),
	)
}

func (store *HbaseStore) initialize(zkquorum, table string) (err error) {
	store.Client = gohbase.NewClient(zkquorum)
	store.table = []byte(table)
	store.cfKv = "kv"
	store.cfMetaDir = "meta"
	store.column = "a"

	// check table exists
	key := "whatever"
	headers := map[string][]string{store.cfMetaDir: nil}
	get, err := hrpc.NewGet(context.Background(), store.table, []byte(key), hrpc.Families(headers))
	if err != nil {
		return fmt.Errorf("NewGet returned an error: %v", err)
	}
	_, err = store.Client.Get(get)
	if err != gohbase.TableNotFound {
		return nil
	}

	// create table
	adminClient := gohbase.NewAdminClient(zkquorum)
	cFamilies := []string{store.cfKv, store.cfMetaDir}
	cf := make(map[string]map[string]string, len(cFamilies))
	for _, f := range cFamilies {
		cf[f] = nil
	}
	ct := hrpc.NewCreateTable(context.Background(), []byte(table), cf)
	if err := adminClient.CreateTable(ct); err != nil {
		return err
	}

	return nil
}

func (store *HbaseStore) InsertEntry(ctx context.Context, entry *filer.Entry) error {
	value, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encoding %s %+v: %v", entry.FullPath, entry.Attr, err)
	}
	if len(entry.Chunks) > 50 {
		value = util.MaybeGzipData(value)
	}

	return store.doPut(ctx, store.cfMetaDir, []byte(entry.FullPath), value)
}

func (store *HbaseStore) UpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {
	return store.InsertEntry(ctx, entry)
}

func (store *HbaseStore) FindEntry(ctx context.Context, path util.FullPath) (entry *filer.Entry, err error) {
	value, err := store.doGet(ctx, store.cfMetaDir, []byte(path))
	if err != nil {
		if err == filer.ErrKvNotFound {
			return nil, filer_pb.ErrNotFound
		}
		return nil, err
	}

	entry = &filer.Entry{
		FullPath: path,
	}
	err = entry.DecodeAttributesAndChunks(util.MaybeDecompressData(value))
	if err != nil {
		return entry, fmt.Errorf("decode %s : %v", entry.FullPath, err)
	}
	return entry, nil
}

func (store *HbaseStore) DeleteEntry(ctx context.Context, path util.FullPath) (err error) {
	return store.doDelete(ctx, store.cfMetaDir, []byte(path))
}

func (store *HbaseStore) DeleteFolderChildren(ctx context.Context, path util.FullPath) (err error) {

	family := map[string][]string{store.cfMetaDir: {COLUMN_NAME}}
	expectedPrefix := []byte(path + "/")
	scan, err := hrpc.NewScanRange(ctx, store.table, expectedPrefix, nil, hrpc.Families(family))
	if err != nil {
		return err
	}

	scanner := store.Client.Scan(scan)
	defer scanner.Close()
	for {
		res, err := scanner.Next()
		if err != nil {
			break
		}
		if len(res.Cells) == 0 {
			continue
		}
		cell := res.Cells[0]

		if !bytes.HasPrefix(cell.Row, expectedPrefix) {
			break
		}
		fullpath := util.FullPath(cell.Row)
		dir, _ := fullpath.DirAndName()
		if dir != string(dirPath) {
			continue
		}

		err = store.doDelete(ctx, store.cfMetaDir, cell.Row)
		if err != nil {
			break
		}

	}
	return
}

func (store *HbaseStore) ListDirectoryEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int) ([]*filer.Entry, error) {
	return store.ListDirectoryPrefixedEntries(ctx, dirPath, startFileName, includeStartFile, limit, "")
}

func (store *HbaseStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int, prefix string) ([]*filer.Entry, error) {
	family := map[string][]string{store.cfMetaDir: {COLUMN_NAME}}
	expectedPrefix := []byte(string(dirPath) + "/" + prefix)
	scan, err := hrpc.NewScanRange(ctx, store.table, expectedPrefix, nil, hrpc.Families(family))
	if err != nil {
		return nil, err
	}

	var entries []*filer.Entry
	scanner := store.Client.Scan(scan)
	defer scanner.Close()
	for {
		res, err := scanner.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return entries, err
		}
		if len(res.Cells) == 0 {
			continue
		}
		cell := res.Cells[0]

		if !bytes.HasPrefix(cell.Row, expectedPrefix) {
			break
		}

		fullpath := util.FullPath(cell.Row)
		dir, fileName := fullpath.DirAndName()
		if dir != string(dirPath) {
			continue
		}

		value := cell.Value

		if fileName == startFileName && !includeStartFile {
			continue
		}

		limit--
		if limit < 0 {
			break
		}
		entry := &filer.Entry{
			FullPath: fullpath,
		}
		if decodeErr := entry.DecodeAttributesAndChunks(util.MaybeDecompressData(value)); decodeErr != nil {
			err = decodeErr
			glog.V(0).Infof("list %s : %v", entry.FullPath, err)
			break
		}
		entries = append(entries, entry)
	}

	return entries, nil
}

func (store *HbaseStore) BeginTransaction(ctx context.Context) (context.Context, error) {
	return ctx, nil
}

func (store *HbaseStore) CommitTransaction(ctx context.Context) error {
	return nil
}

func (store *HbaseStore) RollbackTransaction(ctx context.Context) error {
	return nil
}

func (store *HbaseStore) Shutdown() {
	store.Client.Close()
}
<commit_msg>fix compilation<commit_after>package hbase

import (
	"bytes"
	"context"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/tsuna/gohbase"
	"github.com/tsuna/gohbase/hrpc"
	"io"
)

func init() {
	filer.Stores = append(filer.Stores, &HbaseStore{})
}

type HbaseStore struct {
	Client    gohbase.Client
	table     []byte
	cfKv      string
	cfMetaDir string
	column    string
}

func (store *HbaseStore) GetName() string {
	return "hbase"
}

func (store *HbaseStore) Initialize(configuration util.Configuration, prefix string) (err error) {
	return store.initialize(
		configuration.GetString(prefix+"zkquorum"),
		configuration.GetString(prefix+"table"),
	)
}

func (store *HbaseStore) initialize(zkquorum, table string) (err error) {
	store.Client = gohbase.NewClient(zkquorum)
	store.table = []byte(table)
	store.cfKv = "kv"
	store.cfMetaDir = "meta"
	store.column = "a"

	// check table exists
	key := "whatever"
	headers := map[string][]string{store.cfMetaDir: nil}
	get, err := hrpc.NewGet(context.Background(), store.table, []byte(key), hrpc.Families(headers))
	if err != nil {
		return fmt.Errorf("NewGet returned an error: %v", err)
	}
	_, err = store.Client.Get(get)
	if err != gohbase.TableNotFound {
		return nil
	}

	// create table
	adminClient := gohbase.NewAdminClient(zkquorum)
	cFamilies := []string{store.cfKv, store.cfMetaDir}
	cf := make(map[string]map[string]string, len(cFamilies))
	for _, f := range cFamilies {
		cf[f] = nil
	}
	ct := hrpc.NewCreateTable(context.Background(), []byte(table), cf)
	if err := adminClient.CreateTable(ct); err != nil {
		return err
	}

	return nil
}

func (store *HbaseStore) InsertEntry(ctx context.Context, entry *filer.Entry) error {
	value, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encoding %s %+v: %v", entry.FullPath, entry.Attr, err)
	}
	if len(entry.Chunks) > 50 {
		value = util.MaybeGzipData(value)
	}

	return store.doPut(ctx, store.cfMetaDir, []byte(entry.FullPath), value)
}

func (store *HbaseStore) UpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {
	return store.InsertEntry(ctx, entry)
}

func (store *HbaseStore) FindEntry(ctx context.Context, path util.FullPath) (entry *filer.Entry, err error) {
	value, err := store.doGet(ctx, store.cfMetaDir, []byte(path))
	if err != nil {
		if err == filer.ErrKvNotFound {
			return nil, filer_pb.ErrNotFound
		}
		return nil, err
	}

	entry = &filer.Entry{
		FullPath: path,
	}
	err = entry.DecodeAttributesAndChunks(util.MaybeDecompressData(value))
	if err != nil {
		return entry, fmt.Errorf("decode %s : %v", entry.FullPath, err)
	}
	return entry, nil
}

func (store *HbaseStore) DeleteEntry(ctx context.Context, path util.FullPath) (err error) {
	return store.doDelete(ctx, store.cfMetaDir, []byte(path))
}

func (store *HbaseStore) DeleteFolderChildren(ctx context.Context, path util.FullPath) (err error) {

	family := map[string][]string{store.cfMetaDir: {COLUMN_NAME}}
	expectedPrefix := []byte(path + "/")
	scan, err := hrpc.NewScanRange(ctx, store.table, expectedPrefix, nil, hrpc.Families(family))
	if err != nil {
		return err
	}

	scanner := store.Client.Scan(scan)
	defer scanner.Close()
	for {
		res, err := scanner.Next()
		if err != nil {
			break
		}
		if len(res.Cells) == 0 {
			continue
		}
		cell := res.Cells[0]

		if !bytes.HasPrefix(cell.Row, expectedPrefix) {
			break
		}
		fullpath := util.FullPath(cell.Row)
		dir, _ := fullpath.DirAndName()
		if dir != string(dir) {
			continue
		}

		err = store.doDelete(ctx, store.cfMetaDir, cell.Row)
		if err != nil {
			break
		}

	}
	return
}

func (store *HbaseStore) ListDirectoryEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int) ([]*filer.Entry, error) {
	return store.ListDirectoryPrefixedEntries(ctx, dirPath, startFileName, includeStartFile, limit, "")
}

func (store *HbaseStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int, prefix string) ([]*filer.Entry, error) {
	family := map[string][]string{store.cfMetaDir: {COLUMN_NAME}}
	expectedPrefix := []byte(string(dirPath) + "/" + prefix)
	scan, err := hrpc.NewScanRange(ctx, store.table, expectedPrefix, nil, hrpc.Families(family))
	if err != nil {
		return nil, err
	}

	var entries []*filer.Entry
	scanner := store.Client.Scan(scan)
	defer scanner.Close()
	for {
		res, err := scanner.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return entries, err
		}
		if len(res.Cells) == 0 {
			continue
		}
		cell := res.Cells[0]

		if !bytes.HasPrefix(cell.Row, expectedPrefix) {
			break
		}

		fullpath := util.FullPath(cell.Row)
		dir, fileName := fullpath.DirAndName()
		if dir != string(dirPath) {
			continue
		}

		value := cell.Value

		if fileName == startFileName && !includeStartFile {
			continue
		}

		limit--
		if limit < 0 {
			break
		}
		entry := &filer.Entry{
			FullPath: fullpath,
		}
		if decodeErr := entry.DecodeAttributesAndChunks(util.MaybeDecompressData(value)); decodeErr != nil {
			err = decodeErr
			glog.V(0).Infof("list %s : %v", entry.FullPath, err)
			break
		}
		entries = append(entries, entry)
	}

	return entries, nil
}

func (store *HbaseStore) BeginTransaction(ctx context.Context) (context.Context, error) {
	return ctx, nil
}

func (store *HbaseStore) CommitTransaction(ctx context.Context) error {
	return nil
}

func (store *HbaseStore) RollbackTransaction(ctx context.Context) error {
	return nil
}

func (store *HbaseStore) Shutdown() {
	store.Client.Close()
}
<|endoftext|>


<commit_before>package prifi

/*
* This is the internal part of the API. As probably the prifi-service will
* not have an external API, this will not have any API-functions.
 */

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"strconv"
	"errors"

	"github.com/dedis/cothority/app/lib/config"
	"github.com/dedis/cothority/log"
	"github.com/dedis/cothority/network"
	"github.com/dedis/cothority/sda"
	"github.com/lbarman/prifi_dev/sda/protocols"
)

// ServiceName is the name to refer to the Template service from another
// package.
const ServiceName = "PrifiService"

var serviceID sda.ServiceID

// Register Service with SDA
func init() {
	sda.RegisterNewService(ServiceName, newService)
	serviceID = sda.ServiceFactory.ServiceID(ServiceName)
	network.RegisterPacketType(ConnectionResponse{})
}

// This struct contains the state of the service
type Service struct {
	// We need to embed the ServiceProcessor, so that incoming messages
	// are correctly handled.
	*sda.ServiceProcessor
	group *config.Group
	Storage *Storage
	path    string
	role prifi.PriFiRole
}

// This structure will be saved, on the contrary of the 'Service'-structure
// which has per-service information stored
type Storage struct {
	TrusteeID string
}

type ConnectionRequest struct {}
type DisconnectionRequest struct {}
type ConnectionResponse struct {
	Status bool
}

// HandleConnection receives connection requests from other nodes.
// It decides when another PriFi protocol should be started.
func (s *Service) HandleConnection(from *network.ServerIdentity, req *ConnectionRequest) (network.Body, error) {
	log.Info(">>>> Received connection request ! <<<<<")
	// TODO: Check that we are relay, store who wants to start the protocol and start if there are enough participants
	return &ConnectionResponse{true}, nil
}

// HandleDisconnection receives disconnection requests.
// It must stop the current PriFi protocol.
func (s *Service) HandleDisconnection(from *network.ServerIdentity, req *DisconnectionRequest) (network.Body, error) {
	// TODO: This one will be a bit more complicated
	return &ConnectionResponse{true}, nil
}

// StartTrustee has to take a configuration and start the necessary
// protocols to enable the trustee-mode.
func (s *Service) StartTrustee(group *config.Group) error {
	log.Info("Service", s, "running in trustee mode")
	s.group = group
	s.role = prifi.Trustee

	// Inform the relay that we want to join the protocol
	err := sendConnectionRequest(group)
	if err != nil {
		log.Error("Connection failed:", err)
	}

	return nil
}

// StartRelay has to take a configuration and start the necessary
// protocols to enable the relay-mode.
func (s *Service) StartRelay(group *config.Group) error {
	log.Info("Service", s, "running in relay mode")
	s.group = group
	s.role = prifi.Relay
	/*
	var wrapper *prifi.PriFiSDAWrapper
	ids, relayIdentity := mapIdentities(group)

	// Start the PriFi protocol on a flat tree with the relay as root
	tree := group.Roster.GenerateNaryTreeWithRoot(100, &relayIdentity)
	pi, err := s.CreateProtocolService(prifi.ProtocolName, tree)

	if err != nil {
		log.Fatal("Unable to start Prifi protocol:", err)
	}

	// Assert that pi has type PriFiSDAWrapper
	wrapper = pi.(*prifi.PriFiSDAWrapper)

	wrapper.SetConfig(&prifi.PriFiSDAWrapperConfig{
		Identities: ids,
		Role: prifi.Relay,
	})
	wrapper.Start()
	*/
	return nil
}

// StartClient has to take a configuration and start the necessary
// protocols to enable the client-mode.
func (s *Service) StartClient(group *config.Group) error {
	log.Info("Service", s, "running in client mode")
	s.group = group
	s.role = prifi.Client

	// Inform the relay that we want to join the protocol
	err := sendConnectionRequest(group)
	if err != nil {
		log.Error("Connection failed:", err)
	}

	return nil
}

// NewProtocol is called on all nodes of a Tree (except the root, since it is
// the one starting the protocol) so it's the Service that will be called to
// generate the PI on all others node.
// If you use CreateProtocolSDA, this will not be called, as the SDA will
// instantiate the protocol on its own. If you need more control at the
// instantiation of the protocol, use CreateProtocolService, and you can
// give some extra-configuration to your protocol in here.
func (s *Service) NewProtocol(tn *sda.TreeNodeInstance, conf *sda.GenericConfig) (sda.ProtocolInstance, error) {
	log.Lvl5("Setting node configuration from service")

	ids, _ := mapIdentities(s.group)

	pi, err := prifi.NewPriFiSDAWrapperProtocol(tn)
	if err != nil {
		return nil, err
	}

	// Assert that pi has type PriFiSDAWrapper
	wrapper := pi.(*prifi.PriFiSDAWrapper)

	wrapper.SetConfig(&prifi.PriFiSDAWrapperConfig{
		Identities: ids,
		Role: s.role,
	})

	return wrapper, nil
}

// saves the actual identity
func (s *Service) save() {
	log.Lvl3("Saving service")
	b, err := network.MarshalRegisteredType(s.Storage)
	if err != nil {
		log.Error("Couldn't marshal service:", err)
	} else {
		err = ioutil.WriteFile(s.path+"/prifi.bin", b, 0660)
		if err != nil {
			log.Error("Couldn't save file:", err)
		}
	}
}

// Tries to load the configuration and updates if a configuration
// is found, else it returns an error.
func (s *Service) tryLoad() error {
	configFile := s.path + "/identity.bin"
	b, err := ioutil.ReadFile(configFile)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("Error while reading %s: %s", configFile, err)
	}
	if len(b) > 0 {
		_, msg, err := network.UnmarshalRegistered(b)
		if err != nil {
			return fmt.Errorf("Couldn't unmarshal: %s", err)
		}
		log.Lvl3("Successfully loaded")
		s.Storage = msg.(*Storage)
	}
	return nil
}

// newService receives the context and a path where it can write its
// configuration, if desired. As we don't know when the service will exit,
// we need to save the configuration on our own from time to time.
func newService(c *sda.Context, path string) sda.Service {
	log.Lvl4("Calling newService")
	s := &Service{
		ServiceProcessor: sda.NewServiceProcessor(c),
		path:             path,
	}
	if err := s.tryLoad(); err != nil {
		log.Error(err)
	}
	if err:= s.RegisterMessages(s.HandleConnection, s.HandleDisconnection); err != nil {
		log.Fatal("Could not register handlers:", err)
	}
	return s
}

// parseDescription extracts a PriFiIdentity from a string
func parseDescription(description string) (*prifi.PriFiIdentity, error) {
	desc := strings.Split(description, " ")
	if len(desc) == 1 && desc[0] == "relay" {
		return &prifi.PriFiIdentity{
			Role: prifi.Relay,
			Id: 0,
		}, nil
	} else if len(desc) == 2 {
		id, err := strconv.Atoi(desc[1]); if err != nil {
			return nil, errors.New("Unable to parse id:")
		} else {
			pid := prifi.PriFiIdentity{
				Id: id,
			}
			if desc[0] == "client" {
				pid.Role = prifi.Client
			} else if  desc[0] == "trustee" {
				pid.Role = prifi.Trustee
			} else {
				return nil, errors.New("Invalid role.")
			}
			return &pid, nil
		}
	} else {
		return nil, errors.New("Invalid description.")
	}
}

// mapIdentities reads the group configuration to assign PriFi roles
// to server identities and return them with the server
// identity of the relay.
func mapIdentities(group *config.Group) (map[network.Address]prifi.PriFiIdentity, network.ServerIdentity) {
	m := make(map[network.Address]prifi.PriFiIdentity)
	var relay network.ServerIdentity

	// Read the description of the nodes in the config file to assign them PriFi roles.
	nodeList := group.Roster.List
	for i := 0; i < len(nodeList); i++ {
		si := nodeList[i]
		id, err := parseDescription(group.GetDescription(si))
		if err != nil {
			log.Info("Cannot parse node description, skipping:", err)
		} else {
			m[si.Address] = *id
			if id.Role == prifi.Relay {
				relay = *si
			}
		}
	}

	// Check that there is exactly one relay and at least one trustee and client
	t, c, r := 0, 0, 0

	for _, v := range m {
		switch v.Role {
		case prifi.Relay: r++
		case prifi.Client: c++
		case prifi.Trustee: t++
		}
	}

	if !(t > 0 && c > 0 && r == 1) {
		log.Fatal("Config file does not contain exactly one relay and at least one trustee and client.")
	}

	return m, relay
}

func sendConnectionRequest(group *config.Group) error {
	client := sda.NewClient(ServiceName)
	_, relayIdentity := mapIdentities(group)
	reply, err := client.Send(&relayIdentity, &ConnectionRequest{})
	if e := network.ErrMsg(reply, err); e != nil {
		return e
	} else {
		res := reply.Msg.(ConnectionResponse)

		if res.Status {
			return nil
		} else {
			return errors.New("Connection request refused by the relay.")
		}
	}
}
<commit_msg>Relay starts PriFi as soon as enought clients and trustees are conected<commit_after>package prifi

/*
* This is the internal part of the API. As probably the prifi-service will
* not have an external API, this will not have any API-functions.
 */

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"strconv"
	"errors"

	"github.com/dedis/cothority/app/lib/config"
	"github.com/dedis/cothority/log"
	"github.com/dedis/cothority/network"
	"github.com/dedis/cothority/sda"
	"github.com/lbarman/prifi_dev/sda/protocols"
)

// ServiceName is the name to refer to the Template service from another
// package.
const ServiceName = "PrifiService"

var serviceID sda.ServiceID

// Register Service with SDA
func init() {
	sda.RegisterNewService(ServiceName, newService)
	serviceID = sda.ServiceFactory.ServiceID(ServiceName)
	network.RegisterPacketType(ConnectionResponse{})
}

type WaitQueue struct {
	nTrustee int
	nClient int
}

// This struct contains the state of the service
type Service struct {
	// We need to embed the ServiceProcessor, so that incoming messages
	// are correctly handled.
	*sda.ServiceProcessor
	group *config.Group
	Storage *Storage
	path    string
	role prifi.PriFiRole
	identityMap *map[network.Address]prifi.PriFiIdentity
	relayIdentity *network.ServerIdentity
	waitQueue *WaitQueue
}

// This structure will be saved, on the contrary of the 'Service'-structure
// which has per-service information stored
type Storage struct {
	TrusteeID string
}

type ConnectionRequest struct {
	Role prifi.PriFiRole
}
type DisconnectionRequest struct {}
type ConnectionResponse struct {
	Status bool
}

// HandleConnection receives connection requests from other nodes.
// It decides when another PriFi protocol should be started.
func (s *Service) HandleConnection(from *network.ServerIdentity, req *ConnectionRequest) (network.Body, error) {
	log.Lvl3("Received connection request")

	if s.role != prifi.Relay {
		return nil, errors.New("This host is not the relay.")
	}

	// No integrity checks are performed here for simplicity.
	// If a client sends a connection request with a role which
	// is not it's own or sends multiple connection requests,
	// this will be detected when the PriFi protocol attempts to start.
	switch req.Role {
	case prifi.Client: s.waitQueue.nClient++
	case prifi.Trustee: s.waitQueue.nTrustee++
	default: log.Error("Connection request from invalid node, ignoring")
	}

	if s.waitQueue.nClient >= 2 && s.waitQueue.nTrustee >= 1 {
		s.startPriFi()
	}

	return &ConnectionResponse{true}, nil
}

// HandleDisconnection receives disconnection requests.
// It must stop the current PriFi protocol.
func (s *Service) HandleDisconnection(from *network.ServerIdentity, req *DisconnectionRequest) (network.Body, error) {
	// TODO: This one will be a bit more complicated
	return &ConnectionResponse{true}, nil
}

// StartTrustee has to take a configuration and start the necessary
// protocols to enable the trustee-mode.
func (s *Service) StartTrustee(group *config.Group) error {
	log.Info("Service", s, "running in trustee mode")
	s.role = prifi.Trustee
	s.readGroup(group)

	// Inform the relay that we want to join the protocol
	err := s.sendConnectionRequest()
	if err != nil {
		log.Error("Connection failed:", err)
	}

	return nil
}

// StartRelay has to take a configuration and start the necessary
// protocols to enable the relay-mode.
func (s *Service) StartRelay(group *config.Group) error {
	log.Info("Service", s, "running in relay mode")
	s.role = prifi.Relay
	s.readGroup(group)
	s.waitQueue = &WaitQueue{0, 0}

	return nil
}

// StartClient has to take a configuration and start the necessary
// protocols to enable the client-mode.
func (s *Service) StartClient(group *config.Group) error {
	log.Info("Service", s, "running in client mode")
	s.role = prifi.Client
	s.readGroup(group)

	// Inform the relay that we want to join the protocol
	err := s.sendConnectionRequest()
	if err != nil {
		log.Error("Connection failed:", err)
	}

	return nil
}

// NewProtocol is called on all nodes of a Tree (except the root, since it is
// the one starting the protocol) so it's the Service that will be called to
// generate the PI on all others node.
// If you use CreateProtocolSDA, this will not be called, as the SDA will
// instantiate the protocol on its own. If you need more control at the
// instantiation of the protocol, use CreateProtocolService, and you can
// give some extra-configuration to your protocol in here.
func (s *Service) NewProtocol(tn *sda.TreeNodeInstance, conf *sda.GenericConfig) (sda.ProtocolInstance, error) {
	log.Lvl5("Setting node configuration from service")

	pi, err := prifi.NewPriFiSDAWrapperProtocol(tn)
	if err != nil {
		return nil, err
	}

	// Assert that pi has type PriFiSDAWrapper
	wrapper := pi.(*prifi.PriFiSDAWrapper)

	wrapper.SetConfig(&prifi.PriFiSDAWrapperConfig{
		Identities: *s.identityMap,
		Role: s.role,
	})

	return wrapper, nil
}

// saves the actual identity
func (s *Service) save() {
	log.Lvl3("Saving service")
	b, err := network.MarshalRegisteredType(s.Storage)
	if err != nil {
		log.Error("Couldn't marshal service:", err)
	} else {
		err = ioutil.WriteFile(s.path+"/prifi.bin", b, 0660)
		if err != nil {
			log.Error("Couldn't save file:", err)
		}
	}
}

// Tries to load the configuration and updates if a configuration
// is found, else it returns an error.
func (s *Service) tryLoad() error {
	configFile := s.path + "/identity.bin"
	b, err := ioutil.ReadFile(configFile)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("Error while reading %s: %s", configFile, err)
	}
	if len(b) > 0 {
		_, msg, err := network.UnmarshalRegistered(b)
		if err != nil {
			return fmt.Errorf("Couldn't unmarshal: %s", err)
		}
		log.Lvl3("Successfully loaded")
		s.Storage = msg.(*Storage)
	}
	return nil
}

// newService receives the context and a path where it can write its
// configuration, if desired. As we don't know when the service will exit,
// we need to save the configuration on our own from time to time.
func newService(c *sda.Context, path string) sda.Service {
	log.Lvl4("Calling newService")
	s := &Service{
		ServiceProcessor: sda.NewServiceProcessor(c),
		path:             path,
	}
	if err := s.tryLoad(); err != nil {
		log.Error(err)
	}
	if err:= s.RegisterMessages(s.HandleConnection, s.HandleDisconnection); err != nil {
		log.Fatal("Could not register handlers:", err)
	}
	return s
}

// parseDescription extracts a PriFiIdentity from a string
func parseDescription(description string) (*prifi.PriFiIdentity, error) {
	desc := strings.Split(description, " ")
	if len(desc) == 1 && desc[0] == "relay" {
		return &prifi.PriFiIdentity{
			Role: prifi.Relay,
			Id: 0,
		}, nil
	} else if len(desc) == 2 {
		id, err := strconv.Atoi(desc[1]); if err != nil {
			return nil, errors.New("Unable to parse id:")
		} else {
			pid := prifi.PriFiIdentity{
				Id: id,
			}
			if desc[0] == "client" {
				pid.Role = prifi.Client
			} else if  desc[0] == "trustee" {
				pid.Role = prifi.Trustee
			} else {
				return nil, errors.New("Invalid role.")
			}
			return &pid, nil
		}
	} else {
		return nil, errors.New("Invalid description.")
	}
}

// mapIdentities reads the group configuration to assign PriFi roles
// to server identities and return them with the server
// identity of the relay.
func mapIdentities(group *config.Group) (map[network.Address]prifi.PriFiIdentity, network.ServerIdentity) {
	m := make(map[network.Address]prifi.PriFiIdentity)
	var relay network.ServerIdentity

	// Read the description of the nodes in the config file to assign them PriFi roles.
	nodeList := group.Roster.List
	for i := 0; i < len(nodeList); i++ {
		si := nodeList[i]
		id, err := parseDescription(group.GetDescription(si))
		if err != nil {
			log.Info("Cannot parse node description, skipping:", err)
		} else {
			m[si.Address] = *id
			if id.Role == prifi.Relay {
				relay = *si
			}
		}
	}

	// Check that there is exactly one relay and at least one trustee and client
	t, c, r := 0, 0, 0

	for _, v := range m {
		switch v.Role {
		case prifi.Relay: r++
		case prifi.Client: c++
		case prifi.Trustee: t++
		}
	}

	if !(t > 0 && c > 0 && r == 1) {
		log.Fatal("Config file does not contain exactly one relay and at least one trustee and client.")
	}

	return m, relay
}

// readGroup reads the group description and sets up the Service struct fields
// accordingly. It *MUST* be called first when the node is started.
func (s *Service) readGroup(group *config.Group) {
	ids, relayId := mapIdentities(group)
	s.identityMap = &ids
	s.relayIdentity = &relayId
	s.group = group
}

// sendConnectionRequest sends a connection request to the relay.
// It is called by the client and trustee services at startup to
// announce themselves to the relay.
func (s *Service) sendConnectionRequest() error {
	client := sda.NewClient(ServiceName)
	reply, err := client.Send(s.relayIdentity, &ConnectionRequest{s.role})
	if e := network.ErrMsg(reply, err); e != nil {
		return e
	} else {
		res := reply.Msg.(ConnectionResponse)

		if res.Status {
			return nil
		} else {
			return errors.New("Connection request refused by the relay.")
		}
	}
	return nil
}

// startPriFi starts a PriFi protocol. It is called
// by the relay as soon as enough participants are
// ready (one trustee and two clients).
func (s *Service) startPriFi() {
	if s.role != prifi.Relay {
		log.Error("Trying to start PriFi protocol from a non-relay node.")
		return;
	}

	var wrapper *prifi.PriFiSDAWrapper

	// Start the PriFi protocol on a flat tree with the relay as root
	tree := s.group.Roster.GenerateNaryTreeWithRoot(100, s.relayIdentity)
	pi, err := s.CreateProtocolService(prifi.ProtocolName, tree)

	if err != nil {
		log.Fatal("Unable to start Prifi protocol:", err)
	}

	// Assert that pi has type PriFiSDAWrapper
	wrapper = pi.(*prifi.PriFiSDAWrapper)

	wrapper.SetConfig(&prifi.PriFiSDAWrapperConfig{
		Identities: *s.identityMap,
		Role: prifi.Relay,
	})
	wrapper.Start()
}
<|endoftext|>


<commit_before>/* Copyright (C) 2015-2020 김운하 (unha.kim@ghts.org)

이 파일은 GHTS의 일부입니다.

이 프로그램은 자유 소프트웨어입니다.
소프트웨어의 피양도자는 자유 소프트웨어 재단이 공표한 GNU LGPL 2.1판
규정에 따라 프로그램을 개작하거나 재배포할 수 있습니다.

이 프로그램은 유용하게 사용될 수 있으리라는 희망에서 배포되고 있지만,
특정한 목적에 적합하다거나, 이익을 안겨줄 수 있다는 묵시적인 보증을 포함한
어떠한 형태의 보증도 제공하지 않습니다.
보다 자세한 사항에 대해서는 GNU LGPL 2.1판을 참고하시기 바랍니다.
GNU LGPL 2.1판은 이 프로그램과 함께 제공됩니다.
만약, 이 문서가 누락되어 있다면 자유 소프트웨어 재단으로 문의하시기 바랍니다.
(자유 소프트웨어 재단 : Free Software Foundation, Inc.,
59 Temple Place - Suite 330, Boston, MA 02111-1307, USA)

Copyright (C) 2015-2020년 UnHa Kim (unha.kim@ghts.org)

This file is part of GHTS.

GHTS is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, version 2.1 of the License.

GHTS is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with GHTS.  If not, see <http://www.gnu.org/licenses/>. */

package c

// #cgo CFLAGS: -Wall
// #include <stdlib.h>
import "C"
import (
	"golang.org/x/text/encoding/korean"
	"strings"
	"unsafe"
)

func F2C문자열(go문자열 string) *C.char {
	return C.CString(go문자열)
}

func F2Go문자열(c문자열_포인터 unsafe.Pointer) string {
	return C.GoString((*C.char)(c문자열_포인터))
}

func F2문자열_EUC_KR(c문자열_포인터 unsafe.Pointer) string {
	var 바이트_모음 []byte
	길이 := 64

	for {
		바이트_모음 = F2Go바이트_모음(c문자열_포인터, 길이)

		null문자_인덱스 := strings.Index(string(바이트_모음), "\x00")

		if null문자_인덱스 >= 0 {
			break
		}

		길이 += 64
	}

	return strings.TrimSpace(f2문자열_EUC_KR(바이트_모음))
}

func f2문자열_EUC_KR(바이트_모음 []byte) string {
	null문자_인덱스 := strings.Index(string(바이트_모음), "\x00")

	if null문자_인덱스 >= 0 {
		바이트_모음 = 바이트_모음[:null문자_인덱스]
	}

	바이트_모음_utf8, 에러 := korean.EUCKR.NewDecoder().Bytes(바이트_모음)
	if 에러 != nil {
		if len(바이트_모음) > 0 {
			return f2문자열_EUC_KR(바이트_모음[:len(바이트_모음)-1])
		}

		return string(바이트_모음)
	}

	return string(바이트_모음_utf8)
}

func F2Go바이트_모음(c데이터 unsafe.Pointer, 길이 int) []byte {
	return C.GoBytes(c데이터, C.int(길이))
}

func F메모리_해제(포인터 unsafe.Pointer) {
	C.free(포인터)
}
<commit_msg>테스트 에러 디버깅<commit_after>/* Copyright (C) 2015-2020 김운하 (unha.kim@ghts.org)

이 파일은 GHTS의 일부입니다.

이 프로그램은 자유 소프트웨어입니다.
소프트웨어의 피양도자는 자유 소프트웨어 재단이 공표한 GNU LGPL 2.1판
규정에 따라 프로그램을 개작하거나 재배포할 수 있습니다.

이 프로그램은 유용하게 사용될 수 있으리라는 희망에서 배포되고 있지만,
특정한 목적에 적합하다거나, 이익을 안겨줄 수 있다는 묵시적인 보증을 포함한
어떠한 형태의 보증도 제공하지 않습니다.
보다 자세한 사항에 대해서는 GNU LGPL 2.1판을 참고하시기 바랍니다.
GNU LGPL 2.1판은 이 프로그램과 함께 제공됩니다.
만약, 이 문서가 누락되어 있다면 자유 소프트웨어 재단으로 문의하시기 바랍니다.
(자유 소프트웨어 재단 : Free Software Foundation, Inc.,
59 Temple Place - Suite 330, Boston, MA 02111-1307, USA)

Copyright (C) 2015-2020년 UnHa Kim (unha.kim@ghts.org)

This file is part of GHTS.

GHTS is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, version 2.1 of the License.

GHTS is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with GHTS.  If not, see <http://www.gnu.org/licenses/>. */

package c

// #cgo CFLAGS: -Wall
// #include <stdlib.h>
import "C"
import (
	"golang.org/x/text/encoding/korean"
	"strings"
	"unsafe"
)

func F2C문자열(go문자열 string) *C.char {
	return C.CString(go문자열)
}

func F2Go문자열(c문자열_포인터 unsafe.Pointer) string {
	return C.GoString((*C.char)(c문자열_포인터))
}

func F2문자열_EUC_KR(c문자열_포인터 unsafe.Pointer) string {
	const 단위 = 32

	var 바이트_모음 []byte
	길이 := 단위

	for {
		바이트_모음 = F2Go바이트_모음(c문자열_포인터, 길이)

		null문자_인덱스 := strings.Index(string(바이트_모음), "\x00")

		if null문자_인덱스 >= 0 {
			break
		}

		길이 += 단위
	}

	return strings.TrimSpace(f2문자열_EUC_KR(바이트_모음))
}

func f2문자열_EUC_KR(바이트_모음 []byte) string {
	null문자_인덱스 := strings.Index(string(바이트_모음), "\x00")

	if null문자_인덱스 >= 0 {
		바이트_모음 = 바이트_모음[:null문자_인덱스]
	}

	바이트_모음_utf8, 에러 := korean.EUCKR.NewDecoder().Bytes(바이트_모음)
	if 에러 != nil {
		if len(바이트_모음) > 0 {
			return f2문자열_EUC_KR(바이트_모음[:len(바이트_모음)-1])
		}

		return string(바이트_모음)
	}

	return string(바이트_모음_utf8)
}

func F2Go바이트_모음(c데이터 unsafe.Pointer, 길이 int) []byte {
	return C.GoBytes(c데이터, C.int(길이))
}

func F메모리_해제(포인터 unsafe.Pointer) {
	C.free(포인터)
}
<|endoftext|>


<commit_before>package e2e

import (
	"io/ioutil"
	"net/http"
	"testing"
	"time"
)

func TestWebdash(t *testing.T) {
	setLogOutput(t)
	// Get the webdash health check endpoint
	address := fun.Conf.Server.HTTPAddress()
	client := http.Client{Timeout: 10 * time.Second}
	resp, err := client.Get(address + "/health.html")
	if err != nil {
		t.Fatal(err)
	}

	// Get the response body
	defer resp.Body.Close()
	b, berr := ioutil.ReadAll(resp.Body)
	if berr != nil {
		t.Fatal(berr)
	}

	if string(b) != "OK\n" {
		t.Fatal("Webdash health check fail", string(b))
	}
}
<commit_msg>webdash: fix health check test<commit_after>package e2e

import (
	"io/ioutil"
	"net/http"
	"testing"
	"time"
)

func TestWebdash(t *testing.T) {
	setLogOutput(t)
	// Get the webdash health check endpoint
	address := fun.Conf.Server.HTTPAddress()
	client := http.Client{Timeout: 10 * time.Second}
	resp, err := client.Get(address + "/static/health.html")
	if err != nil {
		t.Fatal(err)
	}

	// Get the response body
	defer resp.Body.Close()
	b, berr := ioutil.ReadAll(resp.Body)
	if berr != nil {
		t.Fatal(berr)
	}

	if string(b) != "OK\n" {
		t.Fatal("Webdash health check fail", string(b))
	}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/nu7hatch/gouuid"

	"github.com/cloudfoundry-incubator/auctioneer"
	"github.com/cloudfoundry-incubator/auctioneer/auctionmetricemitterdelegate"
	"github.com/cloudfoundry-incubator/auctioneer/auctionrunnerdelegate"
	"github.com/cloudfoundry-incubator/auctioneer/handlers"
	"github.com/cloudfoundry-incubator/bbs"
	"github.com/cloudfoundry-incubator/cf-debug-server"
	cf_lager "github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry-incubator/consuladapter"
	"github.com/cloudfoundry-incubator/locket"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/localip"

	"github.com/cloudfoundry-incubator/auction/auctionrunner"
	"github.com/cloudfoundry-incubator/auction/auctiontypes"
	"github.com/cloudfoundry/dropsonde"
	"github.com/cloudfoundry/gunk/workpool"
	"github.com/pivotal-golang/clock"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
)

var communicationTimeout = flag.Duration(
	"communicationTimeout",
	10*time.Second,
	"Timeout applied to all HTTP requests.",
)

var cellStateTimeout = flag.Duration(
	"cellStateTimeout",
	1*time.Second,
	"Timeout applied to HTTP requests to the Cell State endpoint.",
)

var consulCluster = flag.String(
	"consulCluster",
	"",
	"comma-separated list of consul server addresses (ip:port)",
)

var dropsondePort = flag.Int(
	"dropsondePort",
	3457,
	"port the local metron agent is listening on",
)

var lockTTL = flag.Duration(
	"lockTTL",
	locket.LockTTL,
	"TTL for service lock",
)

var lockRetryInterval = flag.Duration(
	"lockRetryInterval",
	locket.RetryInterval,
	"interval to wait before retrying a failed lock acquisition",
)

var listenAddr = flag.String(
	"listenAddr",
	"0.0.0.0:9016",
	"host:port to serve auction and LRP stop requests on",
)

var bbsAddress = flag.String(
	"bbsAddress",
	"",
	"Address to the BBS Server",
)

var bbsCACert = flag.String(
	"bbsCACert",
	"",
	"path to certificate authority cert used for mutually authenticated TLS BBS communication",
)

var bbsClientCert = flag.String(
	"bbsClientCert",
	"",
	"path to client cert used for mutually authenticated TLS BBS communication",
)

var bbsClientKey = flag.String(
	"bbsClientKey",
	"",
	"path to client key used for mutually authenticated TLS BBS communication",
)

var bbsClientSessionCacheSize = flag.Int(
	"bbsClientSessionCacheSize",
	0,
	"Capacity of the ClientSessionCache option on the TLS configuration. If zero, golang's default will be used",
)

var bbsMaxIdleConnsPerHost = flag.Int(
	"bbsMaxIdleConnsPerHost",
	0,
	"Controls the maximum number of idle (keep-alive) connctions per host. If zero, golang's default will be used",
)

var auctionRunnerWorkers = flag.Int(
	"auctionRunnerWorkers",
	1000,
	"Max concurrency for cell operations in the auction runner",
)

const (
	auctionRunnerTimeout = 10 * time.Second
	dropsondeOrigin      = "auctioneer"
	serverProtocol       = "http"
)

func main() {
	cf_debug_server.AddFlags(flag.CommandLine)
	cf_lager.AddFlags(flag.CommandLine)
	flag.Parse()

	cf_http.Initialize(*communicationTimeout)

	logger, reconfigurableSink := cf_lager.New("auctioneer")
	initializeDropsonde(logger)

	if err := validateBBSAddress(); err != nil {
		logger.Fatal("invalid-bbs-address", err)
	}

	client, err := consuladapter.NewClient(*consulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	sessionMgr := consuladapter.NewSessionManager(client)
	consulSession, err := consuladapter.NewSession("auctioneer", *lockTTL, client, sessionMgr)
	if err != nil {
		logger.Fatal("consul-session-failed", err)
	}

	clock := clock.NewClock()
	auctioneerServiceClient := auctioneer.NewServiceClient(consulSession, clock)

	auctionRunner := initializeAuctionRunner(logger, *cellStateTimeout, initializeBBSClient(logger))
	auctionServer := initializeAuctionServer(logger, auctionRunner)
	lockMaintainer := initializeLockMaintainer(logger, auctioneerServiceClient)

	members := grouper.Members{
		{"lock-maintainer", lockMaintainer},
		{"auction-runner", auctionRunner},
		{"auction-server", auctionServer},
	}

	if dbgAddr := cf_debug_server.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", cf_debug_server.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))

	logger.Info("started")

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func initializeAuctionRunner(logger lager.Logger, cellStateTimeout time.Duration, bbsClient bbs.Client) auctiontypes.AuctionRunner {
	httpClient := cf_http.NewClient()
	stateClient := cf_http.NewCustomTimeoutClient(cellStateTimeout)
	repClientFactory := rep.NewClientFactory(httpClient, stateClient)

	delegate := auctionrunnerdelegate.New(repClientFactory, bbsClient, logger)
	metricEmitter := auctionmetricemitterdelegate.New()
	workPool, err := workpool.NewWorkPool(*auctionRunnerWorkers)
	if err != nil {
		logger.Fatal("failed-to-construct-auction-runner-workpool", err, lager.Data{"num-workers": *auctionRunnerWorkers}) // should never happen
	}

	return auctionrunner.New(
		delegate,
		metricEmitter,
		clock.NewClock(),
		workPool,
		logger,
	)
}

func initializeDropsonde(logger lager.Logger) {
	dropsondeDestination := fmt.Sprint("localhost:", *dropsondePort)
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func initializeAuctionServer(logger lager.Logger, runner auctiontypes.AuctionRunner) ifrit.Runner {
	return http_server.New(*listenAddr, handlers.New(runner, logger))
}

func initializeLockMaintainer(logger lager.Logger, serviceClient auctioneer.ServiceClient) ifrit.Runner {
	uuid, err := uuid.NewV4()
	if err != nil {
		logger.Fatal("Couldn't generate uuid", err)
	}

	localIP, err := localip.LocalIP()
	if err != nil {
		logger.Fatal("Couldn't determine local IP", err)
	}

	port := strings.Split(*listenAddr, ":")[1]
	address := fmt.Sprintf("%s://%s:%s", serverProtocol, localIP, port)
	auctioneerPresence := auctioneer.NewPresence(uuid.String(), address)

	lockMaintainer, err := serviceClient.NewAuctioneerLockRunner(logger, auctioneerPresence, *lockRetryInterval)
	if err != nil {
		logger.Fatal("Couldn't create lock maintainer", err)
	}

	return lockMaintainer
}

func validateBBSAddress() error {
	if *bbsAddress == "" {
		return errors.New("bbsAddress is required")
	}
	return nil
}

func initializeBBSClient(logger lager.Logger) bbs.Client {
	bbsURL, err := url.Parse(*bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(*bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(*bbsAddress, *bbsCACert, *bbsClientCert, *bbsClientKey, *bbsClientSessionCacheSize, *bbsMaxIdleConnsPerHost)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}
<commit_msg>Add startingContainerWeight flag<commit_after>package main

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

	"github.com/nu7hatch/gouuid"

	"github.com/cloudfoundry-incubator/auctioneer"
	"github.com/cloudfoundry-incubator/auctioneer/auctionmetricemitterdelegate"
	"github.com/cloudfoundry-incubator/auctioneer/auctionrunnerdelegate"
	"github.com/cloudfoundry-incubator/auctioneer/handlers"
	"github.com/cloudfoundry-incubator/bbs"
	"github.com/cloudfoundry-incubator/cf-debug-server"
	cf_lager "github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry-incubator/consuladapter"
	"github.com/cloudfoundry-incubator/locket"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/localip"

	"github.com/cloudfoundry-incubator/auction/auctionrunner"
	"github.com/cloudfoundry-incubator/auction/auctiontypes"
	"github.com/cloudfoundry/dropsonde"
	"github.com/cloudfoundry/gunk/workpool"
	"github.com/pivotal-golang/clock"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
)

var communicationTimeout = flag.Duration(
	"communicationTimeout",
	10*time.Second,
	"Timeout applied to all HTTP requests.",
)

var cellStateTimeout = flag.Duration(
	"cellStateTimeout",
	1*time.Second,
	"Timeout applied to HTTP requests to the Cell State endpoint.",
)

var consulCluster = flag.String(
	"consulCluster",
	"",
	"comma-separated list of consul server addresses (ip:port)",
)

var dropsondePort = flag.Int(
	"dropsondePort",
	3457,
	"port the local metron agent is listening on",
)

var lockTTL = flag.Duration(
	"lockTTL",
	locket.LockTTL,
	"TTL for service lock",
)

var lockRetryInterval = flag.Duration(
	"lockRetryInterval",
	locket.RetryInterval,
	"interval to wait before retrying a failed lock acquisition",
)

var listenAddr = flag.String(
	"listenAddr",
	"0.0.0.0:9016",
	"host:port to serve auction and LRP stop requests on",
)

var bbsAddress = flag.String(
	"bbsAddress",
	"",
	"Address to the BBS Server",
)

var bbsCACert = flag.String(
	"bbsCACert",
	"",
	"path to certificate authority cert used for mutually authenticated TLS BBS communication",
)

var bbsClientCert = flag.String(
	"bbsClientCert",
	"",
	"path to client cert used for mutually authenticated TLS BBS communication",
)

var bbsClientKey = flag.String(
	"bbsClientKey",
	"",
	"path to client key used for mutually authenticated TLS BBS communication",
)

var bbsClientSessionCacheSize = flag.Int(
	"bbsClientSessionCacheSize",
	0,
	"Capacity of the ClientSessionCache option on the TLS configuration. If zero, golang's default will be used",
)

var bbsMaxIdleConnsPerHost = flag.Int(
	"bbsMaxIdleConnsPerHost",
	0,
	"Controls the maximum number of idle (keep-alive) connctions per host. If zero, golang's default will be used",
)

var auctionRunnerWorkers = flag.Int(
	"auctionRunnerWorkers",
	1000,
	"Max concurrency for cell operations in the auction runner",
)

var startingContainerWeight = flag.Float64(
	"startingContainerWeight",
	0.25,
	"Factor to bias against cells with starting containers (0.0 - 1.0)",
)

const (
	auctionRunnerTimeout = 10 * time.Second
	dropsondeOrigin      = "auctioneer"
	serverProtocol       = "http"
)

func main() {
	cf_debug_server.AddFlags(flag.CommandLine)
	cf_lager.AddFlags(flag.CommandLine)
	flag.Parse()

	cf_http.Initialize(*communicationTimeout)

	logger, reconfigurableSink := cf_lager.New("auctioneer")
	initializeDropsonde(logger)

	if err := validateBBSAddress(); err != nil {
		logger.Fatal("invalid-bbs-address", err)
	}

	client, err := consuladapter.NewClient(*consulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	sessionMgr := consuladapter.NewSessionManager(client)
	consulSession, err := consuladapter.NewSession("auctioneer", *lockTTL, client, sessionMgr)
	if err != nil {
		logger.Fatal("consul-session-failed", err)
	}

	clock := clock.NewClock()
	auctioneerServiceClient := auctioneer.NewServiceClient(consulSession, clock)

	auctionRunner := initializeAuctionRunner(logger, *cellStateTimeout,
		initializeBBSClient(logger), *startingContainerWeight)
	auctionServer := initializeAuctionServer(logger, auctionRunner)
	lockMaintainer := initializeLockMaintainer(logger, auctioneerServiceClient)

	members := grouper.Members{
		{"lock-maintainer", lockMaintainer},
		{"auction-runner", auctionRunner},
		{"auction-server", auctionServer},
	}

	if dbgAddr := cf_debug_server.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", cf_debug_server.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))

	logger.Info("started")

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func initializeAuctionRunner(logger lager.Logger, cellStateTimeout time.Duration, bbsClient bbs.Client, startingContainerWeight float64) auctiontypes.AuctionRunner {
	httpClient := cf_http.NewClient()
	stateClient := cf_http.NewCustomTimeoutClient(cellStateTimeout)
	repClientFactory := rep.NewClientFactory(httpClient, stateClient)

	delegate := auctionrunnerdelegate.New(repClientFactory, bbsClient, logger)
	metricEmitter := auctionmetricemitterdelegate.New()
	workPool, err := workpool.NewWorkPool(*auctionRunnerWorkers)
	if err != nil {
		logger.Fatal("failed-to-construct-auction-runner-workpool", err, lager.Data{"num-workers": *auctionRunnerWorkers}) // should never happen
	}

	return auctionrunner.New(
		logger,
		delegate,
		metricEmitter,
		clock.NewClock(),
		workPool,
		startingContainerWeight,
	)
}

func initializeDropsonde(logger lager.Logger) {
	dropsondeDestination := fmt.Sprint("localhost:", *dropsondePort)
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func initializeAuctionServer(logger lager.Logger, runner auctiontypes.AuctionRunner) ifrit.Runner {
	return http_server.New(*listenAddr, handlers.New(runner, logger))
}

func initializeLockMaintainer(logger lager.Logger, serviceClient auctioneer.ServiceClient) ifrit.Runner {
	uuid, err := uuid.NewV4()
	if err != nil {
		logger.Fatal("Couldn't generate uuid", err)
	}

	localIP, err := localip.LocalIP()
	if err != nil {
		logger.Fatal("Couldn't determine local IP", err)
	}

	port := strings.Split(*listenAddr, ":")[1]
	address := fmt.Sprintf("%s://%s:%s", serverProtocol, localIP, port)
	auctioneerPresence := auctioneer.NewPresence(uuid.String(), address)

	lockMaintainer, err := serviceClient.NewAuctioneerLockRunner(logger, auctioneerPresence, *lockRetryInterval)
	if err != nil {
		logger.Fatal("Couldn't create lock maintainer", err)
	}

	return lockMaintainer
}

func validateBBSAddress() error {
	if *bbsAddress == "" {
		return errors.New("bbsAddress is required")
	}
	return nil
}

func initializeBBSClient(logger lager.Logger) bbs.Client {
	bbsURL, err := url.Parse(*bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(*bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(*bbsAddress, *bbsCACert, *bbsClientCert, *bbsClientKey, *bbsClientSessionCacheSize, *bbsMaxIdleConnsPerHost)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}
<|endoftext|>


<commit_before>/*
Package runfile implements sensorbee runfile command. This command is provided
separately from run command to reduce the footprint of runfile command itself.
For example, if sensorbee doesn't need to provide sensorbee run command and only
offers sensorbee runfile, the footprint of sensorbee command might be reduced a
lot due to low functionality of this command.
*/
package runfile

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"pfi/sensorbee/sensorbee/bql"
	"pfi/sensorbee/sensorbee/bql/parser"
	"pfi/sensorbee/sensorbee/bql/udf"
	"pfi/sensorbee/sensorbee/core"
	"pfi/sensorbee/sensorbee/data"
	"pfi/sensorbee/sensorbee/server/config"
	"pfi/sensorbee/sensorbee/server/udsstorage"
)

// SetUp sets up a command for running single BQL file.
func SetUp() cli.Command {
	cmd := cli.Command{
		Name:        "runfile",
		Usage:       "run a BQL file",
		Description: "runfile command runs a BQL file",
		Action:      Run,
	}

	cmd.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "config, c",
			Value:  "",
			Usage:  "file path of a config file in YAML format (only logging and storage sections are used)",
			EnvVar: "SENSORBEE_CONFIG",
		},
		cli.StringSliceFlag{
			Name:  "save-uds, s",
			Value: &cli.StringSlice{},
			Usage: "save UDSs after all tuples are processed",
		},
	}
	return cmd
}

// Run runs "runfile" command.
func Run(c *cli.Context) {
	// TODO: Merge this implementation with cmd/run

	if len(c.Args()) != 1 {
		cli.ShowSubcommandHelp(c)
		os.Exit(1)
	}

	conf, err := func() (*config.Config, error) {
		if c.IsSet("config") {
			p := c.String("config")
			in, err := ioutil.ReadFile(p)
			if err != nil {
				return nil, fmt.Errorf("Cannot read the config file %v: %v\n", p, err)
			}

			var yml map[string]interface{}
			if err := yaml.Unmarshal(in, &yml); err != nil {
				return nil, fmt.Errorf("Cannot parse the config file %v: %v\n", p, err)
			}
			m, err := data.NewMap(yml)
			if err != nil {
				return nil, fmt.Errorf("The config file %v has invalid values: %v\n", p, err)
			}
			return config.New(m)

		} else {
			// Currently there's no required parameters. However, when a required
			// parameter is added to Config, remove this else block and add a
			// default value like "/etc/sensorbee/config.yaml" to the config option.
			return config.New(data.Map{})
		}
	}()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	w, err := conf.Logging.CreateWriter()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	logger := logrus.New()
	logger.Out = w

	udsStorage, err := setUpUDSStorage(&conf.Storage.UDS)
	if err != nil {
		logger.WithField("err", err).Error("Cannot set up a storage for UDSs")
		os.Exit(1)
	}

	logger.Info("Setting up a topology")
	bqlFile := c.Args()[0]
	tb, err := setUpTopology(logger, conf, udsStorage)
	if err != nil {
		logger.WithField("err", err).Error("Cannot set up the topology")
		os.Exit(1)
	}

	if err := setUpBQLStmt(tb, bqlFile); err != nil {
		logger.WithFields(logrus.Fields{
			"err":      err,
			"bql_file": bqlFile,
		}).Error("Cannot set up BQL statement")
		os.Exit(1)
	}

	defer func() {
		logger.Info("Waiting for all nodes to finish processing tuples")
		if err := tb.Topology().Stop(); err != nil {
			logger.WithField("err", err).Error("Cannot stop the topology")
			os.Exit(1)
		}

		saveUDSs := c.StringSlice("save-uds")
		if err := saveStates(tb, saveUDSs); err != nil {
			logger.WithFields(logrus.Fields{
				"err":      err,
				"topology": tb.Topology().Name(),
			}).Error("Cannot save a UDS")
			os.Exit(1)
		}
		// TODO: Terminate all shared states
	}()

	logger.Info("Starting the topology")
	for name, s := range tb.Topology().Sources() {
		if err := s.Resume(); err != nil {
			logger.WithFields(logrus.Fields{
				"err":    err,
				"source": name,
			}).Error("Cannot resume the source")
			os.Exit(1)
		}
	}

	for _, s := range tb.Topology().Sources() {
		// TODO: error check if necessary
		s.State().Wait(core.TSStopped)
	}
	logger.Info("All sources has been stopped.")

}

func setUpUDSStorage(conf *config.UDSStorage) (udf.UDSStorage, error) {
	// TODO: merge this implementation with server/context.go
	// Parameters are already validated in conf
	switch conf.Type {
	case "in_memory":
		return udf.NewInMemoryUDSStorage(), nil
	case "fs":
		dir, _ := data.AsString(conf.Params["dir"])
		var tempDir string
		if v, ok := conf.Params["temp_dir"]; ok {
			tempDir, _ = data.AsString(v)
		}
		return udsstorage.NewFS(dir, tempDir)
	default:
		return nil, fmt.Errorf("unsupported uds storage type: %v", conf.Type)
	}
}

func setUpTopology(logger *logrus.Logger, conf *config.Config, us udf.UDSStorage) (
	*bql.TopologyBuilder, error) {
	cc := &core.ContextConfig{
		Logger: logger,
	}
	cc.Flags.DroppedTupleLog.Set(conf.Logging.LogDroppedTuples)
	cc.Flags.DroppedTupleSummarization.Set(conf.Logging.SummarizeDroppedTuples)

	name := "runfile" // FIXME: bql file name is better?
	tp := core.NewDefaultTopology(core.NewContext(cc), name)
	tb, err := bql.NewTopologyBuilder(tp)
	if err != nil {
		return nil, fmt.Errorf("cannot create a new topology builder: %v", err)
	}
	tb.UDSStorage = us

	return tb, nil
}

func setUpBQLStmt(tb *bql.TopologyBuilder, bqlFile string) error {
	queries, err := func() (string, error) {
		f, err := os.Open(bqlFile)
		if err != nil {
			return "", err
		}
		defer f.Close()
		b, err := ioutil.ReadAll(f)
		if err != nil {
			return "", err
		}
		return string(b), nil
	}()
	if err != nil {
		return err
	}

	bp := parser.New()
	// TODO: provide better parse error reporting using ParseStmt instead of ParseStmts
	stmts, err := bp.ParseStmts(string(queries))
	if err != nil {
		return err
	}

	for _, stmt := range stmts {
		// TODO: if stmt is CREATE SOURCE, create it with PAUSED
		if n, err := tb.AddStmt(stmt); err != nil {
			tb.Topology().Context().ErrLog(err).WithField("stmt", stmt).Error(
				"Cannot add a statement to the topology")
			return err // FIXME: logger output "err" two twice
		} else if n != nil && n.Type() == core.NTSource {
			sn, _ := n.(core.SourceNode)
			if _, ok := sn.Source().(core.RewindableSource); ok {
				return fmt.Errorf(`rewindable source "%v" isn't supported`, n.Name())
			}
		}
	}
	return nil
}

func saveStates(tb *bql.TopologyBuilder, saveUDSs []string) error {
	states, err := tb.Topology().Context().SharedStates.List()
	if err != nil {
		return err
	}
	for _, name := range saveUDSs {
		state, ok := states[name]
		if !ok {
			err := fmt.Errorf("UDS '%v' is not found", name)
			tb.Topology().Context().ErrLog(err).WithField("uds", name).Error(
				"Cannot save the UDS")
			continue
		}
		target, ok := state.(core.SavableSharedState)
		if !ok {
			err := fmt.Errorf("UDS '%v' does not support save", name)
			tb.Topology().Context().ErrLog(err).WithField("uds", name).Error(
				"Cannot save the UDS")
			continue
		}
		// TODO get tag
		w, err := tb.UDSStorage.Save(tb.Topology().Name(), name, "default")
		if err != nil {
			tb.Topology().Context().ErrLog(err).WithField("uds", name).Error(
				"Cannot save the UDS")
			continue
		}
		defer w.Commit() // should be catch error
		// TODO get parameters
		err = target.Save(tb.Topology().Context(), w, data.Map{})
		if err != nil {
			tb.Topology().Context().ErrLog(err).WithField("uds", name).Error(
				"Cannot save the UDS")
			continue
		}
	}
	return nil
}
<commit_msg>change not to use StringSliceFlag, support to save all UDS<commit_after>/*
Package runfile implements sensorbee runfile command. This command is provided
separately from run command to reduce the footprint of runfile command itself.
For example, if sensorbee doesn't need to provide sensorbee run command and only
offers sensorbee runfile, the footprint of sensorbee command might be reduced a
lot due to low functionality of this command.
*/
package runfile

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"pfi/sensorbee/sensorbee/bql"
	"pfi/sensorbee/sensorbee/bql/parser"
	"pfi/sensorbee/sensorbee/bql/udf"
	"pfi/sensorbee/sensorbee/core"
	"pfi/sensorbee/sensorbee/data"
	"pfi/sensorbee/sensorbee/server/config"
	"pfi/sensorbee/sensorbee/server/udsstorage"
	"strings"
)

// SetUp sets up a command for running single BQL file.
func SetUp() cli.Command {
	cmd := cli.Command{
		Name:        "runfile",
		Usage:       "run a BQL file",
		Description: "runfile command runs a BQL file",
		Action:      Run,
	}

	cmd.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "config, c",
			Value:  "",
			Usage:  "file path of a config file in YAML format (only logging and storage sections are used)",
			EnvVar: "SENSORBEE_CONFIG",
		},
		cli.StringFlag{
			Name:  "save-uds, s",
			Value: "",
			Usage: "save UDSs after all tuples are processed",
		},
	}
	return cmd
}

// Run runs "runfile" command.
func Run(c *cli.Context) {
	// TODO: Merge this implementation with cmd/run
	if len(c.Args()) != 1 {
		cli.ShowSubcommandHelp(c)
		os.Exit(1)
	}

	conf, err := func() (*config.Config, error) {
		if c.IsSet("config") {
			p := c.String("config")
			in, err := ioutil.ReadFile(p)
			if err != nil {
				return nil, fmt.Errorf("Cannot read the config file %v: %v\n", p, err)
			}

			var yml map[string]interface{}
			if err := yaml.Unmarshal(in, &yml); err != nil {
				return nil, fmt.Errorf("Cannot parse the config file %v: %v\n", p, err)
			}
			m, err := data.NewMap(yml)
			if err != nil {
				return nil, fmt.Errorf("The config file %v has invalid values: %v\n", p, err)
			}
			return config.New(m)

		} else {
			// Currently there's no required parameters. However, when a required
			// parameter is added to Config, remove this else block and add a
			// default value like "/etc/sensorbee/config.yaml" to the config option.
			return config.New(data.Map{})
		}
	}()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	w, err := conf.Logging.CreateWriter()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	logger := logrus.New()
	logger.Out = w

	udsStorage, err := setUpUDSStorage(&conf.Storage.UDS)
	if err != nil {
		logger.WithField("err", err).Error("Cannot set up a storage for UDSs")
		os.Exit(1)
	}

	logger.Info("Setting up a topology")
	bqlFile := c.Args()[0]
	tb, err := setUpTopology(logger, conf, udsStorage)
	if err != nil {
		logger.WithField("err", err).Error("Cannot set up the topology")
		os.Exit(1)
	}

	if err := setUpBQLStmt(tb, bqlFile); err != nil {
		logger.WithFields(logrus.Fields{
			"err":      err,
			"bql_file": bqlFile,
		}).Error("Cannot set up BQL statement")
		os.Exit(1)
	}

	defer func() {
		logger.Info("Waiting for all nodes to finish processing tuples")
		if err := tb.Topology().Stop(); err != nil {
			logger.WithField("err", err).Error("Cannot stop the topology")
			os.Exit(1)
		}

		if c.IsSet("save-uds") {
			saveUDSList := c.String("save-uds")
			if err := saveStates(tb, saveUDSList); err != nil {
				logger.WithFields(logrus.Fields{
					"err":      err,
					"topology": tb.Topology().Name(),
				}).Error("Cannot save a UDS")
				os.Exit(1)
			}
		}
		// TODO: Terminate all shared states
	}()

	logger.Info("Starting the topology")
	for name, s := range tb.Topology().Sources() {
		if err := s.Resume(); err != nil {
			logger.WithFields(logrus.Fields{
				"err":    err,
				"source": name,
			}).Error("Cannot resume the source")
			os.Exit(1)
		}
	}

	for _, s := range tb.Topology().Sources() {
		// TODO: error check if necessary
		s.State().Wait(core.TSStopped)
	}
	logger.Info("All sources has been stopped.")

}

func setUpUDSStorage(conf *config.UDSStorage) (udf.UDSStorage, error) {
	// TODO: merge this implementation with server/context.go
	// Parameters are already validated in conf
	switch conf.Type {
	case "in_memory":
		return udf.NewInMemoryUDSStorage(), nil
	case "fs":
		dir, _ := data.AsString(conf.Params["dir"])
		var tempDir string
		if v, ok := conf.Params["temp_dir"]; ok {
			tempDir, _ = data.AsString(v)
		}
		return udsstorage.NewFS(dir, tempDir)
	default:
		return nil, fmt.Errorf("unsupported uds storage type: %v", conf.Type)
	}
}

func setUpTopology(logger *logrus.Logger, conf *config.Config, us udf.UDSStorage) (
	*bql.TopologyBuilder, error) {
	cc := &core.ContextConfig{
		Logger: logger,
	}
	cc.Flags.DroppedTupleLog.Set(conf.Logging.LogDroppedTuples)
	cc.Flags.DroppedTupleSummarization.Set(conf.Logging.SummarizeDroppedTuples)

	name := "runfile" // FIXME: bql file name is better?
	tp := core.NewDefaultTopology(core.NewContext(cc), name)
	tb, err := bql.NewTopologyBuilder(tp)
	if err != nil {
		return nil, fmt.Errorf("cannot create a new topology builder: %v", err)
	}
	tb.UDSStorage = us

	return tb, nil
}

func setUpBQLStmt(tb *bql.TopologyBuilder, bqlFile string) error {
	queries, err := func() (string, error) {
		f, err := os.Open(bqlFile)
		if err != nil {
			return "", err
		}
		defer f.Close()
		b, err := ioutil.ReadAll(f)
		if err != nil {
			return "", err
		}
		return string(b), nil
	}()
	if err != nil {
		return err
	}

	bp := parser.New()
	// TODO: provide better parse error reporting using ParseStmt instead of ParseStmts
	stmts, err := bp.ParseStmts(string(queries))
	if err != nil {
		return err
	}

	for _, stmt := range stmts {
		// TODO: if stmt is CREATE SOURCE, create it with PAUSED
		if n, err := tb.AddStmt(stmt); err != nil {
			tb.Topology().Context().ErrLog(err).WithField("stmt", stmt).Error(
				"Cannot add a statement to the topology")
			return err // FIXME: logger output "err" two twice
		} else if n != nil && n.Type() == core.NTSource {
			sn, _ := n.(core.SourceNode)
			if _, ok := sn.Source().(core.RewindableSource); ok {
				return fmt.Errorf(`rewindable source "%v" isn't supported`, n.Name())
			}
		}
	}
	return nil
}

func saveStates(tb *bql.TopologyBuilder, saveUDSList string) error {
	states, err := tb.Topology().Context().SharedStates.List()
	if err != nil {
		return err
	}
	// if save UDS list is empty then all UDS will be saved.
	saveUDSs := []string{}
	if saveUDSList == "" {
		for k := range states {
			saveUDSs = append(saveUDSs, k)
		}
	} else {
		saveUDSs = strings.Split(saveUDSList, ",")
	}
	for _, name := range saveUDSs {
		state, ok := states[name]
		if !ok {
			err := fmt.Errorf("UDS '%v' is not found", name)
			tb.Topology().Context().ErrLog(err).WithField("uds", name).Error(
				"Cannot save the UDS")
			continue
		}
		target, ok := state.(core.SavableSharedState)
		if !ok {
			err := fmt.Errorf("UDS '%v' does not support save", name)
			tb.Topology().Context().ErrLog(err).WithField("uds", name).Error(
				"Cannot save the UDS")
			continue
		}
		// TODO get tag
		w, err := tb.UDSStorage.Save(tb.Topology().Name(), name, "default")
		if err != nil {
			tb.Topology().Context().ErrLog(err).WithField("uds", name).Error(
				"Cannot save the UDS")
			continue
		}
		defer w.Commit() // should be catch error
		// TODO get parameters
		err = target.Save(tb.Topology().Context(), w, data.Map{})
		if err != nil {
			tb.Topology().Context().ErrLog(err).WithField("uds", name).Error(
				"Cannot save the UDS")
			continue
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strings"

	"github.com/sourcegraph/makex"
)

var file = flag.String("f", "Makefile", "path to Makefile")
var cwd = flag.String("C", "", "change to this directory before doing anything")
var dryRun = flag.Bool("n", false, "dry run (don't actually run any commands)")
var jobs = flag.Int("j", runtime.GOMAXPROCS(0), "number of jobs to run in parallel")
var expand = flag.Bool("x", true, "expand globs in makefile prereqs")

func main() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, `makex is an experimental, incomplete implementation of make in Go.

Usage:

        makex [options] [target] ...

If no targets are specified, the first target that appears in the makefile is
used.

The options are:
`)
		flag.PrintDefaults()
		os.Exit(1)
	}

	flag.Parse()
	conf := makex.Default
	conf.ParallelJobs = *jobs

	data, err := ioutil.ReadFile(*file)
	if err != nil {
		conf.Log.Fatal(err)
	}

	if *cwd != "" {
		err := os.Chdir(*cwd)
		if err != nil {
			conf.Log.Fatal(err)
		}
	}

	mf, err := makex.Parse(data)
	if err != nil {
		conf.Log.Fatal(err)
	}

	goals := flag.Args()
	if len(goals) == 0 {
		// Find the first rule that doesn't begin with a ".".
		for _, rule := range mf.Rules {
			target := rule.Target()
			if !strings.HasPrefix(target, ".") {
				goals = []string{target}
				break
			}
		}
	}

	if *expand {
		mf, err = conf.Expand(mf)
		if err != nil {
			log.Fatal(err)
		}
	}

	mk := conf.NewMaker(mf, goals...)

	targetSets, err := mk.TargetSetsNeedingBuild()
	if err != nil {
		conf.Log.Fatal(err)
	}

	if len(targetSets) == 0 {
		fmt.Println("Nothing to do.")
	}

	if *dryRun {
		for _, targetSet := range targetSets {
			fmt.Println(targetSet)
		}
		return
	}

	err = mk.Run()
	if err != nil {
		conf.Log.Fatal(err)
	}
}
<commit_msg>Nicer printing of dry run<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strings"

	"github.com/sourcegraph/makex"
)

var file = flag.String("f", "Makefile", "path to Makefile")
var cwd = flag.String("C", "", "change to this directory before doing anything")
var dryRun = flag.Bool("n", false, "dry run (don't actually run any commands)")
var jobs = flag.Int("j", runtime.GOMAXPROCS(0), "number of jobs to run in parallel")
var expand = flag.Bool("x", true, "expand globs in makefile prereqs")

func main() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, `makex is an experimental, incomplete implementation of make in Go.

Usage:

        makex [options] [target] ...

If no targets are specified, the first target that appears in the makefile is
used.

The options are:
`)
		flag.PrintDefaults()
		os.Exit(1)
	}

	flag.Parse()
	conf := makex.Default
	conf.ParallelJobs = *jobs

	data, err := ioutil.ReadFile(*file)
	if err != nil {
		conf.Log.Fatal(err)
	}

	if *cwd != "" {
		err := os.Chdir(*cwd)
		if err != nil {
			conf.Log.Fatal(err)
		}
	}

	mf, err := makex.Parse(data)
	if err != nil {
		conf.Log.Fatal(err)
	}

	goals := flag.Args()
	if len(goals) == 0 {
		// Find the first rule that doesn't begin with a ".".
		for _, rule := range mf.Rules {
			target := rule.Target()
			if !strings.HasPrefix(target, ".") {
				goals = []string{target}
				break
			}
		}
	}

	if *expand {
		mf, err = conf.Expand(mf)
		if err != nil {
			log.Fatal(err)
		}
	}

	mk := conf.NewMaker(mf, goals...)

	targetSets, err := mk.TargetSetsNeedingBuild()
	if err != nil {
		conf.Log.Fatal(err)
	}

	if len(targetSets) == 0 {
		fmt.Println("Nothing to do.")
	}

	if *dryRun {
		for i, targetSet := range targetSets {
			if i != 0 {
				fmt.Println()
			}
			fmt.Printf("========= TARGET SET %d (%d targets)\n", i, len(targetSet))
			for _, target := range targetSet {
				fmt.Println(" - ", target)
			}
		}
		return
	}

	err = mk.Run()
	if err != nil {
		conf.Log.Fatal(err)
	}
}
<|endoftext|>


<commit_before><commit_msg>"Jeff: Come to the dark side... together we will rule the galaxy."<commit_after><|endoftext|>