<commit_before>package widget

import (
	"bytes"
	"fmt"
	"html/template"
	"os"
	"path/filepath"
	"strings"

	"github.com/qor/qor/utils"
)

// Render find widget by name, render it based on current context
func (widgets *Widgets) Render(widgetName string, widgetGroupName string) template.HTML {
	return widgets.NewContext(nil).Render(widgetName, widgetGroupName)
}

// NewContext create new context for widgets
func (widgets *Widgets) NewContext(context *Context) *Context {
	if context == nil {
		context = &Context{}
	}

	if context.DB == nil {
		context.DB = widgets.Config.DB
	}

	if context.Options == nil {
		context.Options = map[string]interface{}{}
	}

	if context.FuncMaps == nil {
		context.FuncMaps = template.FuncMap{}
	}

	for key, fc := range widgets.funcMaps {
		if _, ok := context.FuncMaps[key]; !ok {
			context.FuncMaps[key] = fc
		}
	}

	context.Widgets = widgets
	return context
}

// Funcs return view functions map
func (context *Context) Funcs(funcMaps template.FuncMap) *Context {
	if context.FuncMaps == nil {
		context.FuncMaps = template.FuncMap{}
	}

	for key, fc := range funcMaps {
		context.FuncMaps[key] = fc
	}

	return context
}

// FuncMap return funcmap
func (context *Context) FuncMap() template.FuncMap {
	funcMap := template.FuncMap{}

	funcMap["render_widget"] = func(widgetName string, widgetGroupName ...string) template.HTML {
		var groupName string
		if len(widgetGroupName) == 0 {
			groupName = ""
		} else {
			groupName = widgetGroupName[0]
		}
		return context.Render(widgetName, groupName)
	}

	return funcMap
}

// Render register widget itself content
func (w *Widget) Render(context *Context, file string) template.HTML {
	var (
		err     error
		content []byte
		tmpl    *template.Template
	)

	if file == "" {
		file = w.Templates[0]
	}

	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("Get error when render file %v: %v", file, r)
			utils.ExitWithMsg(err)
		}
	}()

	if content, err = context.Widgets.AssetFS.Asset(file + ".tmpl"); err == nil {
		if tmpl, err = template.New(filepath.Base(file)).Funcs(context.FuncMaps).Parse(string(content)); err == nil {
			var result = bytes.NewBufferString("")
			if err = tmpl.Execute(result, context.Options); err == nil {
				return template.HTML(result.String())
			}
		}
	}

	return template.HTML(err.Error())
}

// RegisterViewPath register views directory
func (widgets *Widgets) RegisterViewPath(p string) {
	if filepath.IsAbs(p) {
		viewPaths = append(viewPaths, p)
		widgets.AssetFS.RegisterPath(p)
	} else {
		for _, gopath := range strings.Split(os.Getenv("GOPATH"), ":") {
			viewPaths = append(viewPaths, filepath.Join(gopath, "src", p))
			widgets.AssetFS.RegisterPath(filepath.Join(gopath, "src", p))
		}
	}
}

// LoadPreviewAssets will return assets tag used for Preview
func (widgets *Widgets) LoadPreviewAssets() template.HTML {
	tags := ""
	for _, asset := range widgets.Config.PreviewAssets {
		extension := filepath.Ext(asset)
		if extension == ".css" {
			tags += fmt.Sprintf("<link rel=\"stylesheet\" type=\"text/css\" href=\"%v\">\n", asset)
		} else if extension == ".js" {
			tags += fmt.Sprintf("<script src=\"%v\"></script>\n", asset)
		} else {
			tags += fmt.Sprintf("%v\n", asset)
		}
	}
	return template.HTML(tags)
}
<commit_msg>Use utils.GOPATH to get GOPATH<commit_after>package widget

import (
	"bytes"
	"fmt"
	"html/template"
	"path/filepath"

	"github.com/qor/qor/utils"
)

// Render find widget by name, render it based on current context
func (widgets *Widgets) Render(widgetName string, widgetGroupName string) template.HTML {
	return widgets.NewContext(nil).Render(widgetName, widgetGroupName)
}

// NewContext create new context for widgets
func (widgets *Widgets) NewContext(context *Context) *Context {
	if context == nil {
		context = &Context{}
	}

	if context.DB == nil {
		context.DB = widgets.Config.DB
	}

	if context.Options == nil {
		context.Options = map[string]interface{}{}
	}

	if context.FuncMaps == nil {
		context.FuncMaps = template.FuncMap{}
	}

	for key, fc := range widgets.funcMaps {
		if _, ok := context.FuncMaps[key]; !ok {
			context.FuncMaps[key] = fc
		}
	}

	context.Widgets = widgets
	return context
}

// Funcs return view functions map
func (context *Context) Funcs(funcMaps template.FuncMap) *Context {
	if context.FuncMaps == nil {
		context.FuncMaps = template.FuncMap{}
	}

	for key, fc := range funcMaps {
		context.FuncMaps[key] = fc
	}

	return context
}

// FuncMap return funcmap
func (context *Context) FuncMap() template.FuncMap {
	funcMap := template.FuncMap{}

	funcMap["render_widget"] = func(widgetName string, widgetGroupName ...string) template.HTML {
		var groupName string
		if len(widgetGroupName) == 0 {
			groupName = ""
		} else {
			groupName = widgetGroupName[0]
		}
		return context.Render(widgetName, groupName)
	}

	return funcMap
}

// Render register widget itself content
func (w *Widget) Render(context *Context, file string) template.HTML {
	var (
		err     error
		content []byte
		tmpl    *template.Template
	)

	if file == "" {
		file = w.Templates[0]
	}

	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("Get error when render file %v: %v", file, r)
			utils.ExitWithMsg(err)
		}
	}()

	if content, err = context.Widgets.AssetFS.Asset(file + ".tmpl"); err == nil {
		if tmpl, err = template.New(filepath.Base(file)).Funcs(context.FuncMaps).Parse(string(content)); err == nil {
			var result = bytes.NewBufferString("")
			if err = tmpl.Execute(result, context.Options); err == nil {
				return template.HTML(result.String())
			}
		}
	}

	return template.HTML(err.Error())
}

// RegisterViewPath register views directory
func (widgets *Widgets) RegisterViewPath(p string) {
	if filepath.IsAbs(p) {
		viewPaths = append(viewPaths, p)
		widgets.AssetFS.RegisterPath(p)
	} else {
		for _, gopath := range utils.GOPATH() {
			viewPaths = append(viewPaths, filepath.Join(gopath, "src", p))
			widgets.AssetFS.RegisterPath(filepath.Join(gopath, "src", p))
		}
	}
}

// LoadPreviewAssets will return assets tag used for Preview
func (widgets *Widgets) LoadPreviewAssets() template.HTML {
	tags := ""
	for _, asset := range widgets.Config.PreviewAssets {
		extension := filepath.Ext(asset)
		if extension == ".css" {
			tags += fmt.Sprintf("<link rel=\"stylesheet\" type=\"text/css\" href=\"%v\">\n", asset)
		} else if extension == ".js" {
			tags += fmt.Sprintf("<script src=\"%v\"></script>\n", asset)
		} else {
			tags += fmt.Sprintf("%v\n", asset)
		}
	}
	return template.HTML(tags)
}
<|endoftext|>


<commit_before><commit_msg>Add another exception for special lookups<commit_after><|endoftext|>


<commit_before>package storage

import (
	"bytes"
	"configuration"
	"fmt"
	"os"

	mdb "github.com/szferi/gomdb"
)

const MDB_NAME = "lmdb"

func init() {
	registerEngine(MDB_NAME, Initializer{
		NewMDBConfiguration,
		NewMDB,
	})
}

type MDBConfiguration struct {
	MapSize configuration.Size `toml:"map-size"`
}

type MDB struct {
	env  *mdb.Env
	db   mdb.DBI
	path string
}

func NewMDBConfiguration() interface{} {
	return &MDBConfiguration{}
}

func NewMDB(path string, config interface{}) (Engine, error) {
	c, ok := config.(*MDBConfiguration)
	if !ok {
		return nil, fmt.Errorf("Got config of type %T instead of %T", config, MDBConfiguration{})
	}

	if c.MapSize == 0 {
		c.MapSize = 10 * 1024 * 1024 * 1024
	}

	env, err := mdb.NewEnv()
	if err != nil {
		return MDB{}, err
	}

	// TODO: max dbs should be configurable
	if err := env.SetMaxDBs(1); err != nil {
		return MDB{}, err
	}
	if err := env.SetMapSize(uint64(c.MapSize)); err != nil {
		return MDB{}, err
	}

	if _, err := os.Stat(path); err != nil {
		err = os.MkdirAll(path, 0755)
		if err != nil {
			return MDB{}, err
		}
	}

	err = env.Open(path, mdb.WRITEMAP|mdb.MAPASYNC|mdb.CREATE, 0755)
	if err != nil {
		return MDB{}, err
	}

	tx, err := env.BeginTxn(nil, 0)
	if err != nil {
		return MDB{}, err
	}

	dbi, err := tx.DBIOpen(nil, mdb.CREATE)
	if err != nil {
		return MDB{}, err
	}

	if err := tx.Commit(); err != nil {
		return MDB{}, err
	}

	db := MDB{
		env:  env,
		db:   dbi,
		path: path,
	}

	return db, nil
}

func (db MDB) Put(key, value []byte) error {
	return db.BatchPut([]Write{Write{key, value}})
}

func (db MDB) BatchPut(writes []Write) error {
	tx, err := db.env.BeginTxn(nil, 0)
	if err != nil {
		return err
	}

	for _, w := range writes {
		var err error
		if w.Value == nil {
			err = tx.Del(db.db, w.Key, nil)
		} else {
			err = tx.Put(db.db, w.Key, w.Value, 0)
		}

		if err != nil && err != mdb.NotFound {
			tx.Abort()
			return err
		}
	}

	return tx.Commit()
}

func (db MDB) Get(key []byte) ([]byte, error) {
	tx, err := db.env.BeginTxn(nil, mdb.RDONLY)
	if err != nil {
		return nil, err
	}
	defer tx.Commit()

	v, err := tx.Get(db.db, key)
	if err == mdb.NotFound {
		return nil, nil
	}
	return v, err
}

func (db MDB) Del(start, finish []byte) error {
	tx, err := db.env.BeginTxn(nil, 0)
	if err != nil {
		return err
	}
	defer tx.Commit()

	itr := db.iterator(true)
	defer itr.Close()

	count := 0
	for itr.Seek(start); itr.Valid(); itr.Next() {
		key := itr.Key()
		if bytes.Compare(key, finish) > 0 {
			break
		}

		// TODO: We should be using one cursor instead of two
		// transactions, but deleting using a cursor, crashes
		err = tx.Del(db.db, itr.key, nil)
		if err != nil {
			return err
		}
		count++
	}
	return nil
}

type MDBIterator struct {
	key   []byte
	value []byte
	c     *mdb.Cursor
	tx    *mdb.Txn
	valid bool
	err   error
}

func (itr *MDBIterator) Key() []byte {
	return itr.key
}

func (itr *MDBIterator) Value() []byte {
	return itr.value
}

func (itr *MDBIterator) Valid() bool {
	return itr.valid
}

func (itr *MDBIterator) Error() error {
	return itr.err
}

func (itr *MDBIterator) getCurrent() {
	itr.key, itr.value, itr.err = itr.c.Get(nil, mdb.GET_CURRENT)
	itr.setState()
}

func (itr *MDBIterator) Seek(key []byte) {
	itr.key, itr.value, itr.err = itr.c.Get(key, mdb.SET_RANGE)
	itr.setState()
}
func (itr *MDBIterator) Next() {
	itr.key, itr.value, itr.err = itr.c.Get(nil, mdb.NEXT)
	itr.setState()
}

func (itr *MDBIterator) Prev() {
	itr.key, itr.value, itr.err = itr.c.Get(nil, mdb.PREV)
	itr.setState()
}

func (itr *MDBIterator) setState() {
	if itr.err != nil {
		if itr.err == mdb.NotFound {
			itr.err = nil
		}
		itr.valid = false
	}
}

func (itr *MDBIterator) Close() error {
	if err := itr.c.Close(); err != nil {
		itr.tx.Commit()
		return err
	}
	return itr.tx.Commit()
}

func (_ MDB) Name() string {
	return MDB_NAME
}

func (db MDB) Path() string {
	return db.path
}

func (db MDB) Iterator() Iterator {
	return db.iterator(true)
}

func (db MDB) Compact() {
}

func (db MDB) iterator(rdonly bool) *MDBIterator {
	flags := uint(0)
	if rdonly {
		flags = mdb.RDONLY
	}
	tx, err := db.env.BeginTxn(nil, flags)
	if err != nil {
		return &MDBIterator{nil, nil, nil, nil, false, err}
	}

	c, err := tx.CursorOpen(db.db)
	if err != nil {
		tx.Abort()
		return &MDBIterator{nil, nil, nil, nil, false, err}
	}

	return &MDBIterator{nil, nil, c, tx, true, nil}
}

func (db MDB) Close() {
	db.env.DBIClose(db.db)
	if err := db.env.Close(); err != nil {
		panic(err)
	}
}
<commit_msg>Fix delete<commit_after>package storage

import (
	"bytes"
	"configuration"
	"fmt"
	"os"

	mdb "github.com/szferi/gomdb"
)

const MDB_NAME = "lmdb"

func init() {
	registerEngine(MDB_NAME, Initializer{
		NewMDBConfiguration,
		NewMDB,
	})
}

type MDBConfiguration struct {
	MapSize configuration.Size `toml:"map-size"`
}

type MDB struct {
	env  *mdb.Env
	db   mdb.DBI
	path string
}

func NewMDBConfiguration() interface{} {
	return &MDBConfiguration{}
}

func NewMDB(path string, config interface{}) (Engine, error) {
	c, ok := config.(*MDBConfiguration)
	if !ok {
		return nil, fmt.Errorf("Got config of type %T instead of %T", config, MDBConfiguration{})
	}

	if c.MapSize == 0 {
		c.MapSize = 10 * 1024 * 1024 * 1024
	}

	env, err := mdb.NewEnv()
	if err != nil {
		return MDB{}, err
	}

	// TODO: max dbs should be configurable
	if err := env.SetMaxDBs(1); err != nil {
		return MDB{}, err
	}
	if err := env.SetMapSize(uint64(c.MapSize)); err != nil {
		return MDB{}, err
	}

	if _, err := os.Stat(path); err != nil {
		err = os.MkdirAll(path, 0755)
		if err != nil {
			return MDB{}, err
		}
	}

	err = env.Open(path, mdb.WRITEMAP|mdb.MAPASYNC|mdb.CREATE, 0755)
	if err != nil {
		return MDB{}, err
	}

	tx, err := env.BeginTxn(nil, 0)
	if err != nil {
		return MDB{}, err
	}

	dbi, err := tx.DBIOpen(nil, mdb.CREATE)
	if err != nil {
		return MDB{}, err
	}

	if err := tx.Commit(); err != nil {
		return MDB{}, err
	}

	db := MDB{
		env:  env,
		db:   dbi,
		path: path,
	}

	return db, nil
}

func (db MDB) Put(key, value []byte) error {
	return db.BatchPut([]Write{Write{key, value}})
}

func (db MDB) BatchPut(writes []Write) error {
	tx, err := db.env.BeginTxn(nil, 0)
	if err != nil {
		return err
	}

	for _, w := range writes {
		var err error
		if w.Value == nil {
			err = tx.Del(db.db, w.Key, nil)
		} else {
			err = tx.Put(db.db, w.Key, w.Value, 0)
		}

		if err != nil && err != mdb.NotFound {
			tx.Abort()
			return err
		}
	}

	return tx.Commit()
}

func (db MDB) Get(key []byte) ([]byte, error) {
	tx, err := db.env.BeginTxn(nil, mdb.RDONLY)
	if err != nil {
		return nil, err
	}
	defer tx.Commit()

	v, err := tx.Get(db.db, key)
	if err == mdb.NotFound {
		return nil, nil
	}
	return v, err
}

func (db MDB) Del(start, finish []byte) error {
	itr := db.iterator(false)
	defer itr.Close()

	count := 0
	for itr.Seek(start); itr.Valid(); itr.Next() {
		key := itr.Key()
		if bytes.Compare(key, finish) > 0 {
			break
		}

		err := itr.c.Del(0)
		if err != nil {
			return err
		}
		count++
	}
	return nil
}

type MDBIterator struct {
	key   []byte
	value []byte
	c     *mdb.Cursor
	tx    *mdb.Txn
	valid bool
	err   error
}

func (itr *MDBIterator) Key() []byte {
	return itr.key
}

func (itr *MDBIterator) Value() []byte {
	return itr.value
}

func (itr *MDBIterator) Valid() bool {
	return itr.valid
}

func (itr *MDBIterator) Error() error {
	return itr.err
}

func (itr *MDBIterator) getCurrent() {
	itr.key, itr.value, itr.err = itr.c.Get(nil, mdb.GET_CURRENT)
	itr.setState()
}

func (itr *MDBIterator) Seek(key []byte) {
	itr.key, itr.value, itr.err = itr.c.Get(key, mdb.SET_RANGE)
	itr.setState()
}
func (itr *MDBIterator) Next() {
	itr.key, itr.value, itr.err = itr.c.Get(nil, mdb.NEXT)
	itr.setState()
}

func (itr *MDBIterator) Prev() {
	itr.key, itr.value, itr.err = itr.c.Get(nil, mdb.PREV)
	itr.setState()
}

func (itr *MDBIterator) setState() {
	if itr.err != nil {
		if itr.err == mdb.NotFound {
			itr.err = nil
		}
		itr.valid = false
	}
}

func (itr *MDBIterator) Close() error {
	if err := itr.c.Close(); err != nil {
		itr.tx.Commit()
		return err
	}
	return itr.tx.Commit()
}

func (_ MDB) Name() string {
	return MDB_NAME
}

func (db MDB) Path() string {
	return db.path
}

func (db MDB) Iterator() Iterator {
	return db.iterator(true)
}

func (db MDB) Compact() {
}

func (db MDB) iterator(rdonly bool) *MDBIterator {
	flags := uint(0)
	if rdonly {
		flags = mdb.RDONLY
	}
	tx, err := db.env.BeginTxn(nil, flags)
	if err != nil {
		return &MDBIterator{nil, nil, nil, nil, false, err}
	}

	c, err := tx.CursorOpen(db.db)
	if err != nil {
		tx.Abort()
		return &MDBIterator{nil, nil, nil, nil, false, err}
	}

	return &MDBIterator{nil, nil, c, tx, true, nil}
}

func (db MDB) Close() {
	db.env.DBIClose(db.db)
	if err := db.env.Close(); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/cloudflare/redoctober/core"
)

// RemoteServer represents a remote RedOctober server.
type RemoteServer struct {
	client        *http.Client
	serverAddress string
}

// NewRemoteServer generates a RemoteServer with the server address and
// the root CA the server uses to authenticate itself.
func NewRemoteServer(serverAddress, CAFile string) (*RemoteServer, error) {

	// populate a root CA pool from file
	rootCAs := x509.NewCertPool()
	pemBytes, err := ioutil.ReadFile(CAFile)
	if err != nil {
		return nil, errors.New("fail to read CA file: " + err.Error())
	}
	ok := rootCAs.AppendCertsFromPEM(pemBytes)
	if !ok {
		return nil, errors.New("fail to populate CA root pool.")
	}

	tr := &http.Transport{
		TLSClientConfig:    &tls.Config{RootCAs: rootCAs},
		DisableCompression: true,
	}
	server := &RemoteServer{
		client:        &http.Client{Transport: tr},
		serverAddress: serverAddress,
	}
	return server, nil
}

// getURL creates URL for a specific path of the RemoteServer
func (c *RemoteServer) getURL(path string) string {
	return fmt.Sprintf("https://%s%s", c.serverAddress, path)

}

// doAction sends req to the remote server and returns the response
func (c *RemoteServer) doAction(action string, req []byte) ([]byte, error) {
	buf := bytes.NewBuffer(req)
	url := c.getURL("/" + action)
	resp, err := c.client.Post(url, "application/json", buf)
	if err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	resp.Body.Close()

	return body, nil

}

// unmarshalResponseData is a helper function that unmarshal response bytes
// into ResponseData object.
func unmarshalResponseData(respBytes []byte) (*core.ResponseData, error) {
	response := new(core.ResponseData)
	err := json.Unmarshal(respBytes, response)
	if err != nil {
		return nil, err
	}

	if response.Status != "ok" {
		return nil, errors.New(response.Status)
	}

	return response, nil
}

// Create creates an admin account at the remote server
func (c *RemoteServer) Create(req core.CreateRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("create", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}

// Summary returns the summary reported by the remote server
func (c *RemoteServer) Summary(req core.SummaryRequest) (*core.SummaryData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("summary", reqBytes)
	if err != nil {
		return nil, err
	}

	response := new(core.SummaryData)
	err = json.Unmarshal(respBytes, response)
	if err != nil {
		return nil, err
	}

	if response.Status != "ok" {
		return nil, errors.New(response.Status)
	}
	return response, nil
}

// Delegate issues a delegate request to the remote server
func (c *RemoteServer) Delegate(req core.DelegateRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("delegate", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}

// Modify issues a modify request to the remote server
func (c *RemoteServer) Modify(req core.ModifyRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("modify", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}

// Encrypt issues an encrypt request to the remote server
func (c *RemoteServer) Encrypt(req core.EncryptRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("encrypt", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}

// Decrypt issues an decrypt request to the remote server
func (c *RemoteServer) Decrypt(req core.DecryptRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("decrypt", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)

}

// DecryptIntoData issues an decrypt request to the remote server and extract
// the decrypted data from the response
func (c *RemoteServer) DecryptIntoData(req core.DecryptRequest) ([]byte, error) {
	responseData, err := c.Decrypt(req)
	if err != nil {
		return nil, err
	}

	d := new(core.DecryptWithDelegates)
	err = json.Unmarshal(responseData.Response, d)
	if err != nil {
		return nil, err
	}

	return d.Data, nil

}

// Password issues an password request to the remote server
func (c *RemoteServer) Password(req []byte) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("delegate", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}
<commit_msg>client update. Trusted CAFile should be optional<commit_after>package client

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/cloudflare/redoctober/core"
)

// RemoteServer represents a remote RedOctober server.
type RemoteServer struct {
	client        *http.Client
	serverAddress string
}

// NewRemoteServer generates a RemoteServer with the server address and
// the root CA the server uses to authenticate itself.
func NewRemoteServer(serverAddress, CAFile string) (*RemoteServer, error) {

	var rootCAs *x509.CertPool
	// populate a root CA pool from input CAfile
	// otherwise, use the system's default root CA set
	if CAFile != "" {
		rootCAs = x509.NewCertPool()
		pemBytes, err := ioutil.ReadFile(CAFile)
		if err != nil {
			return nil, errors.New("fail to read CA file: " + err.Error())
		}
		ok := rootCAs.AppendCertsFromPEM(pemBytes)
		if !ok {
			return nil, errors.New("fail to populate CA root pool.")
		}
	}

	tr := &http.Transport{
		TLSClientConfig:    &tls.Config{RootCAs: rootCAs},
		DisableCompression: true,
	}
	server := &RemoteServer{
		client:        &http.Client{Transport: tr},
		serverAddress: serverAddress,
	}
	return server, nil
}

// getURL creates URL for a specific path of the RemoteServer
func (c *RemoteServer) getURL(path string) string {
	return fmt.Sprintf("https://%s%s", c.serverAddress, path)

}

// doAction sends req to the remote server and returns the response
func (c *RemoteServer) doAction(action string, req []byte) ([]byte, error) {
	buf := bytes.NewBuffer(req)
	url := c.getURL("/" + action)
	resp, err := c.client.Post(url, "application/json", buf)
	if err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	resp.Body.Close()

	return body, nil

}

// unmarshalResponseData is a helper function that unmarshal response bytes
// into ResponseData object.
func unmarshalResponseData(respBytes []byte) (*core.ResponseData, error) {
	response := new(core.ResponseData)
	err := json.Unmarshal(respBytes, response)
	if err != nil {
		return nil, err
	}

	if response.Status != "ok" {
		return nil, errors.New(response.Status)
	}

	return response, nil
}

// Create creates an admin account at the remote server
func (c *RemoteServer) Create(req core.CreateRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("create", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}

// Summary returns the summary reported by the remote server
func (c *RemoteServer) Summary(req core.SummaryRequest) (*core.SummaryData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("summary", reqBytes)
	if err != nil {
		return nil, err
	}

	response := new(core.SummaryData)
	err = json.Unmarshal(respBytes, response)
	if err != nil {
		return nil, err
	}

	if response.Status != "ok" {
		return nil, errors.New(response.Status)
	}
	return response, nil
}

// Delegate issues a delegate request to the remote server
func (c *RemoteServer) Delegate(req core.DelegateRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("delegate", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}

// Modify issues a modify request to the remote server
func (c *RemoteServer) Modify(req core.ModifyRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("modify", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}

// Encrypt issues an encrypt request to the remote server
func (c *RemoteServer) Encrypt(req core.EncryptRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("encrypt", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}

// Decrypt issues an decrypt request to the remote server
func (c *RemoteServer) Decrypt(req core.DecryptRequest) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("decrypt", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)

}

// DecryptIntoData issues an decrypt request to the remote server and extract
// the decrypted data from the response
func (c *RemoteServer) DecryptIntoData(req core.DecryptRequest) ([]byte, error) {
	responseData, err := c.Decrypt(req)
	if err != nil {
		return nil, err
	}

	d := new(core.DecryptWithDelegates)
	err = json.Unmarshal(responseData.Response, d)
	if err != nil {
		return nil, err
	}

	return d.Data, nil

}

// Password issues an password request to the remote server
func (c *RemoteServer) Password(req []byte) (*core.ResponseData, error) {
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := c.doAction("delegate", reqBytes)
	if err != nil {
		return nil, err
	}

	return unmarshalResponseData(respBytes)
}
<|endoftext|>


<commit_before>package engine

import (
	"testing"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/protocol/go"
)

func verifyPassphraseChange(tc libkb.TestContext, u *FakeUser, newPassphrase string) {
	_, err := tc.G.LoginState().VerifyPlaintextPassphrase(newPassphrase)
	if err != nil {
		tc.T.Fatal(err)
	}

	_, err = tc.G.LoginState().VerifyPlaintextPassphrase(u.Passphrase)
	if err == nil {
		tc.T.Fatal("old passphrase passed verification")
	}
}

func assertLoadSecretKeys(tc libkb.TestContext, u *FakeUser, msg string) {
	me, err := libkb.LoadMe(libkb.LoadUserArg{})
	if err != nil {
		tc.T.Fatalf("%s: %s", msg, err)
	}
	skarg := libkb.SecretKeyArg{
		Me:      me,
		KeyType: libkb.DeviceSigningKeyType,
	}
	sigKey, _, err := tc.G.Keyrings.GetSecretKeyWithPrompt(nil, skarg, u.NewSecretUI(), "testing sig")
	if err != nil {
		tc.T.Fatalf("%s: %s", msg, err)
	}
	if sigKey == nil {
		tc.T.Fatalf("%s: got nil signing key", msg)
	}

	skarg.KeyType = libkb.DeviceEncryptionKeyType
	encKey, _, err := tc.G.Keyrings.GetSecretKeyWithPrompt(nil, skarg, u.NewSecretUI(), "testing enc")
	if err != nil {
		tc.T.Fatalf("%s: %s", msg, err)
	}
	if encKey == nil {
		tc.T.Fatalf("%s: got nil encryption key", msg)
	}
}

// Test changing the passphrase when user knows current
// passphrase.
func TestPassphraseChangeKnown(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")
	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		OldPassphrase: u.Passphrase,
		Passphrase:    newPassphrase,
	}

	// using an empty secret ui to make sure existing pp doesn't come from ui prompt:
	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change known")
}

// Test changing the passphrase when user knows current
// passphrase, prompt for it.
func TestPassphraseChangeKnownPrompt(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	// clear the passphrase stream cache to force a prompt
	// for the existing passphrase.
	tc.G.LoginState().Account(func(a *libkb.Account) {
		a.ClearStreamCache()
	}, "clear stream cache")

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
	}
	secui := u.NewSecretUI()
	ctx := &Context{
		SecretUI: secui,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	if !secui.CalledGetKBPassphrase {
		t.Errorf("get kb passphrase not called")
	}

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change known prompt")
}

// Test changing the passphrase after logging in via pubkey.
func TestPassphraseChangeAfterPubkeyLogin(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	// this should do a pubkey login
	Logout(tc)

	secui := u.NewSecretUI()
	u.LoginWithSecretUI(secui, tc.G)
	if !secui.CalledGetSecret {
		t.Errorf("get secret not called")
	}

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
	}
	ctx := &Context{
		SecretUI: secui,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change after pubkey login")
}

// Test changing the passphrase when previous pp stream available.
func TestPassphraseChangeKnownNotSupplied(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")
	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
	}
	secui := &libkb.TestSecretUI{}
	ctx := &Context{
		SecretUI: secui,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	if secui.CalledGetKBPassphrase {
		t.Errorf("get kb passphrase called")
	}

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change known, not supplied")
}

// Test changing the passphrase when user forgets current
// passphrase.
func TestPassphraseChangeUnknown(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	// this has a flaw:  the passphrase stream cache is available.
	// it is being used to unlock the secret key to generate the
	// change passphrase proof.
	//

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change unknown")
}

// Test changing the passphrase when user forgets current
// passphrase and there's no passphrase stream cache.
// No backup key available.
func TestPassphraseChangeUnknownNoPSCache(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	f := func(arg *SignupEngineRunArg) {
		arg.SkipPaper = true
	}

	u, _ := CreateAndSignupFakeUserCustomArg(tc, "paper", f)

	tc.G.LoginState().Account(func(a *libkb.Account) {
		a.ClearStreamCache()
	}, "clear stream cache")

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewPassphraseChange(arg, tc.G)
	err := RunEngine(eng, ctx)
	if err == nil {
		t.Fatal("passphrase change should have failed")
	}
	if _, ok := err.(libkb.NoPaperKeysError); !ok {
		t.Fatalf("unexpected error: %s (%T)", err, err)
	}

	assertLoadSecretKeys(tc, u, "passphrase change unknown, no ps cache")
}

// Test changing the passphrase when user forgets current
// passphrase and there's no passphrase stream cache.
// Backup key exists
func TestPassphraseChangeUnknownBackupKey(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		LoginUI:  libkb.TestLoginUI{},
		SecretUI: &libkb.TestSecretUI{},
	}
	beng := NewPaperKey(tc.G)
	if err := RunEngine(beng, ctx); err != nil {
		t.Fatal(err)
	}
	backupPassphrase := beng.Passphrase()
	ctx.SecretUI = &libkb.TestSecretUI{BackupPassphrase: backupPassphrase}

	tc.G.LoginState().Account(func(a *libkb.Account) {
		a.ClearStreamCache()
	}, "clear stream cache")

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change unknown, backup key")
}

// Test changing the passphrase when user forgets current
// passphrase and is logged out, but has a backup key.
func TestPassphraseChangeLoggedOutBackupKey(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	assertLoadSecretKeys(tc, u, "logged out w/ backup key, before passphrase change")

	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		LoginUI:  libkb.TestLoginUI{},
		SecretUI: &libkb.TestSecretUI{},
	}
	beng := NewPaperKey(tc.G)
	if err := RunEngine(beng, ctx); err != nil {
		t.Fatal(err)
	}
	backupPassphrase := beng.Passphrase()
	ctx.SecretUI = &libkb.TestSecretUI{BackupPassphrase: backupPassphrase}

	Logout(tc)

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "logged out w/ backup key, after passphrase change")
}

// Test changing the passphrase when user forgets current passphrase
// and is logged out, but has a backup key (generated by a secret from
// the secret store).
func TestPassphraseChangeLoggedOutBackupKeySecretStore(t *testing.T) {
	// TODO: Get this working on non-OS X platforms (by mocking
	// out the SecretStore).
	if !libkb.HasSecretStore() {
		t.Skip("Skipping test since there is no secret store")
	}

	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := NewFakeUserOrBust(tc.T, "login")
	signupArg := MakeTestSignupEngineRunArg(u)
	signupArg.StoreSecret = true
	_ = SignupFakeUserWithArg(tc, u, signupArg)

	// This needs to happen *before* resetting the login state, as
	// this call will cause the login state to be reloaded.
	assertLoadSecretKeys(tc, u, "logged out w/ backup key, before passphrase change")

	tc.ResetLoginState()

	secretUI := libkb.TestSecretUI{}
	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		LoginUI:  libkb.TestLoginUI{},
		SecretUI: &secretUI,
	}
	beng := NewBackup(tc.G)
	if err := RunEngine(beng, ctx); err != nil {
		t.Fatal(err)
	}

	if secretUI.CalledGetSecret {
		t.Fatal("GetSecret() unexpectedly called")
	}

	backupPassphrase := beng.Passphrase()
	ctx.SecretUI = &libkb.TestSecretUI{BackupPassphrase: backupPassphrase}

	Logout(tc)

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "logged out w/ backup key, after passphrase change")
}
<commit_msg>fixed for new func name<commit_after>package engine

import (
	"testing"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/protocol/go"
)

func verifyPassphraseChange(tc libkb.TestContext, u *FakeUser, newPassphrase string) {
	_, err := tc.G.LoginState().VerifyPlaintextPassphrase(newPassphrase)
	if err != nil {
		tc.T.Fatal(err)
	}

	_, err = tc.G.LoginState().VerifyPlaintextPassphrase(u.Passphrase)
	if err == nil {
		tc.T.Fatal("old passphrase passed verification")
	}
}

func assertLoadSecretKeys(tc libkb.TestContext, u *FakeUser, msg string) {
	me, err := libkb.LoadMe(libkb.LoadUserArg{})
	if err != nil {
		tc.T.Fatalf("%s: %s", msg, err)
	}
	skarg := libkb.SecretKeyArg{
		Me:      me,
		KeyType: libkb.DeviceSigningKeyType,
	}
	sigKey, _, err := tc.G.Keyrings.GetSecretKeyWithPrompt(nil, skarg, u.NewSecretUI(), "testing sig")
	if err != nil {
		tc.T.Fatalf("%s: %s", msg, err)
	}
	if sigKey == nil {
		tc.T.Fatalf("%s: got nil signing key", msg)
	}

	skarg.KeyType = libkb.DeviceEncryptionKeyType
	encKey, _, err := tc.G.Keyrings.GetSecretKeyWithPrompt(nil, skarg, u.NewSecretUI(), "testing enc")
	if err != nil {
		tc.T.Fatalf("%s: %s", msg, err)
	}
	if encKey == nil {
		tc.T.Fatalf("%s: got nil encryption key", msg)
	}
}

// Test changing the passphrase when user knows current
// passphrase.
func TestPassphraseChangeKnown(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")
	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		OldPassphrase: u.Passphrase,
		Passphrase:    newPassphrase,
	}

	// using an empty secret ui to make sure existing pp doesn't come from ui prompt:
	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change known")
}

// Test changing the passphrase when user knows current
// passphrase, prompt for it.
func TestPassphraseChangeKnownPrompt(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	// clear the passphrase stream cache to force a prompt
	// for the existing passphrase.
	tc.G.LoginState().Account(func(a *libkb.Account) {
		a.ClearStreamCache()
	}, "clear stream cache")

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
	}
	secui := u.NewSecretUI()
	ctx := &Context{
		SecretUI: secui,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	if !secui.CalledGetKBPassphrase {
		t.Errorf("get kb passphrase not called")
	}

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change known prompt")
}

// Test changing the passphrase after logging in via pubkey.
func TestPassphraseChangeAfterPubkeyLogin(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	// this should do a pubkey login
	Logout(tc)

	secui := u.NewSecretUI()
	u.LoginWithSecretUI(secui, tc.G)
	if !secui.CalledGetSecret {
		t.Errorf("get secret not called")
	}

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
	}
	ctx := &Context{
		SecretUI: secui,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change after pubkey login")
}

// Test changing the passphrase when previous pp stream available.
func TestPassphraseChangeKnownNotSupplied(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")
	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
	}
	secui := &libkb.TestSecretUI{}
	ctx := &Context{
		SecretUI: secui,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	if secui.CalledGetKBPassphrase {
		t.Errorf("get kb passphrase called")
	}

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change known, not supplied")
}

// Test changing the passphrase when user forgets current
// passphrase.
func TestPassphraseChangeUnknown(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	// this has a flaw:  the passphrase stream cache is available.
	// it is being used to unlock the secret key to generate the
	// change passphrase proof.
	//

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change unknown")
}

// Test changing the passphrase when user forgets current
// passphrase and there's no passphrase stream cache.
// No backup key available.
func TestPassphraseChangeUnknownNoPSCache(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	f := func(arg *SignupEngineRunArg) {
		arg.SkipPaper = true
	}

	u, _ := CreateAndSignupFakeUserCustomArg(tc, "paper", f)

	tc.G.LoginState().Account(func(a *libkb.Account) {
		a.ClearStreamCache()
	}, "clear stream cache")

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewPassphraseChange(arg, tc.G)
	err := RunEngine(eng, ctx)
	if err == nil {
		t.Fatal("passphrase change should have failed")
	}
	if _, ok := err.(libkb.NoPaperKeysError); !ok {
		t.Fatalf("unexpected error: %s (%T)", err, err)
	}

	assertLoadSecretKeys(tc, u, "passphrase change unknown, no ps cache")
}

// Test changing the passphrase when user forgets current
// passphrase and there's no passphrase stream cache.
// Backup key exists
func TestPassphraseChangeUnknownBackupKey(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		LoginUI:  libkb.TestLoginUI{},
		SecretUI: &libkb.TestSecretUI{},
	}
	beng := NewPaperKey(tc.G)
	if err := RunEngine(beng, ctx); err != nil {
		t.Fatal(err)
	}
	backupPassphrase := beng.Passphrase()
	ctx.SecretUI = &libkb.TestSecretUI{BackupPassphrase: backupPassphrase}

	tc.G.LoginState().Account(func(a *libkb.Account) {
		a.ClearStreamCache()
	}, "clear stream cache")

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "passphrase change unknown, backup key")
}

// Test changing the passphrase when user forgets current
// passphrase and is logged out, but has a backup key.
func TestPassphraseChangeLoggedOutBackupKey(t *testing.T) {
	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := CreateAndSignupFakeUser(tc, "login")

	assertLoadSecretKeys(tc, u, "logged out w/ backup key, before passphrase change")

	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		LoginUI:  libkb.TestLoginUI{},
		SecretUI: &libkb.TestSecretUI{},
	}
	beng := NewPaperKey(tc.G)
	if err := RunEngine(beng, ctx); err != nil {
		t.Fatal(err)
	}
	backupPassphrase := beng.Passphrase()
	ctx.SecretUI = &libkb.TestSecretUI{BackupPassphrase: backupPassphrase}

	Logout(tc)

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "logged out w/ backup key, after passphrase change")
}

// Test changing the passphrase when user forgets current passphrase
// and is logged out, but has a backup key (generated by a secret from
// the secret store).
func TestPassphraseChangeLoggedOutBackupKeySecretStore(t *testing.T) {
	// TODO: Get this working on non-OS X platforms (by mocking
	// out the SecretStore).
	if !libkb.HasSecretStore() {
		t.Skip("Skipping test since there is no secret store")
	}

	tc := SetupEngineTest(t, "PassphraseChange")
	defer tc.Cleanup()

	u := NewFakeUserOrBust(tc.T, "login")
	signupArg := MakeTestSignupEngineRunArg(u)
	signupArg.StoreSecret = true
	_ = SignupFakeUserWithArg(tc, u, signupArg)

	// This needs to happen *before* resetting the login state, as
	// this call will cause the login state to be reloaded.
	assertLoadSecretKeys(tc, u, "logged out w/ backup key, before passphrase change")

	tc.ResetLoginState()

	secretUI := libkb.TestSecretUI{}
	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		LoginUI:  libkb.TestLoginUI{},
		SecretUI: &secretUI,
	}
	beng := NewPaperKey(tc.G)
	if err := RunEngine(beng, ctx); err != nil {
		t.Fatal(err)
	}

	if secretUI.CalledGetSecret {
		t.Fatal("GetSecret() unexpectedly called")
	}

	backupPassphrase := beng.Passphrase()
	ctx.SecretUI = &libkb.TestSecretUI{BackupPassphrase: backupPassphrase}

	Logout(tc)

	newPassphrase := "password"
	arg := &keybase1.PassphraseChangeArg{
		Passphrase: newPassphrase,
		Force:      true,
	}
	eng := NewPassphraseChange(arg, tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	verifyPassphraseChange(tc, u, newPassphrase)

	u.Passphrase = newPassphrase
	assertLoadSecretKeys(tc, u, "logged out w/ backup key, after passphrase change")
}
<|endoftext|>


<commit_before>package cache

import (
	"sync"
	"time"

	"github.com/fatih/set"
)

var (
	socketSubscriptionsMap      = make(map[string]*SubscriptionSet)
	socketSubscriptionsMapMutex sync.Mutex
)

type SubscriptionSet struct {
	set      *set.Set
	socketID string
}

func NewSubscriptionSet(socketID string) (*SubscriptionSet, error) {
	s := &SubscriptionSet{
		set:      set.New(),
		socketID: socketID,
	}
	socketSubscriptionsMapMutex.Lock()
	defer socketSubscriptionsMapMutex.Unlock()

	socketSubscriptionsMap[socketID] = s
	return s, nil
}

func (s *SubscriptionSet) Each(f func(item interface{}) bool) error {
	s.set.Each(f)
	// each doesnt return anything
	return nil
}

func (s *SubscriptionSet) Subscribe(routingKeyPrefix string) error {
	s.set.Add(routingKeyPrefix)
	// add doesnt return any error
	return nil
}

func (s *SubscriptionSet) Resubscribe(socketID string) (bool, error) {
	socketSubscription, ok := socketSubscriptionsMap[socketID]
	if !ok {
		return false, nil
	}

	socketSubscription.Each(func(routingKeyPrefix interface{}) bool {
		if err := s.Subscribe(routingKeyPrefix.(string)); err != nil {
			return false
		}
		return true
	})
	return true, nil
}

func (s *SubscriptionSet) Unsubscribe(routingKeyPrefix string) error {
	s.set.Remove(routingKeyPrefix)
	// remove doesnt return any error
	return nil
}

func (s *SubscriptionSet) Has(routingKeyPrefix string) (bool, error) {
	// has only returns bool
	return s.set.Has(routingKeyPrefix), nil
}

func (s *SubscriptionSet) Len() (int, error) {
	// size only returns count
	return s.set.Size(), nil
}

func (s *SubscriptionSet) ClearWithTimeout() error {
	time.AfterFunc(5*time.Minute, func() {
		socketSubscriptionsMapMutex.Lock()
		delete(socketSubscriptionsMap, s.socketID)
		socketSubscriptionsMapMutex.Unlock()
	})
	return nil
}
<commit_msg>Broker/Cache: use variadic function call for Subscsribe<commit_after>package cache

import (
	"sync"
	"time"

	"github.com/fatih/set"
)

var (
	socketSubscriptionsMap      = make(map[string]*SubscriptionSet)
	socketSubscriptionsMapMutex sync.Mutex
)

type SubscriptionSet struct {
	set      *set.Set
	socketID string
}

func NewSubscriptionSet(socketID string) (*SubscriptionSet, error) {
	s := &SubscriptionSet{
		set:      set.New(),
		socketID: socketID,
	}
	socketSubscriptionsMapMutex.Lock()
	defer socketSubscriptionsMapMutex.Unlock()

	socketSubscriptionsMap[socketID] = s
	return s, nil
}

func (s *SubscriptionSet) Each(f func(item interface{}) bool) error {
	s.set.Each(f)
	// each doesnt return anything
	return nil
}

func (s *SubscriptionSet) Subscribe(routingKeyPrefixes ...string) error {
	for _, routingKeyPrefix := range routingKeyPrefixes {
		s.set.Add(routingKeyPrefix)
	}
	// add doesnt return any error
	return nil
}

func (s *SubscriptionSet) Resubscribe(socketID string) (bool, error) {
	socketSubscription, ok := socketSubscriptionsMap[socketID]
	if !ok {
		return false, nil
	}

	socketSubscription.Each(func(routingKeyPrefix interface{}) bool {
		if err := s.Subscribe(routingKeyPrefix.(string)); err != nil {
			return false
		}
		return true
	})
	return true, nil
}

func (s *SubscriptionSet) Unsubscribe(routingKeyPrefix string) error {
	s.set.Remove(routingKeyPrefix)
	// remove doesnt return any error
	return nil
}

func (s *SubscriptionSet) Has(routingKeyPrefix string) (bool, error) {
	// has only returns bool
	return s.set.Has(routingKeyPrefix), nil
}

func (s *SubscriptionSet) Len() (int, error) {
	// size only returns count
	return s.set.Size(), nil
}

func (s *SubscriptionSet) ClearWithTimeout() error {
	time.AfterFunc(5*time.Minute, func() {
		socketSubscriptionsMapMutex.Lock()
		delete(socketSubscriptionsMap, s.socketID)
		socketSubscriptionsMapMutex.Unlock()
	})
	return nil
}
<|endoftext|>


<commit_before>package main

import (
    "os"
    "fmt"
    "log"
    "bytes"
    "strings"
    "path/filepath"
    //"path"
    "encoding/gob"
    "net"
    "mob/proto"
    //"os/signal"
    //"syscall"
    //"mob/client/music"
    //"github.com/tcolgate/mp3"
)

func main() {
    fmt.Println("Starting client ...")
    conn, err := net.Dial("tcp", os.Args[1])
    if err != nil {
        // handle error
    }

    enc := gob.NewEncoder(conn) // Will write to network.
    dec := gob.NewDecoder(conn) // Will read from network.

    // Encode (send) some values.
    enc.Encode(proto.Client_Init_Packet{getSongNames()})

    for {
        var frame_packet proto.Mp3_Frame_Packet

        dec.Decode(&frame_packet)

        fmt.Println(frame_packet.Seqnum)
        fmt.Println(frame_packet.Mp3_frame.Size())
    }

    // Shell commands:
    // list
    // play
    // exit
    // join

    //for {
        // get stdin input
        // switch on string
    //    break
    //}
}

/*
func handleJoin() {

}

func handleList() {

}

func handlePlay() {

}

func handleExit() {

}*/

// Returns csv of all song names in the songs folder.
func getSongNames() string {
    var buffer bytes.Buffer
    filepath.Walk("../songs", func (p string, i os.FileInfo, err error) error {
        if err != nil {
            log.Println(err)
            return nil
        }

        s := filepath.Base(p)
        if strings.Compare(s, "songs") != 0 {
            buffer.WriteString(s + ";")
        }

        return nil
    })

    return buffer.String()
}
<commit_msg>Adding RTT calc haven't tested<commit_after>package main

import (
    "os"
    "fmt"
    "log"
    "bytes"
    "strings"
    "path/filepath"
    //"path"
    "encoding/gob"
    "net"
    "mob/proto"
    "github.com/mikioh/tcp"
    "github.com/mikioh/tcpinfo"
    "time"
    //"os/signal"
    //"syscall"
    //"mob/client/music"
    //"github.com/tcolgate/mp3"
)

func main() {
    fmt.Println("Starting client ...")
    conn, err := net.Dial("tcp", os.Args[1])
    if err != nil {
        // handle error
    }

    enc := gob.NewEncoder(conn) // Will write to network.
    dec := gob.NewDecoder(conn) // Will read from network.

    // Encode (send) some values.
    enc.Encode(proto.Client_Init_Packet{getSongNames()})

    for {
        var frame_packet proto.Mp3_Frame_Packet

        dec.Decode(&frame_packet)

        fmt.Println(frame_packet.Seqnum)
        //fmt.Println(frame_packet.Mp3_frame.Size())
    }

    // Shell commands:
    // list
    // play
    // exit
    // join

    //for {
        // get stdin input
        // switch on string
    //    break
    //}
}

/*
func handleJoin() {

}

func handleList() {

}

func handlePlay() {

}

func handleExit() {

}*/

// Returns csv of all song names in the songs folder.
func getSongNames() string {
    var buffer bytes.Buffer
    filepath.Walk("../songs", func (p string, i os.FileInfo, err error) error {
        if err != nil {
            log.Println(err)
            return nil
        }

        s := filepath.Base(p)
        if strings.Compare(s, "songs") != 0 {
            buffer.WriteString(s + ";")
        }

        return nil
    })

    return buffer.String()
}

// get RTT in terms of milliseconds between current node and specified IP
func getRTTBetweenNodes(address string) float32 {
    c, err := net.Dial("tcp", address)

    if err != nil {
	       // error handling
    }
    defer c.Close()

    tc, err := tcp.NewConn(c)
    if err != nil {
    	// error handling
    }
    var o tcpinfo.Info
    var b [256]byte
    i, err := tc.Option(o.Level(), o.Name(), b[:])
    if err != nil {
    	// error handling
    }
    txt, err := json.Marshal(i)
    if err != nil {
    	// error handling
    }

    return (time.Seconds(txt.rtt) / 1000)
}
<|endoftext|>


<commit_before>package fuseklient

import (
	"fmt"
	"math/rand"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
	"golang.org/x/net/context"
	"golang.org/x/net/trace"
)

type TraceFS struct {
	Id string
	*KodingNetworkFS
}

func NewTraceFS(k *KodingNetworkFS) *TraceFS {
	return &TraceFS{
		Id:              randSeq(6),
		KodingNetworkFS: k,
	}
}

func (t *TraceFS) Mount() (*fuse.MountedFileSystem, error) {
	r := t.newTrace("Mount", "Path=%s", t.MountPath)
	defer r.Finish()

	server := fuseutil.NewFileSystemServer(t)
	fs, err := fuse.Mount(t.MountPath, server, t.MountConfig)
	if err != nil {
		r.LazyPrintf("KodingNetworkFS#Mount: err '%s'", err)
		r.SetError()

		return nil, err
	}

	return fs, nil
}

func (t *TraceFS) Unmount() error {
	r := t.newTrace("Unmount", "Path=%s", t.MountPath)
	defer r.Finish()

	if err := t.KodingNetworkFS.Unmount(); err != nil {
		r.LazyPrintf("KodingNetworkFS#Unmount err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) GetInodeAttributes(ctx context.Context, op *fuseops.GetInodeAttributesOp) error {
	r := t.newTrace("GetInodeAttributes", "ID=%d", op.Inode)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.GetInodeAttributes(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#GetInodeAttributes err:%s", err)
		r.SetError()

		return err
	}

	a := op.Attributes
	r.LazyPrintf(
		"res: size=%d, mode=%s atime=%s mtime=%s", a.Size, a.Mode, a.Atime, a.Mtime,
	)

	return nil
}

func (t *TraceFS) LookUpInode(ctx context.Context, op *fuseops.LookUpInodeOp) error {
	r := t.newTrace("LookUpInode", "ParentID=%d Name=%s", op.Parent, op.Name)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.LookUpInode(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#LookUpInode err:%s", err)
		r.SetError()

		return err
	}

	e := op.Entry
	a := e.Attributes
	r.LazyPrintf(
		"res: size=%d, mode=%s atime=%s mtime=%s", a.Size, a.Mode, a.Atime, a.Mtime,
	)

	return nil
}

func (t *TraceFS) OpenDir(ctx context.Context, op *fuseops.OpenDirOp) error {
	r := t.newTrace("OpenDir", "ID=%d, HandleID=%d", op.Inode, op.Handle)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.OpenDir(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#OpenDir err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) ReadDir(ctx context.Context, op *fuseops.ReadDirOp) error {
	r := t.newTrace("ReadDir", "ID=%d Offset=%d", op.Inode, op.Offset)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.ReadDir(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#ReadDir err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf("res: read=%d bytes", op.BytesRead)

	return nil
}

func (t *TraceFS) MkDir(ctx context.Context, op *fuseops.MkDirOp) error {
	r := t.newTrace("MkDir", "ParentID=%d Name=%s", op.Parent, op.Name)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.MkDir(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#MkDir err:%s", err)
		r.SetError()

		return err
	}

	e := op.Entry
	a := e.Attributes
	r.LazyPrintf(
		"res: size=%d, mode=%s atime=%s mtime=%s", a.Size, a.Mode, a.Atime, a.Mtime,
	)

	return nil
}

func (t *TraceFS) Rename(ctx context.Context, op *fuseops.RenameOp) error {
	r := t.newTrace("Rename", "Old=%v,%s New=%v,%s", op.OldParent, op.OldName, op.NewParent, op.NewName)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.Rename(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#Rename err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) RmDir(ctx context.Context, op *fuseops.RmDirOp) error {
	r := t.newTrace("RmDir", "Parent=%d Name=%s", op.Parent, op.Name)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.RmDir(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#RmDir err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) OpenFile(ctx context.Context, op *fuseops.OpenFileOp) error {
	r := t.newTrace("OpenFile", "ID=%v", op.Inode)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.OpenFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#OpenFile err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf("res: handle=%d", op.Handle)

	return nil
}

func (t *TraceFS) ReadFile(ctx context.Context, op *fuseops.ReadFileOp) error {
	r := t.newTrace("ReadFile", "ID=%v Offset=%v", op.Inode, op.Offset)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.ReadFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#ReadFile err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf("res: read=%d bytes at offset=%d", op.BytesRead, op.Offset)

	return nil
}

func (t *TraceFS) WriteFile(ctx context.Context, op *fuseops.WriteFileOp) error {
	r := t.newTrace("WriteFile", "ID=%v DataLen=%v Offset=%v", op.Inode, len(op.Data), op.Offset)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.WriteFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#WriteFile err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf("result: wrote=%d bytes at offset=%d", len(op.Data), op.Offset)

	return nil
}

func (t *TraceFS) CreateFile(ctx context.Context, op *fuseops.CreateFileOp) error {
	r := t.newTrace("CreateFile", "Parent=%v Name=%s Mode=%s", op.Parent, op.Name, op.Mode)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.CreateFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#CreateFile err:%s", err)
		r.SetError()

		return err
	}

	e := op.Entry
	a := e.Attributes
	r.LazyPrintf(
		"res: size=%d, mode=%s atime=%s mtime=%s", a.Size, a.Mode, a.Atime, a.Mtime,
	)

	return nil
}

func (t *TraceFS) SetInodeAttributes(ctx context.Context, op *fuseops.SetInodeAttributesOp) error {
	r := t.newTrace("SetInodeAttributes", "ID=%v Size=%d Mode=%s", op.Inode, op.Size, op.Mode)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.SetInodeAttributes(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#SetInodeAttributes err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf(
		"res: size=%d, mode=%s atime=%s mtime=%s", op.Size, op.Mode, op.Atime, op.Mtime,
	)

	return nil
}

func (t *TraceFS) FlushFile(ctx context.Context, op *fuseops.FlushFileOp) error {
	r := t.newTrace("FlushFile", "ID=%d Handle=%d", op.Inode, op.Handle)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.FlushFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#FlushFile err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) SyncFile(ctx context.Context, op *fuseops.SyncFileOp) error {
	r := t.newTrace("SyncFile", "ID=%v Handle=%v", op.Inode, op.Handle)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.SyncFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#SyncFile err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) Unlink(ctx context.Context, op *fuseops.UnlinkOp) error {
	r := t.newTrace("Unlink", "Parent=%v Name=%s", op.Parent, op.Name)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.Unlink(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#Unlink err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) StatFS(ctx context.Context, op *fuseops.StatFSOp) error {
	r := t.newTrace("StatFS", "")
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.StatFS(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#StatFS err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf(
		"res: blockSize=%d blocks=%d blocksFree=%d, blocksAvailable=%d",
		op.BlockSize, op.Blocks, op.BlocksFree, op.BlocksAvailable,
	)

	return nil
}

func (t *TraceFS) ReleaseFileHandle(ctx context.Context, op *fuseops.ReleaseFileHandleOp) error {
	r := t.newTrace("ReleaseFileHandle", "Handle=%v", op.Handle)
	defer r.Finish()

	ctx = trace.NewContext(ctx, r)
	if err := t.KodingNetworkFS.ReleaseFileHandle(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#ReleaseFileHandle err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) newTrace(name, ft string, args ...interface{}) trace.Trace {
	argsFmt := fmt.Sprintf(ft, args...)

	v := fmt.Sprintf("%s-%s %s", t.Id, t.MountConfig.FSName, argsFmt)

	return trace.New(name, v)
}

///// Helpers

var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func randSeq(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}
<commit_msg>fuseklient: DRY tracing calls in tracefs<commit_after>package fuseklient

import (
	"fmt"
	"math/rand"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
	"golang.org/x/net/context"
	"golang.org/x/net/trace"
)

type TraceFS struct {
	Id string
	*KodingNetworkFS
}

func NewTraceFS(k *KodingNetworkFS) *TraceFS {
	return &TraceFS{
		Id:              randSeq(6),
		KodingNetworkFS: k,
	}
}

func (t *TraceFS) Mount() (*fuse.MountedFileSystem, error) {
	r, _ := t.newTrace(nil, "Mount", "Path=%s", t.MountPath)
	defer r.Finish()

	server := fuseutil.NewFileSystemServer(t)
	fs, err := fuse.Mount(t.MountPath, server, t.MountConfig)
	if err != nil {
		r.LazyPrintf("KodingNetworkFS#Mount: err '%s'", err)
		r.SetError()

		return nil, err
	}

	return fs, nil
}

func (t *TraceFS) Unmount() error {
	r, _ := t.newTrace(nil, "Unmount", "Path=%s", t.MountPath)
	defer r.Finish()

	if err := t.KodingNetworkFS.Unmount(); err != nil {
		r.LazyPrintf("KodingNetworkFS#Unmount err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) GetInodeAttributes(ctx context.Context, op *fuseops.GetInodeAttributesOp) error {
	r, ctx := t.newTrace(ctx, "GetInodeAttributes", "ID=%d", op.Inode)
	defer r.Finish()

	if err := t.KodingNetworkFS.GetInodeAttributes(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#GetInodeAttributes err:%s", err)
		r.SetError()

		return err
	}

	logAttrs(r, op.Attributes)

	return nil
}

func (t *TraceFS) LookUpInode(ctx context.Context, op *fuseops.LookUpInodeOp) error {
	r, ctx := t.newTrace(ctx, "LookUpInode", "ParentID=%d Name=%s", op.Parent, op.Name)
	defer r.Finish()

	if err := t.KodingNetworkFS.LookUpInode(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#LookUpInode err:%s", err)
		r.SetError()

		return err
	}

	logAttrs(r, op.Entry.Attributes)

	return nil
}

func (t *TraceFS) OpenDir(ctx context.Context, op *fuseops.OpenDirOp) error {
	r, ctx := t.newTrace(ctx, "OpenDir", "ID=%d, HandleID=%d", op.Inode, op.Handle)
	defer r.Finish()

	if err := t.KodingNetworkFS.OpenDir(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#OpenDir err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) ReadDir(ctx context.Context, op *fuseops.ReadDirOp) error {
	r, ctx := t.newTrace(ctx, "ReadDir", "ID=%d Offset=%d", op.Inode, op.Offset)
	defer r.Finish()

	if err := t.KodingNetworkFS.ReadDir(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#ReadDir err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf("res: read=%d bytes", op.BytesRead)

	return nil
}

func (t *TraceFS) MkDir(ctx context.Context, op *fuseops.MkDirOp) error {
	r, ctx := t.newTrace(ctx, "MkDir", "ParentID=%d Name=%s", op.Parent, op.Name)
	defer r.Finish()

	if err := t.KodingNetworkFS.MkDir(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#MkDir err:%s", err)
		r.SetError()

		return err
	}

	logAttrs(r, op.Entry.Attributes)

	return nil
}

func (t *TraceFS) Rename(ctx context.Context, op *fuseops.RenameOp) error {
	r, ctx := t.newTrace(ctx, "Rename", "Old=%v,%s New=%v,%s", op.OldParent, op.OldName, op.NewParent, op.NewName)
	defer r.Finish()

	if err := t.KodingNetworkFS.Rename(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#Rename err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) RmDir(ctx context.Context, op *fuseops.RmDirOp) error {
	r, ctx := t.newTrace(ctx, "RmDir", "Parent=%d Name=%s", op.Parent, op.Name)
	defer r.Finish()

	if err := t.KodingNetworkFS.RmDir(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#RmDir err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) OpenFile(ctx context.Context, op *fuseops.OpenFileOp) error {
	r, ctx := t.newTrace(ctx, "OpenFile", "ID=%v", op.Inode)
	defer r.Finish()

	if err := t.KodingNetworkFS.OpenFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#OpenFile err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf("res: handle=%d", op.Handle)

	return nil
}

func (t *TraceFS) ReadFile(ctx context.Context, op *fuseops.ReadFileOp) error {
	r, ctx := t.newTrace(ctx, "ReadFile", "ID=%v Offset=%v", op.Inode, op.Offset)
	defer r.Finish()

	if err := t.KodingNetworkFS.ReadFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#ReadFile err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf("res: read=%d bytes at offset=%d", op.BytesRead, op.Offset)

	return nil
}

func (t *TraceFS) WriteFile(ctx context.Context, op *fuseops.WriteFileOp) error {
	r, ctx := t.newTrace(ctx, "WriteFile", "ID=%v DataLen=%v Offset=%v", op.Inode, len(op.Data), op.Offset)
	defer r.Finish()

	if err := t.KodingNetworkFS.WriteFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#WriteFile err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf("result: wrote=%d bytes at offset=%d", len(op.Data), op.Offset)

	return nil
}

func (t *TraceFS) CreateFile(ctx context.Context, op *fuseops.CreateFileOp) error {
	r, ctx := t.newTrace(ctx, "CreateFile", "Parent=%v Name=%s Mode=%s", op.Parent, op.Name, op.Mode)
	defer r.Finish()

	if err := t.KodingNetworkFS.CreateFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#CreateFile err:%s", err)
		r.SetError()

		return err
	}

	logAttrs(r, op.Entry.Attributes)

	return nil
}

func (t *TraceFS) SetInodeAttributes(ctx context.Context, op *fuseops.SetInodeAttributesOp) error {
	r, ctx := t.newTrace(ctx, "SetInodeAttributes", "ID=%v Size=%d Mode=%s", op.Inode, op.Size, op.Mode)
	defer r.Finish()

	if err := t.KodingNetworkFS.SetInodeAttributes(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#SetInodeAttributes err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf(
		"res: size=%d, mode=%s atime=%s mtime=%s", op.Size, op.Mode, op.Atime, op.Mtime,
	)

	return nil
}

func (t *TraceFS) FlushFile(ctx context.Context, op *fuseops.FlushFileOp) error {
	r, ctx := t.newTrace(ctx, "FlushFile", "ID=%d Handle=%d", op.Inode, op.Handle)
	defer r.Finish()

	if err := t.KodingNetworkFS.FlushFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#FlushFile err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) SyncFile(ctx context.Context, op *fuseops.SyncFileOp) error {
	r, ctx := t.newTrace(ctx, "SyncFile", "ID=%v Handle=%v", op.Inode, op.Handle)
	defer r.Finish()

	if err := t.KodingNetworkFS.SyncFile(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#SyncFile err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) Unlink(ctx context.Context, op *fuseops.UnlinkOp) error {
	r, ctx := t.newTrace(ctx, "Unlink", "Parent=%v Name=%s", op.Parent, op.Name)
	defer r.Finish()

	if err := t.KodingNetworkFS.Unlink(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#Unlink err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

func (t *TraceFS) StatFS(ctx context.Context, op *fuseops.StatFSOp) error {
	r, ctx := t.newTrace(ctx, "StatFS", "")
	defer r.Finish()

	if err := t.KodingNetworkFS.StatFS(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#StatFS err:%s", err)
		r.SetError()

		return err
	}

	r.LazyPrintf(
		"res: blockSize=%d blocks=%d blocksFree=%d, blocksAvailable=%d",
		op.BlockSize, op.Blocks, op.BlocksFree, op.BlocksAvailable,
	)

	return nil
}

func (t *TraceFS) ReleaseFileHandle(ctx context.Context, op *fuseops.ReleaseFileHandleOp) error {
	r, ctx := t.newTrace(ctx, "ReleaseFileHandle", "Handle=%v", op.Handle)
	defer r.Finish()

	if err := t.KodingNetworkFS.ReleaseFileHandle(ctx, op); err != nil {
		r.LazyPrintf("KodingNetworkFS#ReleaseFileHandle err:%s", err)
		r.SetError()

		return err
	}

	return nil
}

///// Helpers

///// trace helpers

func (t *TraceFS) newTrace(ctx context.Context, name, ft string, args ...interface{}) (trace.Trace, context.Context) {
	argsFmt := fmt.Sprintf(ft, args...)
	v := fmt.Sprintf("%s-%s %s", t.Id, t.MountConfig.FSName, argsFmt)

	if ctx == nil {
		ctx = context.TODO()
	}

	r := trace.New(name, v)
	ctx = trace.NewContext(ctx, r)

	return r, ctx
}

func logAttrs(r trace.Trace, a fuseops.InodeAttributes) {
	r.LazyPrintf(
		"res: size=%d, mode=%s atime=%s mtime=%s", a.Size, a.Mode, a.Atime, a.Mtime,
	)
}

// random string generator

var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func randSeq(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreedto in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vtgate

import (
	"flag"
	"fmt"
	"net"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/mysql"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/callerid"
	"github.com/youtube/vitess/go/vt/servenv"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
	vtgatepb "github.com/youtube/vitess/go/vt/proto/vtgate"
	"github.com/youtube/vitess/go/vt/servenv/grpcutils"
)

var (
	mysqlServerPort               = flag.Int("mysql_server_port", -1, "If set, also listen for MySQL binary protocol connections on this port.")
	mysqlServerBindAddress        = flag.String("mysql_server_bind_address", "", "Binds on this address when listening to MySQL binary protocol. Useful to restrict listening to 'localhost' only for instance.")
	mysqlAuthServerImpl           = flag.String("mysql_auth_server_impl", "static", "Which auth server implementation to use.")
	mysqlAllowClearTextWithoutTLS = flag.Bool("mysql_allow_clear_text_without_tls", false, "If set, the server will allow the use of a clear text password over non-SSL connections.")

	mysqlSslCert = flag.String("mysql_server_ssl_cert", "", "Path to the ssl cert for mysql server plugin SSL")
	mysqlSslKey  = flag.String("mysql_server_ssl_key", "", "Path to ssl key for mysql server plugin SSL")
	mysqlSslCa   = flag.String("mysql_server_ssl_ca", "", "Path to ssl CA for mysql server plugin SSL. If specified, server will require and validate client certs.")

	mysqlSlowConnectWarnThreshold = flag.Duration("mysql_slow_connect_warn_threshold", 0, "Warn if it takes more than the given threshold for a mysql connection to establish")
)

// vtgateHandler implements the Listener interface.
// It stores the Session in the ClientData of a Connection, if a transaction
// is in progress.
type vtgateHandler struct {
	vtg *VTGate
}

func newVtgateHandler(vtg *VTGate) *vtgateHandler {
	return &vtgateHandler{
		vtg: vtg,
	}
}

func (vh *vtgateHandler) NewConnection(c *mysql.Conn) {
}

func (vh *vtgateHandler) ConnectionClosed(c *mysql.Conn) {
	// Rollback if there is an ongoing transaction. Ignore error.
	ctx := context.Background()
	session, _ := c.ClientData.(*vtgatepb.Session)
	if session != nil {
		_, _, _ = vh.vtg.Execute(ctx, session, "rollback", make(map[string]*querypb.BindVariable))
	}
}

func (vh *vtgateHandler) ComQuery(c *mysql.Conn, query []byte, callback func(*sqltypes.Result) error) error {
	// FIXME(alainjobart): Add some kind of timeout to the context.
	ctx := context.Background()

	// Fill in the ImmediateCallerID with the UserData returned by
	// the AuthServer plugin for that user. If nothing was
	// returned, use the User. This lets the plugin map a MySQL
	// user used for authentication to a Vitess User used for
	// Table ACLs and Vitess authentication in general.
	im := c.UserData.Get()
	ef := callerid.NewEffectiveCallerID(
		c.User,                  /* principal: who */
		c.RemoteAddr().String(), /* component: running client process */
		"VTGate MySQL Connector" /* subcomponent: part of the client */)
	ctx = callerid.NewContext(ctx, ef, im)

	session, _ := c.ClientData.(*vtgatepb.Session)
	if session == nil {
		session = &vtgatepb.Session{
			Options: &querypb.ExecuteOptions{
				IncludedFields: querypb.ExecuteOptions_ALL,
			},
		}
		if c.Capabilities&mysql.CapabilityClientFoundRows != 0 {
			session.Options.ClientFoundRows = true
		}
	}
	if c.SchemaName != "" {
		session.TargetString = c.SchemaName
	}
	if session.Options.Workload == querypb.ExecuteOptions_OLAP {
		err := vh.vtg.StreamExecute(ctx, session, string(query), make(map[string]*querypb.BindVariable), callback)
		return mysql.NewSQLErrorFromError(err)
	}
	session, result, err := vh.vtg.Execute(ctx, session, string(query), make(map[string]*querypb.BindVariable))
	c.ClientData = session
	err = mysql.NewSQLErrorFromError(err)
	if err != nil {
		return err
	}
	return callback(result)
}

var mysqlListener *mysql.Listener

// initiMySQLProtocol starts the mysql protocol.
// It should be called only once in a process.
func initMySQLProtocol() {
	// Flag is not set, just return.
	if *mysqlServerPort < 0 {
		return
	}

	// If no VTGate was created, just return.
	if rpcVTGate == nil {
		return
	}

	// Initialize registered AuthServer implementations (or other plugins)
	for _, initFn := range pluginInitializers {
		initFn()
	}
	authServer := mysql.GetAuthServer(*mysqlAuthServerImpl)

	// Create a Listener.
	var err error
	vh := newVtgateHandler(rpcVTGate)
	mysqlListener, err = mysql.NewListener("tcp", net.JoinHostPort(*mysqlServerBindAddress, fmt.Sprintf("%v", *mysqlServerPort)), authServer, vh)
	if err != nil {
		log.Fatalf("mysql.NewListener failed: %v", err)
	}
	if *mysqlSslCert != "" && *mysqlSslKey != "" {
		mysqlListener.TLSConfig, err = grpcutils.TLSServerConfig(*mysqlSslCert, *mysqlSslKey, *mysqlSslCa)
		if err != nil {
			log.Fatalf("grpcutils.TLSServerConfig failed: %v", err)
			return
		}
	}
	mysqlListener.AllowClearTextWithoutTLS = *mysqlAllowClearTextWithoutTLS

	// Check for the connection threshold
	if *mysqlSlowConnectWarnThreshold != 0 {
		log.Infof("setting mysql slow connection threshold to %v", mysqlSlowConnectWarnThreshold)
		mysqlListener.SlowConnectWarnThreshold = *mysqlSlowConnectWarnThreshold
	}

	// And starts listening.
	go func() {
		mysqlListener.Accept()
	}()
}

func init() {
	servenv.OnRun(initMySQLProtocol)

	servenv.OnTerm(func() {
		if mysqlListener != nil {
			mysqlListener.Close()
			mysqlListener = nil
		}
	})
}

var pluginInitializers []func()

// RegisterPluginInitializer lets plugins register themselves to be init'ed at servenv.OnRun-time
func RegisterPluginInitializer(initializer func()) {
	pluginInitializers = append(pluginInitializers, initializer)
}
<commit_msg>set autocommit to true by default for mysql server sessions<commit_after>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreedto in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vtgate

import (
	"flag"
	"fmt"
	"net"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/mysql"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/callerid"
	"github.com/youtube/vitess/go/vt/servenv"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
	vtgatepb "github.com/youtube/vitess/go/vt/proto/vtgate"
	"github.com/youtube/vitess/go/vt/servenv/grpcutils"
)

var (
	mysqlServerPort               = flag.Int("mysql_server_port", -1, "If set, also listen for MySQL binary protocol connections on this port.")
	mysqlServerBindAddress        = flag.String("mysql_server_bind_address", "", "Binds on this address when listening to MySQL binary protocol. Useful to restrict listening to 'localhost' only for instance.")
	mysqlAuthServerImpl           = flag.String("mysql_auth_server_impl", "static", "Which auth server implementation to use.")
	mysqlAllowClearTextWithoutTLS = flag.Bool("mysql_allow_clear_text_without_tls", false, "If set, the server will allow the use of a clear text password over non-SSL connections.")

	mysqlSslCert = flag.String("mysql_server_ssl_cert", "", "Path to the ssl cert for mysql server plugin SSL")
	mysqlSslKey  = flag.String("mysql_server_ssl_key", "", "Path to ssl key for mysql server plugin SSL")
	mysqlSslCa   = flag.String("mysql_server_ssl_ca", "", "Path to ssl CA for mysql server plugin SSL. If specified, server will require and validate client certs.")

	mysqlSlowConnectWarnThreshold = flag.Duration("mysql_slow_connect_warn_threshold", 0, "Warn if it takes more than the given threshold for a mysql connection to establish")
)

// vtgateHandler implements the Listener interface.
// It stores the Session in the ClientData of a Connection, if a transaction
// is in progress.
type vtgateHandler struct {
	vtg *VTGate
}

func newVtgateHandler(vtg *VTGate) *vtgateHandler {
	return &vtgateHandler{
		vtg: vtg,
	}
}

func (vh *vtgateHandler) NewConnection(c *mysql.Conn) {
}

func (vh *vtgateHandler) ConnectionClosed(c *mysql.Conn) {
	// Rollback if there is an ongoing transaction. Ignore error.
	ctx := context.Background()
	session, _ := c.ClientData.(*vtgatepb.Session)
	if session != nil {
		_, _, _ = vh.vtg.Execute(ctx, session, "rollback", make(map[string]*querypb.BindVariable))
	}
}

func (vh *vtgateHandler) ComQuery(c *mysql.Conn, query []byte, callback func(*sqltypes.Result) error) error {
	// FIXME(alainjobart): Add some kind of timeout to the context.
	ctx := context.Background()

	// Fill in the ImmediateCallerID with the UserData returned by
	// the AuthServer plugin for that user. If nothing was
	// returned, use the User. This lets the plugin map a MySQL
	// user used for authentication to a Vitess User used for
	// Table ACLs and Vitess authentication in general.
	im := c.UserData.Get()
	ef := callerid.NewEffectiveCallerID(
		c.User,                  /* principal: who */
		c.RemoteAddr().String(), /* component: running client process */
		"VTGate MySQL Connector" /* subcomponent: part of the client */)
	ctx = callerid.NewContext(ctx, ef, im)

	session, _ := c.ClientData.(*vtgatepb.Session)
	if session == nil {
		session = &vtgatepb.Session{
			Options: &querypb.ExecuteOptions{
				IncludedFields: querypb.ExecuteOptions_ALL,
			},
			Autocommit: true,
		}
		if c.Capabilities&mysql.CapabilityClientFoundRows != 0 {
			session.Options.ClientFoundRows = true
		}
	}
	if c.SchemaName != "" {
		session.TargetString = c.SchemaName
	}
	if session.Options.Workload == querypb.ExecuteOptions_OLAP {
		err := vh.vtg.StreamExecute(ctx, session, string(query), make(map[string]*querypb.BindVariable), callback)
		return mysql.NewSQLErrorFromError(err)
	}
	session, result, err := vh.vtg.Execute(ctx, session, string(query), make(map[string]*querypb.BindVariable))
	c.ClientData = session
	err = mysql.NewSQLErrorFromError(err)
	if err != nil {
		return err
	}
	return callback(result)
}

var mysqlListener *mysql.Listener

// initiMySQLProtocol starts the mysql protocol.
// It should be called only once in a process.
func initMySQLProtocol() {
	// Flag is not set, just return.
	if *mysqlServerPort < 0 {
		return
	}

	// If no VTGate was created, just return.
	if rpcVTGate == nil {
		return
	}

	// Initialize registered AuthServer implementations (or other plugins)
	for _, initFn := range pluginInitializers {
		initFn()
	}
	authServer := mysql.GetAuthServer(*mysqlAuthServerImpl)

	// Create a Listener.
	var err error
	vh := newVtgateHandler(rpcVTGate)
	mysqlListener, err = mysql.NewListener("tcp", net.JoinHostPort(*mysqlServerBindAddress, fmt.Sprintf("%v", *mysqlServerPort)), authServer, vh)
	if err != nil {
		log.Fatalf("mysql.NewListener failed: %v", err)
	}
	if *mysqlSslCert != "" && *mysqlSslKey != "" {
		mysqlListener.TLSConfig, err = grpcutils.TLSServerConfig(*mysqlSslCert, *mysqlSslKey, *mysqlSslCa)
		if err != nil {
			log.Fatalf("grpcutils.TLSServerConfig failed: %v", err)
			return
		}
	}
	mysqlListener.AllowClearTextWithoutTLS = *mysqlAllowClearTextWithoutTLS

	// Check for the connection threshold
	if *mysqlSlowConnectWarnThreshold != 0 {
		log.Infof("setting mysql slow connection threshold to %v", mysqlSlowConnectWarnThreshold)
		mysqlListener.SlowConnectWarnThreshold = *mysqlSlowConnectWarnThreshold
	}

	// And starts listening.
	go func() {
		mysqlListener.Accept()
	}()
}

func init() {
	servenv.OnRun(initMySQLProtocol)

	servenv.OnTerm(func() {
		if mysqlListener != nil {
			mysqlListener.Close()
			mysqlListener = nil
		}
	})
}

var pluginInitializers []func()

// RegisterPluginInitializer lets plugins register themselves to be init'ed at servenv.OnRun-time
func RegisterPluginInitializer(initializer func()) {
	pluginInitializers = append(pluginInitializers, initializer)
}
<|endoftext|>