<commit_before>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"crypto/tls"
	"database/sql/driver"
	"errors"
	"net"
	"strconv"
	"strings"
	"time"
)

type mysqlConn struct {
	buf              buffer
	netConn          net.Conn
	affectedRows     uint64
	insertId         uint64
	cfg              *config
	maxPacketAllowed int
	maxWriteSize     int
	flags            clientFlag
	status           statusFlag
	sequence         uint8
	parseTime        bool
	strict           bool
}

type config struct {
	user              string
	passwd            string
	net               string
	addr              string
	dbname            string
	params            map[string]string
	loc               *time.Location
	tls               *tls.Config
	timeout           time.Duration
	collation         uint8
	allowAllFiles     bool
	allowOldPasswords bool
	clientFoundRows   bool
	columnsWithAlias  bool
	interpolateParams bool
}

// Handles parameters set in DSN after the connection is established
func (mc *mysqlConn) handleParams() (err error) {
	for param, val := range mc.cfg.params {
		switch param {
		// Charset
		case "charset":
			charsets := strings.Split(val, ",")
			for i := range charsets {
				// ignore errors here - a charset may not exist
				err = mc.exec("SET NAMES " + charsets[i])
				if err == nil {
					break
				}
			}
			if err != nil {
				return
			}

		// time.Time parsing
		case "parseTime":
			var isBool bool
			mc.parseTime, isBool = readBool(val)
			if !isBool {
				return errors.New("Invalid Bool value: " + val)
			}

		// Strict mode
		case "strict":
			var isBool bool
			mc.strict, isBool = readBool(val)
			if !isBool {
				return errors.New("Invalid Bool value: " + val)
			}

		// Compression
		case "compress":
			err = errors.New("Compression not implemented yet")
			return

		// System Vars
		default:
			err = mc.exec("SET " + param + "=" + val + "")
			if err != nil {
				return
			}
		}
	}

	return
}

func (mc *mysqlConn) Begin() (driver.Tx, error) {
	if mc.netConn == nil {
		errLog.Print(ErrInvalidConn)
		return nil, driver.ErrBadConn
	}
	err := mc.exec("START TRANSACTION")
	if err == nil {
		return &mysqlTx{mc}, err
	}

	return nil, err
}

func (mc *mysqlConn) Close() (err error) {
	// Makes Close idempotent
	if mc.netConn != nil {
		err = mc.writeCommandPacket(comQuit)
		if err == nil {
			err = mc.netConn.Close()
		} else {
			mc.netConn.Close()
		}
		mc.netConn = nil
	}

	mc.cfg = nil
	mc.buf.rd = nil

	return
}

func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) {
	if mc.netConn == nil {
		errLog.Print(ErrInvalidConn)
		return nil, driver.ErrBadConn
	}
	// Send command
	err := mc.writeCommandPacketStr(comStmtPrepare, query)
	if err != nil {
		return nil, err
	}

	stmt := &mysqlStmt{
		mc: mc,
	}

	// Read Result
	columnCount, err := stmt.readPrepareResultPacket()
	if err == nil {
		if stmt.paramCount > 0 {
			if err = mc.readUntilEOF(); err != nil {
				return nil, err
			}
		}

		if columnCount > 0 {
			err = mc.readUntilEOF()
		}
	}

	return stmt, err
}

// https://github.com/mysql/mysql-server/blob/mysql-5.7.5/libmysql/libmysql.c#L1150-L1156
func (mc *mysqlConn) escapeBytes(buf, v []byte) []byte {
	buf = append(buf, '\'')
	if mc.status&statusNoBackslashEscapes == 0 {
		buf = escapeBytesBackslash(buf, v)
	} else {
		buf = escapeBytesQuotes(buf, v)
	}
	return append(buf, '\'')
}

func (mc *mysqlConn) escapeString(buf []byte, v string) []byte {
	buf = append(buf, '\'')
	if mc.status&statusNoBackslashEscapes == 0 {
		buf = escapeStringBackslash(buf, v)
	} else {
		buf = escapeStringQuotes(buf, v)
	}
	return append(buf, '\'')
}

// estimateParamLength calculates upper bound of string length from types.
func estimateParamLength(args []driver.Value) (int, bool) {
	l := 0
	for _, a := range args {
		switch v := a.(type) {
		case int64, float64:
			// 24 (-1.7976931348623157e+308) may be upper bound. But I'm not sure.
			l += 25
		case bool:
			l += 1 // 0 or 1
		case time.Time:
			l += 30 // '1234-12-23 12:34:56.777777'
		case string:
			l += len(v)*2 + 2
		case []byte:
			l += len(v)*2 + 2
		default:
			return 0, false
		}
	}
	return l, true
}

func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (string, error) {
	estimated, ok := estimateParamLength(args)
	if !ok {
		return "", driver.ErrSkip
	}
	estimated += len(query)

	buf := make([]byte, 0, estimated)
	argPos := 0

	for i := 0; i < len(query); i++ {
		q := strings.IndexByte(query[i:], '?')
		if q == -1 {
			buf = append(buf, query[i:]...)
			break
		}
		buf = append(buf, query[i:i+q]...)
		i += q

		arg := args[argPos]
		argPos++

		if arg == nil {
			buf = append(buf, "NULL"...)
			continue
		}

		switch v := arg.(type) {
		case int64:
			buf = strconv.AppendInt(buf, v, 10)
		case float64:
			buf = strconv.AppendFloat(buf, v, 'g', -1, 64)
		case bool:
			if v {
				buf = append(buf, '1')
			} else {
				buf = append(buf, '0')
			}
		case time.Time:
			if v.IsZero() {
				buf = append(buf, "'0000-00-00'"...)
			} else {
				v := v.In(mc.cfg.loc)
				v = v.Add(time.Nanosecond * 500) // To round under microsecond
				year := v.Year()
				year100 := year / 100
				year1 := year % 100
				month := v.Month()
				day := v.Day()
				hour := v.Hour()
				minute := v.Minute()
				second := v.Second()
				micro := v.Nanosecond() / 1000

				buf = append(buf, []byte{
					'\'',
					digits10[year100], digits01[year100],
					digits10[year1], digits01[year1],
					'-',
					digits10[month], digits01[month],
					'-',
					digits10[day], digits01[day],
					' ',
					digits10[hour], digits01[hour],
					':',
					digits10[minute], digits01[minute],
					':',
					digits10[second], digits01[second],
				}...)

				if micro != 0 {
					micro10000 := micro / 10000
					micro100 := micro / 100 % 100
					micro1 := micro % 100
					buf = append(buf, []byte{
						'.',
						digits10[micro10000], digits01[micro10000],
						digits10[micro100], digits01[micro100],
						digits10[micro1], digits01[micro1],
					}...)
				}
				buf = append(buf, '\'')
			}
		case []byte:
			if v == nil {
				buf = append(buf, "NULL"...)
			} else {
				buf = mc.escapeBytes(buf, v)
			}
		case string:
			buf = mc.escapeString(buf, v)
		default:
			return "", driver.ErrSkip
		}

		if len(buf)+4 > mc.maxPacketAllowed {
			return "", driver.ErrSkip
		}
	}
	if argPos != len(args) {
		return "", driver.ErrSkip
	}
	return string(buf), nil
}

func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, error) {
	if mc.netConn == nil {
		errLog.Print(ErrInvalidConn)
		return nil, driver.ErrBadConn
	}
	if len(args) != 0 {
		if !mc.cfg.interpolateParams {
			return nil, driver.ErrSkip
		}
		// try client-side prepare to reduce roundtrip
		prepared, err := mc.interpolateParams(query, args)
		if err != nil {
			return nil, err
		}
		query = prepared
		args = nil
	}
	mc.affectedRows = 0
	mc.insertId = 0

	err := mc.exec(query)
	if err == nil {
		return &mysqlResult{
			affectedRows: int64(mc.affectedRows),
			insertId:     int64(mc.insertId),
		}, err
	}
	return nil, err
}

// Internal function to execute commands
func (mc *mysqlConn) exec(query string) error {
	// Send command
	err := mc.writeCommandPacketStr(comQuery, query)
	if err != nil {
		return err
	}

	// Read Result
	resLen, err := mc.readResultSetHeaderPacket()
	if err == nil && resLen > 0 {
		if err = mc.readUntilEOF(); err != nil {
			return err
		}

		err = mc.readUntilEOF()
	}

	return err
}

func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, error) {
	if mc.netConn == nil {
		errLog.Print(ErrInvalidConn)
		return nil, driver.ErrBadConn
	}
	if len(args) != 0 {
		if !mc.cfg.interpolateParams {
			return nil, driver.ErrSkip
		}
		// try client-side prepare to reduce roundtrip
		prepared, err := mc.interpolateParams(query, args)
		if err != nil {
			return nil, err
		}
		query = prepared
		args = nil
	}
	// Send command
	err := mc.writeCommandPacketStr(comQuery, query)
	if err == nil {
		// Read Result
		var resLen int
		resLen, err = mc.readResultSetHeaderPacket()
		if err == nil {
			rows := new(textRows)
			rows.mc = mc

			if resLen == 0 {
				// no columns, no more data
				return emptyRows{}, nil
			}
			// Columns
			rows.columns, err = mc.readColumns(resLen)
			return rows, err
		}
	}
	return nil, err
}

// Gets the value of the given MySQL System Variable
// The returned byte slice is only valid until the next read
func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
	// Send command
	if err := mc.writeCommandPacketStr(comQuery, "SELECT @@"+name); err != nil {
		return nil, err
	}

	// Read Result
	resLen, err := mc.readResultSetHeaderPacket()
	if err == nil {
		rows := new(textRows)
		rows.mc = mc

		if resLen > 0 {
			// Columns
			if err := mc.readUntilEOF(); err != nil {
				return nil, err
			}
		}

		dest := make([]driver.Value, resLen)
		if err = rows.readRow(dest); err == nil {
			return dest[0].([]byte), mc.readUntilEOF()
		}
	}
	return nil, err
}
<commit_msg>Inlining mysqlConn.escapeBytes and mysqlConn.escapeString<commit_after>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"crypto/tls"
	"database/sql/driver"
	"errors"
	"net"
	"strconv"
	"strings"
	"time"
)

type mysqlConn struct {
	buf              buffer
	netConn          net.Conn
	affectedRows     uint64
	insertId         uint64
	cfg              *config
	maxPacketAllowed int
	maxWriteSize     int
	flags            clientFlag
	status           statusFlag
	sequence         uint8
	parseTime        bool
	strict           bool
}

type config struct {
	user              string
	passwd            string
	net               string
	addr              string
	dbname            string
	params            map[string]string
	loc               *time.Location
	tls               *tls.Config
	timeout           time.Duration
	collation         uint8
	allowAllFiles     bool
	allowOldPasswords bool
	clientFoundRows   bool
	columnsWithAlias  bool
	interpolateParams bool
}

// Handles parameters set in DSN after the connection is established
func (mc *mysqlConn) handleParams() (err error) {
	for param, val := range mc.cfg.params {
		switch param {
		// Charset
		case "charset":
			charsets := strings.Split(val, ",")
			for i := range charsets {
				// ignore errors here - a charset may not exist
				err = mc.exec("SET NAMES " + charsets[i])
				if err == nil {
					break
				}
			}
			if err != nil {
				return
			}

		// time.Time parsing
		case "parseTime":
			var isBool bool
			mc.parseTime, isBool = readBool(val)
			if !isBool {
				return errors.New("Invalid Bool value: " + val)
			}

		// Strict mode
		case "strict":
			var isBool bool
			mc.strict, isBool = readBool(val)
			if !isBool {
				return errors.New("Invalid Bool value: " + val)
			}

		// Compression
		case "compress":
			err = errors.New("Compression not implemented yet")
			return

		// System Vars
		default:
			err = mc.exec("SET " + param + "=" + val + "")
			if err != nil {
				return
			}
		}
	}

	return
}

func (mc *mysqlConn) Begin() (driver.Tx, error) {
	if mc.netConn == nil {
		errLog.Print(ErrInvalidConn)
		return nil, driver.ErrBadConn
	}
	err := mc.exec("START TRANSACTION")
	if err == nil {
		return &mysqlTx{mc}, err
	}

	return nil, err
}

func (mc *mysqlConn) Close() (err error) {
	// Makes Close idempotent
	if mc.netConn != nil {
		err = mc.writeCommandPacket(comQuit)
		if err == nil {
			err = mc.netConn.Close()
		} else {
			mc.netConn.Close()
		}
		mc.netConn = nil
	}

	mc.cfg = nil
	mc.buf.rd = nil

	return
}

func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) {
	if mc.netConn == nil {
		errLog.Print(ErrInvalidConn)
		return nil, driver.ErrBadConn
	}
	// Send command
	err := mc.writeCommandPacketStr(comStmtPrepare, query)
	if err != nil {
		return nil, err
	}

	stmt := &mysqlStmt{
		mc: mc,
	}

	// Read Result
	columnCount, err := stmt.readPrepareResultPacket()
	if err == nil {
		if stmt.paramCount > 0 {
			if err = mc.readUntilEOF(); err != nil {
				return nil, err
			}
		}

		if columnCount > 0 {
			err = mc.readUntilEOF()
		}
	}

	return stmt, err
}

// estimateParamLength calculates upper bound of string length from types.
func estimateParamLength(args []driver.Value) (int, bool) {
	l := 0
	for _, a := range args {
		switch v := a.(type) {
		case int64, float64:
			// 24 (-1.7976931348623157e+308) may be upper bound. But I'm not sure.
			l += 25
		case bool:
			l += 1 // 0 or 1
		case time.Time:
			l += 30 // '1234-12-23 12:34:56.777777'
		case string:
			l += len(v)*2 + 2
		case []byte:
			l += len(v)*2 + 2
		default:
			return 0, false
		}
	}
	return l, true
}

func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (string, error) {
	estimated, ok := estimateParamLength(args)
	if !ok {
		return "", driver.ErrSkip
	}
	estimated += len(query)

	buf := make([]byte, 0, estimated)
	argPos := 0

	for i := 0; i < len(query); i++ {
		q := strings.IndexByte(query[i:], '?')
		if q == -1 {
			buf = append(buf, query[i:]...)
			break
		}
		buf = append(buf, query[i:i+q]...)
		i += q

		arg := args[argPos]
		argPos++

		if arg == nil {
			buf = append(buf, "NULL"...)
			continue
		}

		switch v := arg.(type) {
		case int64:
			buf = strconv.AppendInt(buf, v, 10)
		case float64:
			buf = strconv.AppendFloat(buf, v, 'g', -1, 64)
		case bool:
			if v {
				buf = append(buf, '1')
			} else {
				buf = append(buf, '0')
			}
		case time.Time:
			if v.IsZero() {
				buf = append(buf, "'0000-00-00'"...)
			} else {
				v := v.In(mc.cfg.loc)
				v = v.Add(time.Nanosecond * 500) // To round under microsecond
				year := v.Year()
				year100 := year / 100
				year1 := year % 100
				month := v.Month()
				day := v.Day()
				hour := v.Hour()
				minute := v.Minute()
				second := v.Second()
				micro := v.Nanosecond() / 1000

				buf = append(buf, []byte{
					'\'',
					digits10[year100], digits01[year100],
					digits10[year1], digits01[year1],
					'-',
					digits10[month], digits01[month],
					'-',
					digits10[day], digits01[day],
					' ',
					digits10[hour], digits01[hour],
					':',
					digits10[minute], digits01[minute],
					':',
					digits10[second], digits01[second],
				}...)

				if micro != 0 {
					micro10000 := micro / 10000
					micro100 := micro / 100 % 100
					micro1 := micro % 100
					buf = append(buf, []byte{
						'.',
						digits10[micro10000], digits01[micro10000],
						digits10[micro100], digits01[micro100],
						digits10[micro1], digits01[micro1],
					}...)
				}
				buf = append(buf, '\'')
			}
		case []byte:
			if v == nil {
				buf = append(buf, "NULL"...)
			} else {
				buf = append(buf, '\'')
				if mc.status&statusNoBackslashEscapes == 0 {
					buf = escapeBytesBackslash(buf, v)
				} else {
					buf = escapeBytesQuotes(buf, v)
				}
				buf = append(buf, '\'')
			}
		case string:
			buf = append(buf, '\'')
			if mc.status&statusNoBackslashEscapes == 0 {
				buf = escapeStringBackslash(buf, v)
			} else {
				buf = escapeStringQuotes(buf, v)
			}
			buf = append(buf, '\'')
		default:
			return "", driver.ErrSkip
		}

		if len(buf)+4 > mc.maxPacketAllowed {
			return "", driver.ErrSkip
		}
	}
	if argPos != len(args) {
		return "", driver.ErrSkip
	}
	return string(buf), nil
}

func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, error) {
	if mc.netConn == nil {
		errLog.Print(ErrInvalidConn)
		return nil, driver.ErrBadConn
	}
	if len(args) != 0 {
		if !mc.cfg.interpolateParams {
			return nil, driver.ErrSkip
		}
		// try client-side prepare to reduce roundtrip
		prepared, err := mc.interpolateParams(query, args)
		if err != nil {
			return nil, err
		}
		query = prepared
		args = nil
	}
	mc.affectedRows = 0
	mc.insertId = 0

	err := mc.exec(query)
	if err == nil {
		return &mysqlResult{
			affectedRows: int64(mc.affectedRows),
			insertId:     int64(mc.insertId),
		}, err
	}
	return nil, err
}

// Internal function to execute commands
func (mc *mysqlConn) exec(query string) error {
	// Send command
	err := mc.writeCommandPacketStr(comQuery, query)
	if err != nil {
		return err
	}

	// Read Result
	resLen, err := mc.readResultSetHeaderPacket()
	if err == nil && resLen > 0 {
		if err = mc.readUntilEOF(); err != nil {
			return err
		}

		err = mc.readUntilEOF()
	}

	return err
}

func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, error) {
	if mc.netConn == nil {
		errLog.Print(ErrInvalidConn)
		return nil, driver.ErrBadConn
	}
	if len(args) != 0 {
		if !mc.cfg.interpolateParams {
			return nil, driver.ErrSkip
		}
		// try client-side prepare to reduce roundtrip
		prepared, err := mc.interpolateParams(query, args)
		if err != nil {
			return nil, err
		}
		query = prepared
		args = nil
	}
	// Send command
	err := mc.writeCommandPacketStr(comQuery, query)
	if err == nil {
		// Read Result
		var resLen int
		resLen, err = mc.readResultSetHeaderPacket()
		if err == nil {
			rows := new(textRows)
			rows.mc = mc

			if resLen == 0 {
				// no columns, no more data
				return emptyRows{}, nil
			}
			// Columns
			rows.columns, err = mc.readColumns(resLen)
			return rows, err
		}
	}
	return nil, err
}

// Gets the value of the given MySQL System Variable
// The returned byte slice is only valid until the next read
func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
	// Send command
	if err := mc.writeCommandPacketStr(comQuery, "SELECT @@"+name); err != nil {
		return nil, err
	}

	// Read Result
	resLen, err := mc.readResultSetHeaderPacket()
	if err == nil {
		rows := new(textRows)
		rows.mc = mc

		if resLen > 0 {
			// Columns
			if err := mc.readUntilEOF(); err != nil {
				return nil, err
			}
		}

		dest := make([]driver.Value, resLen)
		if err = rows.readRow(dest); err == nil {
			return dest[0].([]byte), mc.readUntilEOF()
		}
	}
	return nil, err
}
<|endoftext|>


<commit_before>package xmmsclient

import (
	"errors"
	"math"
	"strings"
)

var defaultSourcePrefs = []string{
	"server",
	"client/*",
	"plugin/playlist",
	"plugin/segment",
	"plugin/nibbler",
	"plugin/id3v2",
	"plugin/*",
	"*",
}

const (
	NoMatch int = math.MaxInt32
)

func match(sourcePrefs []string, source string) int {
	for index, pattern := range sourcePrefs {
		if strings.HasSuffix(pattern, "*") {
			if strings.HasPrefix(source, pattern[:len(pattern)-1]) {
				return index
			}
		} else if source == pattern {
			return index
		}
	}
	return NoMatch
}

// Flatten a {key: {source: value}} to {key: value} based on source preferences.
// Source preferences are defined as a list of strings that may end with '*' to
// match any suffix.
func PropDictToDict(inputValue XmmsDict, sourcePrefs []string) (XmmsDict, error) {
	result := XmmsDict{}
	for key, innerValue := range inputValue {
		innerDict, ok := innerValue.(XmmsDict)
		if !ok {
			return nil, errors.New("Input not a XmmsDict->XmmsDict->XmmsValue")
		}

		bestScore := NoMatch
		for source, value := range innerDict {
			score := match(sourcePrefs, source)
			if score < bestScore {
				result[key] = value
				bestScore = score
			}
		}
	}

	return result, nil
}

// Flatten a PropDict using default source preferences:
//  ["server", "client/*", "plugin/playlist", "plugin/segment",
//   "plugin/nibbler", "plugin/id3v2", "plugin/*", "*"]
func PropDictToDictDefault(input XmmsDict) (XmmsDict, error) {
	return PropDictToDict(input, defaultSourcePrefs)
}
<commit_msg>Make PropDict funcs only accept XmmsDict.<commit_after>package xmmsclient

import (
	"errors"
	"math"
	"strings"
)

var defaultSourcePrefs = []string{
	"server",
	"client/*",
	"plugin/playlist",
	"plugin/segment",
	"plugin/nibbler",
	"plugin/id3v2",
	"plugin/*",
	"*",
}

const (
	NoMatch int = math.MaxInt32
)

func match(sourcePrefs []string, source string) int {
	for index, pattern := range sourcePrefs {
		if strings.HasSuffix(pattern, "*") {
			if strings.HasPrefix(source, pattern[:len(pattern)-1]) {
				return index
			}
		} else if source == pattern {
			return index
		}
	}
	return NoMatch
}

// Flatten a {key: {source: value}} to {key: value} based on source preferences.
// Source preferences are defined as a list of strings that may end with '*' to
// match any suffix.
func PropDictToDict(propDict XmmsDict, sourcePrefs []string) (XmmsDict, error) {
	result := XmmsDict{}
	for key, inner := range propDict {
		sourceDict, ok := inner.(XmmsDict)
		if !ok {
			return nil, errors.New("Input not a XmmsDict->XmmsDict->XmmsValue")
		}

		bestScore := NoMatch
		for source, value := range sourceDict {
			score := match(sourcePrefs, source)
			if score < bestScore {
				result[key] = value
				bestScore = score
			}
		}
	}

	return result, nil
}

// Flatten a PropDict using default source preferences:
//  ["server", "client/*", "plugin/playlist", "plugin/segment",
//   "plugin/nibbler", "plugin/id3v2", "plugin/*", "*"]
func PropDictToDictDefault(propDict XmmsDict) (XmmsDict, error) {
	return PropDictToDict(propDict, defaultSourcePrefs)
}
<|endoftext|>


<commit_before>package meta

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/eyecuelab/kit/flect"
	"github.com/google/jsonapi"
	"github.com/jinzhu/inflection"
	"github.com/spf13/viper"
)

type (
	method       string
	inputType    string
	ActionHolder []*JsonAPIAction

	Extendable interface {
		Extend() error
	}

	JsonAPIAction struct {
		Method string         `json:"method"`
		Name   string         `json:"name"`
		URL    string         `json:"url"`
		Multi  bool           `json:"multi"`
		Fields []JsonAPIField `json:"fields"`
	}

	JsonAPIField struct {
		Name      string        `json:"name"`
		InputType string        `json:"type"`
		Value     interface{}   `json:"value,omitempty"`
		Required  bool          `json:"required"`
		Options   []FieldOption `json:"options,omitempty"`
		Data      *Pagination   `json:"data,omitempty"`
	}

	FieldOption struct {
		Label string                 `json:"label"`
		Value interface{}            `json:"value"`
		Meta  map[string]interface{} `json:"meta,omitempty"`
	}

	Pagination struct {
		Count int `json:"item_count"`
		Max   int `json:"max"`
		Page  int `json:"page"`
	}
)

const (
	DELETE method = "DELETE"
	GET    method = "GET"
	POST   method = "POST"
	PATCH  method = "PATCH"
	PUT    method = "PUT"

	InputText   inputType = "text"
	InputPass   inputType = "password"
	InputBool   inputType = "bool"
	InputNumber inputType = "number"
	InputSelect inputType = "select"
)

// AddAction ...
func (ah *ActionHolder) AddAction(m method, name, urlHelper string, params ...interface{}) *JsonAPIAction {
	a := JsonAPIAction{
		Method: string(m),
		Name:   name,
		URL:    APIURL(fmt.Sprintf(urlHelper, params...)),
		Fields: make([]JsonAPIField, 0),
	}
	*ah = append(*ah, &a)
	return &a
}

// Field ...
func (a *JsonAPIAction) Field(name string, inputType inputType, value interface{}, requred bool) *JsonAPIAction {
	f := JsonAPIField{
		Name:      name,
		InputType: string(inputType),
		Value:     value,
		Required:  requred,
	}
	a.Fields = append(a.Fields, f)
	return a
}

// FieldWithOpts add field with options to action
func (a *JsonAPIAction) FieldWithOpts(name string, inputType inputType, value interface{}, requred bool, options []FieldOption) *JsonAPIAction {
	f := JsonAPIField{
		Name:      name,
		InputType: string(inputType),
		Value:     value,
		Required:  requred,
		Options:   options,
	}
	a.Fields = append(a.Fields, f)
	return a
}

// Pagination add pagination meta to action
func (a *JsonAPIAction) Pagination(data *Pagination) *JsonAPIAction {
	f := JsonAPIField{
		Name:      "page",
		InputType: string(InputNumber),
		Value:     data.Page,
		Required:  false,
		Data:      data,
	}
	a.Fields = append(a.Fields, f)
	return a
}

// RenderActions ...
func (ah *ActionHolder) RenderActions() *jsonapi.Meta {
	return &jsonapi.Meta{"actions": *ah}
}

// APIURL full api url for the path
func APIURL(path string) string {
	return fmt.Sprintf("%s%s", viper.GetString("root_url"), path)
}

// JSONApiSelfLink self link helper
func JSONApiSelfLink(i interface{}) *jsonapi.Links {
	name := reflect.TypeOf(i).Name()
	name = inflection.Plural(strings.ToLower(name))

	id := reflect.ValueOf(i).FieldByName("ID").Interface()

	return &jsonapi.Links{
		"self": jsonapi.Link{
			Href: fmt.Sprintf("/%v/%v", name, id)},
	}
}

// JSONApiRefLink ref link helper
func JSONApiRefLink(i interface{}, rel string) *jsonapi.Links {
	fields := flect.Fields(reflect.ValueOf(i))
	name := reflect.TypeOf(i).Name()
	id := reflect.ValueOf(i).FieldByName("ID").Interface()
	for _, f := range fields {
		value, opts, ok := flect.TagValues(f.Tag, "jsonapi")
		if ok && value == "relation" && opts.HasOption(rel) {
			if _, _, ok := flect.TagValues(f.Tag, "link"); ok {
				return &jsonapi.Links{
					"related": fmt.Sprintf(
						"/%v/%v/%v",
						strings.ToLower(inflection.Plural(name)),
						id,
						strings.ToLower(inflection.Plural(f.Name)),
					),
				}
			}
		}
	}
	return &jsonapi.Links{}
}
<commit_msg>Make FieldOption Label attribute omitempty <commit_after>package meta

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/eyecuelab/kit/flect"
	"github.com/google/jsonapi"
	"github.com/jinzhu/inflection"
	"github.com/spf13/viper"
)

type (
	method       string
	inputType    string
	ActionHolder []*JsonAPIAction

	Extendable interface {
		Extend() error
	}

	JsonAPIAction struct {
		Method string         `json:"method"`
		Name   string         `json:"name"`
		URL    string         `json:"url"`
		Multi  bool           `json:"multi"`
		Fields []JsonAPIField `json:"fields"`
	}

	JsonAPIField struct {
		Name      string        `json:"name"`
		InputType string        `json:"type"`
		Value     interface{}   `json:"value,omitempty"`
		Required  bool          `json:"required"`
		Options   []FieldOption `json:"options,omitempty"`
		Data      *Pagination   `json:"data,omitempty"`
	}

	FieldOption struct {
		Label string                 `json:"label,omitempty"`
		Value interface{}            `json:"value"`
		Meta  map[string]interface{} `json:"meta,omitempty"`
	}

	Pagination struct {
		Count int `json:"item_count"`
		Max   int `json:"max"`
		Page  int `json:"page"`
	}
)

const (
	DELETE method = "DELETE"
	GET    method = "GET"
	POST   method = "POST"
	PATCH  method = "PATCH"
	PUT    method = "PUT"

	InputText   inputType = "text"
	InputPass   inputType = "password"
	InputBool   inputType = "bool"
	InputNumber inputType = "number"
	InputSelect inputType = "select"
)

// AddAction ...
func (ah *ActionHolder) AddAction(m method, name, urlHelper string, params ...interface{}) *JsonAPIAction {
	a := JsonAPIAction{
		Method: string(m),
		Name:   name,
		URL:    APIURL(fmt.Sprintf(urlHelper, params...)),
		Fields: make([]JsonAPIField, 0),
	}
	*ah = append(*ah, &a)
	return &a
}

// Field ...
func (a *JsonAPIAction) Field(name string, inputType inputType, value interface{}, requred bool) *JsonAPIAction {
	f := JsonAPIField{
		Name:      name,
		InputType: string(inputType),
		Value:     value,
		Required:  requred,
	}
	a.Fields = append(a.Fields, f)
	return a
}

// FieldWithOpts add field with options to action
func (a *JsonAPIAction) FieldWithOpts(name string, inputType inputType, value interface{}, requred bool, options []FieldOption) *JsonAPIAction {
	f := JsonAPIField{
		Name:      name,
		InputType: string(inputType),
		Value:     value,
		Required:  requred,
		Options:   options,
	}
	a.Fields = append(a.Fields, f)
	return a
}

// Pagination add pagination meta to action
func (a *JsonAPIAction) Pagination(data *Pagination) *JsonAPIAction {
	f := JsonAPIField{
		Name:      "page",
		InputType: string(InputNumber),
		Value:     data.Page,
		Required:  false,
		Data:      data,
	}
	a.Fields = append(a.Fields, f)
	return a
}

// RenderActions ...
func (ah *ActionHolder) RenderActions() *jsonapi.Meta {
	return &jsonapi.Meta{"actions": *ah}
}

// APIURL full api url for the path
func APIURL(path string) string {
	return fmt.Sprintf("%s%s", viper.GetString("root_url"), path)
}

// JSONApiSelfLink self link helper
func JSONApiSelfLink(i interface{}) *jsonapi.Links {
	name := reflect.TypeOf(i).Name()
	name = inflection.Plural(strings.ToLower(name))

	id := reflect.ValueOf(i).FieldByName("ID").Interface()

	return &jsonapi.Links{
		"self": jsonapi.Link{
			Href: fmt.Sprintf("/%v/%v", name, id)},
	}
}

// JSONApiRefLink ref link helper
func JSONApiRefLink(i interface{}, rel string) *jsonapi.Links {
	fields := flect.Fields(reflect.ValueOf(i))
	name := reflect.TypeOf(i).Name()
	id := reflect.ValueOf(i).FieldByName("ID").Interface()
	for _, f := range fields {
		value, opts, ok := flect.TagValues(f.Tag, "jsonapi")
		if ok && value == "relation" && opts.HasOption(rel) {
			if _, _, ok := flect.TagValues(f.Tag, "link"); ok {
				return &jsonapi.Links{
					"related": fmt.Sprintf(
						"/%v/%v/%v",
						strings.ToLower(inflection.Plural(name)),
						id,
						strings.ToLower(inflection.Plural(f.Name)),
					),
				}
			}
		}
	}
	return &jsonapi.Links{}
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strconv"

	apiv1 "cloud-android-orchestration/api/v1"

	"github.com/gorilla/mux"
)

// The controller implements the web API of the cloud orchestrator. It parses
// and validates requests from the client and passes the information to the
// relevant modules
type Controller struct {
	instanceManager InstanceManager
	sigServer       SignalingServer
	accountManager  AccountManager
}

var DEFAULT_INFRA_CONFIG = apiv1.InfraConfig{
	IceServers: []apiv1.IceServer{
		apiv1.IceServer{URLs: []string{"stun:stun.l.google.com:19302"}},
	},
}

func NewController(im InstanceManager, ss SignalingServer, am AccountManager) *Controller {
	controller := &Controller{im, ss, am}
	controller.SetupRoutes()

	return controller
}

func (c *Controller) ListenAndServe(addr string, handler http.Handler) error {
	return http.ListenAndServe(addr, handler)
}

func (c *Controller) SetupRoutes() {
	router := mux.NewRouter()

	// Signaling Server Routes
	router.Handle("/v1/connections/{connId}/messages", HTTPHandler(c.accountManager.Authenticate(c.Messages))).Methods("GET")
	router.Handle("/v1/connections/{connId}/:forward", HTTPHandler(c.accountManager.Authenticate(c.Forward))).Methods("POST")
	router.Handle("/v1/connections", HTTPHandler(c.accountManager.Authenticate(c.CreateConnection))).Methods("POST")
	router.Handle("/v1/devices/{deviceId}/files{path:/.+}", HTTPHandler(c.accountManager.Authenticate(c.GetDeviceFiles))).Methods("GET")

	// Global routes
	router.HandleFunc("/infra_config", func(w http.ResponseWriter, r *http.Request) {
		// TODO(b/220891296): Make this configurable
		replyJSON(w, DEFAULT_INFRA_CONFIG, http.StatusOK)
	}).Methods("GET")

	router.Handle("/", HTTPHandler(c.accountManager.Authenticate(indexHandler)))

	http.Handle("/", router)
}

// Intercept errors returned by the HTTPHandler and transform them into HTTP
// error responses
func (fn HTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Println(r.Method, " ", r.URL, " ", r.RemoteAddr)
	if err := fn(w, r); err != nil {
		log.Println("Error: ", err)
		var e *AppError
		if errors.As(err, &e) {
			replyJSON(w, e.JSONResponse(), e.StatusCode)
		} else {
			replyJSON(w, apiv1.ErrorMsg{Error: "Internal Server Error"}, http.StatusInternalServerError)
		}
	}
}

func (c *Controller) GetDeviceFiles(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	devId := mux.Vars(r)["deviceId"]
	path := mux.Vars(r)["path"]
	return c.sigServer.ServeDeviceFiles(DeviceFilesRequest{devId, path, w, r}, user)
}

func (c *Controller) CreateConnection(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	var msg apiv1.NewConnMsg
	err := json.NewDecoder(r.Body).Decode(&msg)
	if err != nil {
		return NewBadRequestError("Malformed JSON in request", err)
	}
	log.Println("id: ", msg.DeviceId)
	reply, err := c.sigServer.NewConnection(msg, user)
	if err != nil {
		return fmt.Errorf("Failed to communicate with device: %w", err)
	}
	replyJSON(w, reply.Response, reply.StatusCode)
	return nil
}

func (c *Controller) Messages(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	id := mux.Vars(r)["connId"]
	start, err := intFormValue(r, "start", 0)
	if err != nil {
		return NewBadRequestError("Invalid value for start field", err)
	}
	// -1 means all messages
	count, err := intFormValue(r, "count", -1)
	if err != nil {
		return NewBadRequestError("Invalid value for count field", err)
	}
	reply, err := c.sigServer.Messages(id, start, count, user)
	if err != nil {
		return fmt.Errorf("Failed to get messages: %w", err)
	}
	replyJSON(w, reply.Response, reply.StatusCode)
	return nil
}

func (c *Controller) Forward(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	id := mux.Vars(r)["connId"]
	var msg apiv1.ForwardMsg
	err := json.NewDecoder(r.Body).Decode(&msg)
	if err != nil {
		return NewBadRequestError("Malformed JSON in request", err)
	}
	reply, err := c.sigServer.Forward(id, msg, user)
	if err != nil {
		return fmt.Errorf("Failed to send message to device: %w", err)
	}
	replyJSON(w, reply.Response, reply.StatusCode)
	return nil
}

func indexHandler(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	fmt.Fprintln(w, "Home page")
	return nil
}

// Get an int form value from the request. A form value can be in the query
// string or in the request body in either URL or Multipart encoding. If there
// is no form parameter with the given name the default value is returned. If a
// parameter with that name exists but the value can't be converted to an int an
// error is returned.
func intFormValue(r *http.Request, name string, def int) (int, error) {
	str := r.FormValue(name)
	if str == "" {
		return def, nil
	}
	i, e := strconv.Atoi(str)
	if e != nil {
		return 0, fmt.Errorf("Invalid %s value: %s", name, str)
	}
	return i, nil
}

// Send a JSON http response to the client
func replyJSON(w http.ResponseWriter, obj interface{}, statusCode int) error {
	if statusCode != http.StatusOK {
		w.WriteHeader(statusCode)
	}
	w.Header().Set("Content-Type", "application/json")
	encoder := json.NewEncoder(w)
	return encoder.Encode(obj)
}
<commit_msg>Serve infra config from /v1/ too<commit_after>// Copyright 2022 Google LLC
//
// 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
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strconv"

	apiv1 "cloud-android-orchestration/api/v1"

	"github.com/gorilla/mux"
)

// The controller implements the web API of the cloud orchestrator. It parses
// and validates requests from the client and passes the information to the
// relevant modules
type Controller struct {
	instanceManager InstanceManager
	sigServer       SignalingServer
	accountManager  AccountManager
}

var DEFAULT_INFRA_CONFIG = apiv1.InfraConfig{
	IceServers: []apiv1.IceServer{
		apiv1.IceServer{URLs: []string{"stun:stun.l.google.com:19302"}},
	},
}

func NewController(im InstanceManager, ss SignalingServer, am AccountManager) *Controller {
	controller := &Controller{im, ss, am}
	controller.SetupRoutes()

	return controller
}

func (c *Controller) ListenAndServe(addr string, handler http.Handler) error {
	return http.ListenAndServe(addr, handler)
}

func (c *Controller) SetupRoutes() {
	router := mux.NewRouter()

	// Signaling Server Routes
	router.Handle("/v1/connections/{connId}/messages", HTTPHandler(c.accountManager.Authenticate(c.Messages))).Methods("GET")
	router.Handle("/v1/connections/{connId}/:forward", HTTPHandler(c.accountManager.Authenticate(c.Forward))).Methods("POST")
	router.Handle("/v1/connections", HTTPHandler(c.accountManager.Authenticate(c.CreateConnection))).Methods("POST")
	router.Handle("/v1/devices/{deviceId}/files{path:/.+}", HTTPHandler(c.accountManager.Authenticate(c.GetDeviceFiles))).Methods("GET")

	// Infra route
	router.HandleFunc("/v1/infra_config", func(w http.ResponseWriter, r *http.Request) {
		// TODO(b/220891296): Make this configurable
		replyJSON(w, DEFAULT_INFRA_CONFIG, http.StatusOK)
	}).Methods("GET")

	// Global routes
	router.Handle("/", HTTPHandler(c.accountManager.Authenticate(indexHandler)))

	http.Handle("/", router)
}

// Intercept errors returned by the HTTPHandler and transform them into HTTP
// error responses
func (fn HTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Println(r.Method, " ", r.URL, " ", r.RemoteAddr)
	if err := fn(w, r); err != nil {
		log.Println("Error: ", err)
		var e *AppError
		if errors.As(err, &e) {
			replyJSON(w, e.JSONResponse(), e.StatusCode)
		} else {
			replyJSON(w, apiv1.ErrorMsg{Error: "Internal Server Error"}, http.StatusInternalServerError)
		}
	}
}

func (c *Controller) GetDeviceFiles(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	devId := mux.Vars(r)["deviceId"]
	path := mux.Vars(r)["path"]
	return c.sigServer.ServeDeviceFiles(DeviceFilesRequest{devId, path, w, r}, user)
}

func (c *Controller) CreateConnection(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	var msg apiv1.NewConnMsg
	err := json.NewDecoder(r.Body).Decode(&msg)
	if err != nil {
		return NewBadRequestError("Malformed JSON in request", err)
	}
	log.Println("id: ", msg.DeviceId)
	reply, err := c.sigServer.NewConnection(msg, user)
	if err != nil {
		return fmt.Errorf("Failed to communicate with device: %w", err)
	}
	replyJSON(w, reply.Response, reply.StatusCode)
	return nil
}

func (c *Controller) Messages(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	id := mux.Vars(r)["connId"]
	start, err := intFormValue(r, "start", 0)
	if err != nil {
		return NewBadRequestError("Invalid value for start field", err)
	}
	// -1 means all messages
	count, err := intFormValue(r, "count", -1)
	if err != nil {
		return NewBadRequestError("Invalid value for count field", err)
	}
	reply, err := c.sigServer.Messages(id, start, count, user)
	if err != nil {
		return fmt.Errorf("Failed to get messages: %w", err)
	}
	replyJSON(w, reply.Response, reply.StatusCode)
	return nil
}

func (c *Controller) Forward(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	id := mux.Vars(r)["connId"]
	var msg apiv1.ForwardMsg
	err := json.NewDecoder(r.Body).Decode(&msg)
	if err != nil {
		return NewBadRequestError("Malformed JSON in request", err)
	}
	reply, err := c.sigServer.Forward(id, msg, user)
	if err != nil {
		return fmt.Errorf("Failed to send message to device: %w", err)
	}
	replyJSON(w, reply.Response, reply.StatusCode)
	return nil
}

func indexHandler(w http.ResponseWriter, r *http.Request, user UserInfo) error {
	fmt.Fprintln(w, "Home page")
	return nil
}

// Get an int form value from the request. A form value can be in the query
// string or in the request body in either URL or Multipart encoding. If there
// is no form parameter with the given name the default value is returned. If a
// parameter with that name exists but the value can't be converted to an int an
// error is returned.
func intFormValue(r *http.Request, name string, def int) (int, error) {
	str := r.FormValue(name)
	if str == "" {
		return def, nil
	}
	i, e := strconv.Atoi(str)
	if e != nil {
		return 0, fmt.Errorf("Invalid %s value: %s", name, str)
	}
	return i, nil
}

// Send a JSON http response to the client
func replyJSON(w http.ResponseWriter, obj interface{}, statusCode int) error {
	if statusCode != http.StatusOK {
		w.WriteHeader(statusCode)
	}
	w.Header().Set("Content-Type", "application/json")
	encoder := json.NewEncoder(w)
	return encoder.Encode(obj)
}
<|endoftext|>


<commit_before>package gogo

type ControllerID interface {
	ID() string
}

type ControllerDispatch interface {
	DISPATCH(c *Context)
}

type ControllerIndex interface {
	Index(c *Context)
}

type ControllerShow interface {
	Show(c *Context)
}

type ControllerCreate interface {
	Create(c *Context)
}

type ControllerDestroy interface {
	Destroy(c *Context)
}

type ControllerUpdate interface {
	Update(c *Context)
}

type ControllerExplore interface {
	Explore(c *Context)
}
<commit_msg>What: add comments to controller interfaces<commit_after>package gogo

// ControllerDispatch is an interface that allows custom router for Resource.
type ControllerDispatch interface {
	DISPATCH(c *Context)
}

// ControllerID is an interface that allows custom resource id name with Resource.
type ControllerID interface {
	ID() string
}

// ControllerIndex is an interface that wraps the Index method. It is used by Resource
// for registering GET /resource route handler.
type ControllerIndex interface {
	Index(c *Context)
}

// ControllerShow is an interface that wraps the Show method. It is used by Resource
// for registering GET /resource/:id route handler.
type ControllerShow interface {
	Show(c *Context)
}

// ControllerCreate is an interface that wraps the Create method. It is used by Resource
// for registering POST /resource route handler.
type ControllerCreate interface {
	Create(c *Context)
}

// ControllerUpdate is an interface that wraps the Update method. It is used by Resource
// for registering PUT /resource/:id route handler.
type ControllerUpdate interface {
	Update(c *Context)
}

// ControllerExplore is an interface that wraps the Explore method. It is used by Resource
// for registering HEAD /resource/:id route handler.
type ControllerExplore interface {
	Explore(c *Context)
}

// ControllerDestroy is an interface that wraps the Destroy method. It is used by Resource
// for registering DELETE /resource/:id route handler.
type ControllerDestroy interface {
	Destroy(c *Context)
}
<|endoftext|>


<commit_before>package key

import (
	"fmt"
	"github.com/globocom/config"
	fstesting "github.com/globocom/tsuru/fs/testing"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"os"
	"path"
	"strings"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	origKeyFile string
	rfs         *fstesting.RecordingFs
}

var _ = Suite(&S{})

func (s *S) authKeysContent(c *C) string {
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	return string(b)
}

func (s *S) SetUpSuite(c *C) {
	err := config.ReadConfigFile("../etc/gandalf.conf")
	c.Check(err, IsNil)
}

func (s *S) SetUpTest(c *C) {
	s.rfs = &fstesting.RecordingFs{}
	fsystem = s.rfs
}

func (s *S) TearDownSuite(c *C) {
	fsystem = nil
}

func (s *S) TearDownTest(c *C) {
	ok := s.clearAuthKeyFile()
	c.Assert(ok, Equals, true)
}

func (s *S) clearAuthKeyFile() bool {
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	if err != nil {
		return false
	}
	if err := f.Truncate(0); err != nil {
		return false
	}
	return true
}

func (s *S) TestAuthKeysShouldBeAbsolutePathToUsersAuthorizedKeysByDefault(c *C) {
	home := os.Getenv("HOME")
	expected := path.Join(home, ".ssh", "authorized_keys")
	c.Assert(authKey, Equals, expected)
}

func (s *S) TestShouldAddKeyWithoutError(c *C) {
	key := "somekey blaaaaaaa r2d2@host"
	err := Add(key, "someuser", s.rfs)
	c.Assert(err, IsNil)
}

func (s *S) TestShouldWriteKeyInFile(c *C) {
	key := "somekey blaaaaaaa r2d2@host"
	err := Add(key, "someuser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	c.Assert(got, Equals, formatKey(key, "someuser"))
}

func (s *S) TestShouldAppendKeyInFile(c *C) {
	key1 := "somekey blaaaaaaa r2d2@host"
	err := Add(key1, "someuser", s.rfs)
	c.Assert(err, IsNil)
	key2 := "someotherkey fooo r2d2@host"
	err = Add(key2, "someuser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	expected := fmt.Sprintf(".*%s\n.*%s", key1, key2)
	c.Assert(got, Matches, expected)
}

func (s *S) TestAddShouldWrapKeyWithRestrictions(c *C) {
	key := "somekey bleeeerh r2d2@host"
	expected := fmt.Sprintf("no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command=.* %s", key)
	err := Add(key, "someuser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	c.Assert(got, Matches, expected)
}

func (s *S) TestBulkActionAddKeys(c *C) {
	key1 := "ssh-rsa mykey pippin@nowhere"
	key2 := "ssh-rsa myotherkey pippin@somewhere"
	keys := []string{key1, key2}
	err := bulkAction(Add, keys, "someuser", s.rfs)
	c.Assert(err, IsNil)
	got := strings.Replace(s.authKeysContent(c), "\n", " ", -1)
	c.Assert(got, Matches, ".*"+key1+".*")
	c.Assert(got, Matches, ".*"+key2+".*")
}

func (s *S) TestBulkActionRemoveKeys(c *C) {
	key1 := "ssh-rsa mykey pippin@nowhere"
	key2 := "ssh-rsa myotherkey pippin@somewhere"
	keys := []string{key1, key2}
	err := bulkAction(Add, keys, "someuser", s.rfs)
	got := strings.Replace(s.authKeysContent(c), "\n", " ", -1)
	c.Assert(got, Matches, ".*"+key1+".*")
	c.Assert(got, Matches, ".*"+key2+".*")
	err = bulkAction(Remove, keys, "someuser", s.rfs)
	c.Assert(err, IsNil)
	got = strings.Replace(s.authKeysContent(c), "\n", " ", -1)
	c.Assert(got, Not(Matches), ".*"+key1+".*")
	c.Assert(got, Not(Matches), ".*"+key2+".*")
}

func (s *S) TestBulkAddShouldWriteToAuthorizedKeysFile(c *C) {
	err := BulkAdd([]string{"ssh-rsa mykey pippin@nowhere"}, "someuser", s.rfs)
	c.Assert(err, IsNil)
	keys := s.authKeysContent(c)
	c.Assert(keys, Matches, ".*ssh-rsa mykey pippin@nowhere")
}

func (s *S) TestBulkRemoveShouldRemoveKeysFromAuthorizedKeys(c *C) {
	key := "ssh-rsa mykey pippin@nowhere"
	err := BulkRemove([]string{key}, "someuser", s.rfs)
	c.Assert(err, IsNil)
	keys := s.authKeysContent(c)
	c.Assert(keys, Equals, "")
}

func (s *S) TestRemoveKey(c *C) {
	key1 := "somekey blaaaaaaa r2d2@host"
	err := Add(key1, "someuser", s.rfs)
	c.Assert(err, IsNil)
	key2 := "someotherkey fooo r2d2@host"
	err = Add(key2, "someuser", s.rfs)
	c.Assert(err, IsNil)
	err = Remove(key1, "someuser", s.rfs)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	expected := fmt.Sprintf(".*%s", key2)
	c.Assert(got, Matches, expected)
}

func (s *S) TestRemoveWhenKeyDoesNotExists(c *C) {
	key1 := "somekey blaaaaaaa r2d2@host"
	err := Remove(key1, "anotheruser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	c.Assert(got, Equals, "")
}

func (s *S) TestRemoveWhenExistsOnlyOneKey(c *C) {
	key1 := "somekey blaaaaaaa r2d2@host"
	err := Add(key1, "someuser", s.rfs)
	c.Assert(err, IsNil)
	err = Remove(key1, "someuser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	c.Assert(got, Equals, "")
}

func (s *S) TestFormatKeyShouldAddSshLoginRestrictionsAtBegining(c *C) {
	key := "somekeeey fooo bar@bar.com"
	got := formatKey(key, "brain")
	expected := fmt.Sprintf("no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command=.* %s", key)
	c.Assert(got, Matches, expected)
}

func (s *S) TestFormatKeyShouldAddCommandAfterSshRestrictions(c *C) {
	key := "somekeyyyy fooow bar@bar.com"
	got := formatKey(key, "brain")
	p, err := config.GetString("bin-path")
	c.Assert(err, IsNil)
	expected := fmt.Sprintf(`no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command="%s brain" %s`, p, key)
	c.Assert(got, Equals, expected)
}

func (s *S) TestFormatKeyShouldGetCommandPathFromGandalfConf(c *C) {
	oldConf, err := config.GetString("bin-path")
	c.Assert(err, IsNil)
	config.Set("bin-path", "/foo/bar/hi.go")
	defer config.Set("bin-path", oldConf)
	key := "lol loool bar@bar.com"
	got := formatKey(key, "dash")
	expected := fmt.Sprintf(`no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command="/foo/bar/hi.go dash" %s`, key)
	c.Assert(got, Equals, expected)
}

func (s *S) TestFormatKeyShouldAppendUserNameAsCommandParameter(c *C) {
	key := "ssh-rsa fueeel bar@bar.com"
	p, err := config.GetString("bin-path")
	c.Assert(err, IsNil)
	got := formatKey(key, "someuser")
	expected := fmt.Sprintf(`no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command="%s someuser" %s`, p, key)
	c.Assert(got, Equals, expected)
}
<commit_msg>key: minor refactor in tests<commit_after>package key

import (
	"fmt"
	"github.com/globocom/config"
	fstesting "github.com/globocom/tsuru/fs/testing"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"os"
	"path"
	"strings"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	origKeyFile string
	rfs         *fstesting.RecordingFs
}

var _ = Suite(&S{})

func (s *S) authKeysContent(c *C) string {
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	return string(b)
}

func (s *S) SetUpSuite(c *C) {
	err := config.ReadConfigFile("../etc/gandalf.conf")
	c.Check(err, IsNil)
}

func (s *S) SetUpTest(c *C) {
	s.rfs = &fstesting.RecordingFs{}
	fsystem = s.rfs
}

func (s *S) TearDownSuite(c *C) {
	fsystem = nil
}

func (s *S) TearDownTest(c *C) {
	ok := s.clearAuthKeyFile()
	c.Assert(ok, Equals, true)
}

func (s *S) clearAuthKeyFile() bool {
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	if err != nil {
		return false
	}
	if err := f.Truncate(0); err != nil {
		return false
	}
	return true
}

func (s *S) TestAuthKeysShouldBeAbsolutePathToUsersAuthorizedKeysByDefault(c *C) {
	home := os.Getenv("HOME")
	expected := path.Join(home, ".ssh", "authorized_keys")
	c.Assert(authKey, Equals, expected)
}

func (s *S) TestShouldAddKeyWithoutError(c *C) {
	key := "somekey blaaaaaaa r2d2@host"
	err := Add(key, "someuser", s.rfs)
	c.Assert(err, IsNil)
}

func (s *S) TestShouldWriteKeyInFile(c *C) {
	key := "somekey blaaaaaaa r2d2@host"
	err := Add(key, "someuser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	c.Assert(got, Equals, formatKey(key, "someuser"))
}

func (s *S) TestShouldAppendKeyInFile(c *C) {
	key1 := "somekey blaaaaaaa r2d2@host"
	err := Add(key1, "someuser", s.rfs)
	c.Assert(err, IsNil)
	key2 := "someotherkey fooo r2d2@host"
	err = Add(key2, "someuser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	expected := fmt.Sprintf(".*%s\n.*%s", key1, key2)
	c.Assert(got, Matches, expected)
}

func (s *S) TestAddShouldWrapKeyWithRestrictions(c *C) {
	key := "somekey bleeeerh r2d2@host"
	expected := fmt.Sprintf("no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command=.* %s", key)
	err := Add(key, "someuser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	c.Assert(got, Matches, expected)
}

func (s *S) TestBulkActionAddKeys(c *C) {
	key1 := "ssh-rsa mykey pippin@nowhere"
	key2 := "ssh-rsa myotherkey pippin@somewhere"
	keys := []string{key1, key2}
	err := bulkAction(Add, keys, "someuser", s.rfs)
	c.Assert(err, IsNil)
	got := strings.Replace(s.authKeysContent(c), "\n", " ", -1)
	c.Assert(got, Matches, ".*"+key1+".*")
	c.Assert(got, Matches, ".*"+key2+".*")
}

func (s *S) TestBulkActionRemoveKeys(c *C) {
	key1 := "ssh-rsa mykey pippin@nowhere"
	key2 := "ssh-rsa myotherkey pippin@somewhere"
	keys := []string{key1, key2}
	err := bulkAction(Add, keys, "someuser", s.rfs)
	got := strings.Replace(s.authKeysContent(c), "\n", " ", -1)
	c.Assert(got, Matches, ".*"+key1+".*")
	c.Assert(got, Matches, ".*"+key2+".*")
	err = bulkAction(Remove, keys, "someuser", s.rfs)
	c.Assert(err, IsNil)
	got = strings.Replace(s.authKeysContent(c), "\n", " ", -1)
	c.Assert(got, Not(Matches), ".*"+key1+".*")
	c.Assert(got, Not(Matches), ".*"+key2+".*")
}

func (s *S) TestBulkAddShouldWriteToAuthorizedKeysFile(c *C) {
	err := BulkAdd([]string{"ssh-rsa mykey pippin@nowhere"}, "someuser", s.rfs)
	c.Assert(err, IsNil)
	keys := s.authKeysContent(c)
	c.Assert(keys, Matches, ".*ssh-rsa mykey pippin@nowhere")
}

func (s *S) TestBulkRemoveShouldRemoveKeysFromAuthorizedKeys(c *C) {
	key := "ssh-rsa mykey pippin@nowhere"
	err := BulkRemove([]string{key}, "someuser", s.rfs)
	c.Assert(err, IsNil)
	keys := s.authKeysContent(c)
	c.Assert(keys, Equals, "")
}

func (s *S) TestRemoveKey(c *C) {
	key1 := "somekey blaaaaaaa r2d2@host"
	err := Add(key1, "someuser", s.rfs)
	c.Assert(err, IsNil)
	key2 := "someotherkey fooo r2d2@host"
	err = Add(key2, "someuser", s.rfs)
	c.Assert(err, IsNil)
	err = Remove(key1, "someuser", s.rfs)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	expected := formatKey(key2, "someuser")
	c.Assert(got, Matches, expected)
	expected = formatKey(key1, "someuser")
	c.Assert(got, Not(Matches), expected)
}

func (s *S) TestRemoveWhenKeyDoesNotExists(c *C) {
	key1 := "somekey blaaaaaaa r2d2@host"
	err := Remove(key1, "anotheruser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	c.Assert(got, Equals, "")
}

func (s *S) TestRemoveWhenExistsOnlyOneKey(c *C) {
	key1 := "somekey blaaaaaaa r2d2@host"
	err := Add(key1, "someuser", s.rfs)
	c.Assert(err, IsNil)
	err = Remove(key1, "someuser", s.rfs)
	c.Assert(err, IsNil)
	f, err := s.rfs.OpenFile(authKey, os.O_RDWR, 0755)
	c.Assert(err, IsNil)
	b, err := ioutil.ReadAll(f)
	c.Assert(err, IsNil)
	got := string(b)
	c.Assert(got, Equals, "")
}

func (s *S) TestFormatKeyShouldAddSshLoginRestrictionsAtBegining(c *C) {
	key := "somekeeey fooo bar@bar.com"
	got := formatKey(key, "brain")
	expected := fmt.Sprintf("no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command=.* %s", key)
	c.Assert(got, Matches, expected)
}

func (s *S) TestFormatKeyShouldAddCommandAfterSshRestrictions(c *C) {
	key := "somekeyyyy fooow bar@bar.com"
	got := formatKey(key, "brain")
	p, err := config.GetString("bin-path")
	c.Assert(err, IsNil)
	expected := fmt.Sprintf(`no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command="%s brain" %s`, p, key)
	c.Assert(got, Equals, expected)
}

func (s *S) TestFormatKeyShouldGetCommandPathFromGandalfConf(c *C) {
	oldConf, err := config.GetString("bin-path")
	c.Assert(err, IsNil)
	config.Set("bin-path", "/foo/bar/hi.go")
	defer config.Set("bin-path", oldConf)
	key := "lol loool bar@bar.com"
	got := formatKey(key, "dash")
	expected := fmt.Sprintf(`no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command="/foo/bar/hi.go dash" %s`, key)
	c.Assert(got, Equals, expected)
}

func (s *S) TestFormatKeyShouldAppendUserNameAsCommandParameter(c *C) {
	key := "ssh-rsa fueeel bar@bar.com"
	p, err := config.GetString("bin-path")
	c.Assert(err, IsNil)
	got := formatKey(key, "someuser")
	expected := fmt.Sprintf(`no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty,command="%s someuser" %s`, p, key)
	c.Assert(got, Equals, expected)
}
<|endoftext|>


<commit_before>package runtime

import (
	"context"
	"io"
	"strings"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/concourse/atc/compression"
	"github.com/concourse/concourse/atc/db"
	"go.opentelemetry.io/otel/propagation"
)

type Worker interface {
	Name() string
	FindOrCreateContainer(context.Context, db.ContainerOwner, db.ContainerMetadata, ContainerSpec) (Container, []VolumeMount, error)
	CreateVolumeForArtifact(logger lager.Logger, teamID int) (Volume, db.WorkerArtifact, error)

	LookupContainer(logger lager.Logger, handle string) (Container, bool, error)
	LookupVolume(logger lager.Logger, handle string) (Volume, bool, error)

	DBWorker() db.Worker
}

type Container interface {
	// Run starts a Process on the Container. If the executable (defined in
	// ProcessSpec) does not exist, an ExecutableNotFound error must be
	// returned.
	//
	// The passed in Context only applies to the request to start the Process,
	// so cancelling the Context should not cancel the Process. To cancel the
	// Process, use the Context passed into Process.Wait.
	Run(context.Context, ProcessSpec, ProcessIO) (Process, error)
	// Attach attempts to attach to an existing process by its ID.
	//
	// The passed in Context only applies to the request to attach to the
	// Process, so cancelling the Context should not cancel the Process. To
	// cancel the Process, use the Context passed into Process.Wait.
	Attach(context.Context, string, ProcessIO) (Process, error)

	Properties() (map[string]string, error)
	SetProperty(name string, value string) error

	DBContainer() db.CreatedContainer
}

type ContainerSpec struct {
	TeamID   int
	TeamName string
	JobID    int
	StepName string

	ImageSpec ImageSpec
	Env       []string
	Type      db.ContainerType

	// Working directory for processes run in the container. Must be an absolute path.
	Dir string

	// Inputs to provide to the container. Inputs with a volume local to the
	// selected worker will be made available via a COW volume; others will be
	// streamed.
	Inputs []Input

	// List of cached container paths to re-mount to the volume. If the cached
	// volume doesn't yet exist, an empty volume should be mounted.
	//
	// Paths may be relative (to Dir) or absolute.
	Caches []string

	// Outputs for which volumes should be created and mounted into the container.
	Outputs OutputPaths

	// Resource limits to be set on the container.
	Limits ContainerLimits

	// Whether or not to mount the Certs volume onto the container.
	// TODO: find a more general pattern here
	CertsBindMount bool
}

// ContainerSpec must implement propagation.TextMapCarrier so that it can be
// used to pass tracing metadata to the containers via env vars (specifically,
// the TRACEPARENT var)
var _ propagation.TextMapCarrier = new(ContainerSpec)

func (cs *ContainerSpec) Get(key string) string {
	for _, env := range cs.Env {
		assignment := strings.SplitN("=", env, 2)
		if assignment[0] == strings.ToUpper(key) {
			return assignment[1]
		}
	}
	return ""
}

func (cs *ContainerSpec) Set(key string, value string) {
	varName := strings.ToUpper(key)
	envVar := varName + "=" + value
	for i, env := range cs.Env {
		if strings.SplitN("=", env, 2)[0] == varName {
			cs.Env[i] = envVar
			return
		}
	}
	cs.Env = append(cs.Env, envVar)
}

func (cs *ContainerSpec) Keys() []string {
	// this implementation isn't technically correct - it gives all environment
	// vars as the list of keys (rather than just those set by Set), and it
	// assumes the original keys were all lowercased). from what I can tell,
	// though, this Keys method isn't currently even used, so this doesn't
	// matter right now (but may in the future...)
	keys := make([]string, len(cs.Env))
	for i, env := range cs.Env {
		envName := strings.SplitN("=", env, 2)[0]
		keys[i] = strings.ToLower(envName)
	}
	return keys
}

type ProcessSpec struct {
	// ID is optional - if set, the process can be Attach'd to using this ID.
	// If unset, an ID will be generated, and can be retrieved using
	// Process.ID().
	ID string

	Path string
	Args []string
	Env  []string
	Dir  string
	User string
	TTY  *TTYSpec
}

type TTYSpec struct {
	WindowSize WindowSize
}

type WindowSize struct {
	Columns int
	Rows    int
}

type ProcessIO struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

type Process interface {
	ID() string
	Wait(context.Context) (ProcessResult, error)
	SetTTY(tty TTYSpec) error
}

type ProcessResult struct {
	ExitStatus int
}

type Input struct {
	Volume          Volume
	DestinationPath string
}

// OutputPaths is a mapping from output name to its path in the container.
type OutputPaths map[string]string

type ImageSpec struct {
	ResourceType string
	ImageURL     string
	ImageVolume  Volume
	Privileged   bool
}

type ContainerLimits struct {
	CPU    *uint64
	Memory *uint64
}

type Volume interface {
	Handle() string
	StreamIn(ctx context.Context, path string, compression compression.Compression, reader io.Reader) error
	StreamOut(ctx context.Context, path string, compression compression.Compression) (io.ReadCloser, error)

	InitializeResourceCache(lager.Logger, db.UsedResourceCache) error
	InitializeTaskCache(logger lager.Logger, jobID int, stepName string, path string, privileged bool) error

	DBVolume() db.CreatedVolume
}

type P2PVolume interface {
	Handle() string
	GetStreamInP2PURL(ctx context.Context, path string) (string, error)
	StreamP2POut(ctx context.Context, path string, destURL string, compression compression.Compression) error
}

type VolumeMount struct {
	MountPath string
	Volume    Volume
}
<commit_msg>S: add documentation for runtime types<commit_after>package runtime

import (
	"context"
	"io"
	"strings"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/concourse/atc/compression"
	"github.com/concourse/concourse/atc/db"
	"go.opentelemetry.io/otel/propagation"
)

// Worker represents an entity that aggregates Containers and Volumes, and
// provides a way to construct new/lookup existing Containers and Volumes.
// Depending on the runtime implementation, a Worker may be a single VM, or may
// be a Kubernetes/Nomad cluster, etc.
type Worker interface {
	// Name gives the name of the Worker. The Name is used to uniquely identify
	// a Worker, so no two Workers may share a Name.
	Name() string

	// FindOrCreateContainer looks for a Container specified by the provided
	// db.ContainerOwner, and if none exists, creates one. This also includes
	// finding/creating the associated Volumes and properly mounting them to
	// the Container.
	//
	// It can be thought of as declaratively saying "I want a container
	// matching these specifications" - and the Worker implementation will
	// "make it so", regardless of what Containers/Volumes already exist.
	FindOrCreateContainer(context.Context, db.ContainerOwner, db.ContainerMetadata, ContainerSpec) (Container, []VolumeMount, error)
	// CreateVolumeForArtifact creates a new empty Volume to be used as a
	// WorkerArtifact. This is used for uploading local inputs to a worker via
	// `fly execute -i ...`.
	CreateVolumeForArtifact(logger lager.Logger, teamID int) (Volume, db.WorkerArtifact, error)

	// LookupContainer finds the Container on the Worker by its handle, if it
	// exists.
	LookupContainer(logger lager.Logger, handle string) (Container, bool, error)
	// LookupVolume finds the Volume on the Worker by its handle, if it exists.
	LookupVolume(logger lager.Logger, handle string) (Volume, bool, error)

	DBWorker() db.Worker
}

// Container represents an entity that can execute Processes. It need not
// represent a "true" container - e.g. a Container may represent a Garden
// Container, or it may represent a Kubernetes Pod, or a Nomad Job (or Task?),
// etc.
type Container interface {
	// Run starts a Process on the Container. If the executable (defined in
	// ProcessSpec) does not exist, an ExecutableNotFound error must be
	// returned.
	//
	// The passed in Context only applies to the request to start the Process,
	// so cancelling the Context should not cancel the Process. To cancel the
	// Process, use the Context passed into Process.Wait.
	Run(context.Context, ProcessSpec, ProcessIO) (Process, error)
	// Attach attempts to attach to an existing process by its ID. An error
	// indicates either the ID doesn't correspond to a running process, or that
	// something unexpected went wrong.
	//
	// The passed in Context only applies to the request to attach to the
	// Process, so cancelling the Context should not cancel the Process. To
	// cancel the Process, use the Context passed into Process.Wait.
	Attach(context.Context, string, ProcessIO) (Process, error)

	// Properties gives the key/value pairs that were ascribed to the
	// Container, either via SetProperty() or at creation time.
	Properties() (map[string]string, error)
	// SetProperty adds a new key/value pair to the Container's Properties.
	SetProperty(name string, value string) error

	DBContainer() db.CreatedContainer
}

// ContainerSpec defines how to construct a container.
type ContainerSpec struct {
	// TeamID identifies the team to which the Container belongs.
	TeamID int
	// TeamName is the name of the team to which the Container belongs.
	TeamName string
	// JobID identifies the job in which the Container is running, used for
	// identifying task caches.
	JobID int
	// StepName is the name of the task step, used for identifying task caches.
	// If the Container is not for a task step, this may be left empty.
	StepName string

	// ImageSpec defines where the container image should come from.
	ImageSpec ImageSpec
	// Env is a list of environment variables (of the form "NAME=VALUE") to use
	// for running Processes. These are in addition to, but take precedence
	// over, those defined in the container image's metdata file.
	Env []string
	// Type is the type of step the Container is for (e.g. task, get, etc.).
	Type db.ContainerType

	// Dir is the working directory for processes run in the container. Must be
	// an absolute path.
	Dir string

	// Inputs defines the inputs to provide to the container. The Volumes
	// pointed to in these inputs may be on another worker, in which case they
	// must be streamed.
	Inputs []Input

	// Caches is a list of container paths to cache after execution and
	// re-mount to the volume. If the cached volume doesn't yet exist, an empty
	// volume should be mounted.
	//
	// Paths may be relative (to Dir) or absolute.
	Caches []string

	// Outputs defines a mapping of output names to paths within the container
	// for which volumes should be created and mounted.
	Outputs OutputPaths

	// Limits specifies resource limits to be set on the container.
	Limits ContainerLimits

	// CertsBindMount indicates whether or not to mount the worker's Certs
	// volume onto the container.
	CertsBindMount bool
}

// ContainerSpec must implement propagation.TextMapCarrier so that it can be
// used to pass tracing metadata to the containers via env vars (specifically,
// the TRACEPARENT var)
var _ propagation.TextMapCarrier = new(ContainerSpec)

func (cs *ContainerSpec) Get(key string) string {
	for _, env := range cs.Env {
		assignment := strings.SplitN("=", env, 2)
		if assignment[0] == strings.ToUpper(key) {
			return assignment[1]
		}
	}
	return ""
}

func (cs *ContainerSpec) Set(key string, value string) {
	varName := strings.ToUpper(key)
	envVar := varName + "=" + value
	for i, env := range cs.Env {
		if strings.SplitN("=", env, 2)[0] == varName {
			cs.Env[i] = envVar
			return
		}
	}
	cs.Env = append(cs.Env, envVar)
}

func (cs *ContainerSpec) Keys() []string {
	// this implementation isn't technically correct - it gives all environment
	// vars as the list of keys (rather than just those set by Set), and it
	// assumes the original keys were all lowercased). from what I can tell,
	// though, this Keys method isn't currently even used, so this doesn't
	// matter right now (but may in the future...)
	keys := make([]string, len(cs.Env))
	for i, env := range cs.Env {
		envName := strings.SplitN("=", env, 2)[0]
		keys[i] = strings.ToLower(envName)
	}
	return keys
}

// ProcessSpec defines how a process should be run in the Container.
type ProcessSpec struct {
	// ID is optional - if set, the process can be Attach'd to using this ID.
	// If unset, an ID will be generated, and can be retrieved using
	// Process.ID().
	ID string

	// Path is the executable to run (i.e. $0)
	Path string
	// Args is the list of command-line arguments to provide to the executable
	// (i.e. $1-$x)
	Args []string
	// Env is a list of environment variables (of the form "NAME=VALUE") to
	// provide to the running process. These are in addition to, but take
	// precedence over, those set in ContainerSpec.Env.
	Env []string
	// Dir is the working directory in which to run the process. This takes
	// precedence over ContainerSpec.Dir.
	Dir string
	// User is the username to run the process as. If unset, defaults to the
	// user specified by the container image's metadata.
	User string
	// TTY defines the TTY allocated to the process.
	TTY *TTYSpec
}

// TTYSpec defines the TTY allocated to a Process. It may be set when starting
// the Process via the ProcessSpec, or at runtime using Process.SetTTY().
type TTYSpec struct {
	WindowSize WindowSize
}

// WindowSize represents the size of the TTY allocated to the Process, measured
// in characters.
type WindowSize struct {
	Columns int
	Rows    int
}

type ProcessIO struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

// Process represents a running process in a Container. Containers may have
// multiple Processes running simultaneously, but typically just have one (task
// script, resource script, etc.). Hijacking is run as a separate Process.
type Process interface {
	// ID gives an identifier for the process that can be used in Container.Attach().
	ID() string
	// Wait waits for the process to exit and returns the result. A failed
	// script does not return an error - only a non-zero ExitStatus in the
	// ProcessResult.
	Wait(context.Context) (ProcessResult, error)
	// SetTTY updates the TTY of the Process dynamically, e.g. to resize the
	// terminal.
	SetTTY(tty TTYSpec) error
}

// ProcessResult is the result of executing a Process.
type ProcessResult struct {
	ExitStatus int
}

// Input represents a Volume (typically from a build artifact) to mount to the
// container.
type Input struct {
	// Volume is the volume to mount. It may be on any worker, so may need to
	// be streamed.
	Volume Volume
	// DestinationPath is the path in the container to mount the input.
	//
	// May be absolute or relative to ContainerSpec.Dir.
	DestinationPath string
}

// OutputPaths is a mapping from output name to its path in the container.
type OutputPaths map[string]string

// ImageSpec defines where the container image should come from. For
// containerized platforms, only one of ImageVolume, ImageURL, or ResourceType
// must be set. For non-containerized platforms, each of those fields should be
// left empty.
type ImageSpec struct {
	// ImageVolume is the Volume to use as the container image.
	ImageVolume Volume
	// ImageURL points to a path on the worker to use as the rootfs. This
	// corresponds with `task.rootfs_uri`:
	// https://concourse-ci.org/tasks.html#schema.task.rootfs_uri
	ImageURL string
	// ResourceType is the name of the base resource type to use for the image.
	ResourceType string

	// Privileged indicates whether the container should be privileged. The
	// precise meaning of "privileged" is runtime specific.
	Privileged bool
}

// ContainerLimits defines resource limits for a Container.
type ContainerLimits struct {
	// CPU defines the CPU limit for all Processes run in the Container,
	// measured in shares. Unset means no limit.
	CPU *uint64
	// Memory defines the memory limit for all Processes run in the Container,
	// measured in bytes. Unset means no limit.
	Memory *uint64
}

// Volume represents a data volume that can be mounted to a Container.
type Volume interface {
	// Handle gives the globally unique ID of the Volume.
	Handle() string

	// StreamIn accepts a compressed tar stream and populates the contents of
	// the Volume with the contents under path. This tar stream can be the
	// result of a StreamOut call for another Volume.
	//
	// path is a relative path - "." indicates using the root of the Volume.
	StreamIn(ctx context.Context, path string, compression compression.Compression, reader io.Reader) error
	// StreamOut converts the contents of the Volume under path to a compressed
	// tar stream. The result of StreamOut can be passed in to StreamIn for
	// another Volume.
	//
	// path is a relative path - "." indicates using the root of the Volume.
	StreamOut(ctx context.Context, path string, compression compression.Compression) (io.ReadCloser, error)

	// InitializeResourceCache is called upon a successful run of the get step
	// to register this Volume as a resource cache.
	InitializeResourceCache(lager.Logger, db.UsedResourceCache) error

	// InitializeTaskCache is called upon a successful run of the task step to
	// register this Volume as a task cache.
	InitializeTaskCache(logger lager.Logger, jobID int, stepName string, path string, privileged bool) error

	DBVolume() db.CreatedVolume
}

// P2PVolume is an *optional* interface that may also be satisfied by Volume
// implementations. When streaming contents from one Volume to another, if both
// Volumes implement this interface, then the P2P streaming methods will be
// used.
type P2PVolume interface {
	// Handle gives the globally unique ID of the Volume.
	Handle() string

	// GetStreamInP2PURL gives a URL which, if you POST to with a compressed
	// tar stream (using the encoding format specified by the
	// `Content-Encoding` header parameter - either "gzip" or "zstd"), will
	// stream those contents into the Volume under path.
	//
	// path is a relative path - "." indicates using the root of the Volume.
	GetStreamInP2PURL(ctx context.Context, path string) (string, error)

	// StreamP2POut takes in a URL generated by another P2PVolume's
	// GetStreamInP2PURL call, and makes a POST request with the compressed tar
	// stream of the contents of the Volume under path. The encoding format is
	// specified by the `Content-Encoding` header (either "gzip" or "zstd").
	//
	// path is a relative path - "." indicates using the root of the Volume.
	StreamP2POut(ctx context.Context, path string, destURL string, compression compression.Compression) error
}

// VolumeMount defines a Volume mounted at a particular path in a Container.
type VolumeMount struct {
	// Volume is the mounted Volume.
	Volume Volume
	// MountPath is the absolute path in the Container at which the Volume is
	// mounted.
	MountPath string
}
<|endoftext|>


<commit_before>package client

import "sync"
import "fmt"
import "encoding/json"
import "math/rand"

import common "github.com/couchbase/indexing/secondary/common"
import mclient "github.com/couchbase/indexing/secondary/manager/client"

type metadataClient struct {
	clusterURL  string
	serviceAddr string
	mdClient    *mclient.MetadataProvider
	rw          sync.RWMutex      // protects all fields listed below
	adminports  []string          // list of nodes represented by its adminport.
	queryports  map[string]string // adminport -> queryport
	// sherlock topology management, multi-node & single-partition.
	topology map[string][]*mclient.IndexMetadata // adminport -> indexes
	// shelock load replicas.
	replicas map[common.IndexDefnId][]common.IndexDefnId
	// shelock load balancing.
	loads map[common.IndexDefnId]*loadHeuristics // adminport -> loadHeuristics
}

func newMetaBridgeClient(
	cluster, serviceAddr string) (c *metadataClient, err error) {

	cinfo, err := common.NewClusterInfoCache(common.ClusterUrl(cluster), "default" /*pooln*/)
	if err != nil {
		return nil, err
	}
	if err = cinfo.Fetch(); err != nil {
		return nil, err
	}
	b := &metadataClient{
		clusterURL:  cluster,
		serviceAddr: serviceAddr,
		adminports:  make([]string, 0),
		queryports:  make(map[string]string, 0),
		loads:       make(map[common.IndexDefnId]*loadHeuristics),
	}
	// initialize meta-data-provide.
	b.mdClient, err = mclient.NewMetadataProvider(b.serviceAddr)
	if err != nil {
		return nil, err
	}
	// populate indexers' adminport and queryport
	if b.adminports, err = getIndexerAdminports(cinfo); err != nil {
		return nil, err
	}
	if b.queryports, err = getIndexerQueryports(cinfo); err != nil {
		return nil, err
	}
	// watch all indexers
	for _, adminport := range b.adminports {
		b.mdClient.WatchMetadata(adminport)
	}
	b.Refresh()
	return b, nil
}

// Refresh implement BridgeAccessor{} interface.
func (b *metadataClient) Refresh() ([]*mclient.IndexMetadata, error) {
	mindexes := b.mdClient.ListIndex()
	indexes := make([]*mclient.IndexMetadata, 0, len(mindexes))
	for _, mindex := range mindexes {
		indexes = append(indexes, mindex)
	}
	b.rw.Lock()
	defer b.rw.Unlock()

	b.topology = make(map[string][]*mclient.IndexMetadata)
	for _, adminport := range b.adminports {
		b.topology[adminport] = make([]*mclient.IndexMetadata, 0)
	}
	// gather topology of each index.
	for _, index := range indexes {
		for _, instance := range index.Instances {
			for _, queryport := range instance.Endpts {
				adminport := b.queryport2adminport(string(queryport))
				b.topology[adminport] = append(b.topology[adminport], index)
			}
		}
	}
	// compute replicas
	b.replicas = b.computeReplicas()
	return indexes, nil
}

// Nodes implement BridgeAccessor{} interface.
func (b *metadataClient) Nodes() (map[string]string, error) {
	b.rw.Lock()
	defer b.rw.Unlock()

	nodes := make(map[string]string)
	for adminport, queryport := range b.queryports {
		nodes[adminport] = queryport
	}
	return nodes, nil
}

// CreateIndex implements BridgeAccessor{} interface.
func (b *metadataClient) CreateIndex(
	indexName, bucket, using, exprType, partnExpr, whereExpr string,
	secExprs []string, isPrimary bool,
	planJSON []byte) (common.IndexDefnId, error) {

	var adminport string

	if planJSON != nil && len(planJSON) > 0 {
		plan := make(map[string]interface{})
		err := json.Unmarshal(planJSON, &plan)
		if err != nil {
			return common.IndexDefnId(0), err
		}
		nodes, ok := plan["nodes"].([]interface{})
		if ok {
			if len(nodes) < 1 {
				return common.IndexDefnId(0), ErrorEmptyDeployment
			} else if len(nodes) > 1 {
				return common.IndexDefnId(0), ErrorManyDeployment
			}
			adminport, ok = nodes[0].(string) //topology
			if !ok {
				return common.IndexDefnId(0), ErrorInvalidDeploymentNode
			}
		}
	}
	if adminport == "" {
		// if plan is not provided pick a random indexer node
		n := rand.Intn(len(b.adminports))
		adminport = b.adminports[n]
	}

	defnID, err := b.mdClient.CreateIndex(
		indexName, bucket, using, exprType, partnExpr, whereExpr,
		adminport, secExprs, isPrimary)
	b.Refresh() // refresh so that we too have IndexMetadata table.
	return defnID, err
}

// BuildIndexes implements BridgeAccessor{} interface.
func (b *metadataClient) BuildIndexes(defnIDs []common.IndexDefnId) error {
	_, ok := b.getNodes(defnIDs)
	if !ok {
		return ErrorIndexNotFound
	}

	dispatch := make(map[string][]common.IndexDefnId) // adminport -> []indexes
	for _, defnID := range defnIDs {
		adminport, ok := b.getNode(defnID)
		if !ok {
			return ErrorIndexNotFound
		}
		if _, ok := dispatch[adminport]; !ok {
			dispatch[adminport] = make([]common.IndexDefnId, 0)
		}
		dispatch[adminport] = append(dispatch[adminport], defnID)
	}

	// TODO: This needs to be added at the metadata-client.
	//for adminport, defnIDs := range dispatch {
	//  b.mdClient.BuildIndexes(adminport, defnIDs)
	//}
	return nil
}

// DropIndex implements BridgeAccessor{} interface.
func (b *metadataClient) DropIndex(defnID common.IndexDefnId) error {
	adminport, ok := b.getNode(defnID)
	if !ok {
		return ErrorIndexNotFound
	}
	return b.mdClient.DropIndex(defnID, adminport)
}

// GetScanports implements BridgeAccessor{} interface.
func (b *metadataClient) GetScanports() (queryports []string) {
	b.rw.Lock()
	defer b.rw.Unlock()

	if len(b.queryports) == 0 {
		return nil
	}
	queryports = make([]string, 0, len(b.queryports))
	for _, queryport := range b.queryports {
		queryports = append(queryports, queryport)
	}
	return queryports
}

// GetScanport implements BridgeAccessor{} interface.
func (b *metadataClient) GetScanport(
	defnID common.IndexDefnId) (queryport string, ok bool) {

	defnID = b.pickOptimal(defnID) // defnID (aka index) under least load
	adminport, ok := b.getNode(defnID)
	if !ok {
		return "", false
	}

	b.rw.Lock()
	defer b.rw.Unlock()
	queryport, ok = b.queryports[adminport]
	return queryport, ok
}

// Timeit implement BridgeAccessor{} interface.
func (b *metadataClient) Timeit(defnID uint64, value float64) {
	b.rw.Lock()
	defer b.rw.Unlock()

	id := common.IndexDefnId(defnID)
	if load, ok := b.loads[id]; !ok {
		b.loads[id] = &loadHeuristics{avgLoad: value, count: 1}
	} else {
		// compute incremental average.
		avg, n := load.avgLoad, load.count
		load.avgLoad = (float64(n)*avg + float64(value)) / float64(n+1)
		load.count = n + 1
	}
}

// IndexState implement BridgeAccessor{} interface.
func (b *metadataClient) IndexState(defnID uint64) (common.IndexState, error) {
	b.Refresh()

	b.rw.RLock()
	defer b.rw.RUnlock()

	for _, indexes := range b.topology {
		for _, index := range indexes {
			if index.Definition.DefnId == common.IndexDefnId(defnID) {
				if index.Instances != nil && len(index.Instances) > 0 {
					state := index.Instances[0].State
					if state != common.INDEX_STATE_ACTIVE {
						return state, ErrorIndexNotReady
					}
					return state, nil
				}
				return common.INDEX_STATE_ERROR, ErrorInstanceNotFound
			}
		}
	}
	return common.INDEX_STATE_ERROR, ErrorIndexNotFound
}

// close this bridge, to be called when a new indexer is added or
// an active indexer leaves the cluster or during system shutdown.
func (b *metadataClient) Close() {
	b.mdClient.Close()
}

//--------------------------------
// local functions to map replicas
//--------------------------------

// compute a map of replicas for each index in 2i.
func (b *metadataClient) computeReplicas() map[common.IndexDefnId][]common.IndexDefnId {
	replicaMap := make(map[common.IndexDefnId][]common.IndexDefnId, 0)
	for adminport1, indexes1 := range b.topology {
		for _, index1 := range indexes1 {
			replicas := make([]common.IndexDefnId, 0)
			replicas = append(replicas, index1.Definition.DefnId) // add itself
			for adminport2, indexes2 := range b.topology {
				if adminport1 == adminport2 { // skip colocated indexes
					continue
				}
				for _, index2 := range indexes2 {
					if b.equivalentIndex(index1, index2) { // pick equivalents
						replicas = append(replicas, index2.Definition.DefnId)
					}
				}
			}
			replicaMap[index1.Definition.DefnId] = replicas // map it
		}
	}
	return replicaMap
}

// compare whether two index are equivalent.
func (b *metadataClient) equivalentIndex(
	index1, index2 *mclient.IndexMetadata) bool {

	d1, d2 := index1.Definition, index2.Definition
	if d1.Using != d1.Using ||
		d1.Bucket != d2.Bucket ||
		d1.IsPrimary != d2.IsPrimary ||
		d1.ExprType != d2.ExprType ||
		d1.PartitionScheme != d2.PartitionScheme ||
		d1.PartitionKey != d2.PartitionKey {

		return false
	}

	for _, s1 := range d1.SecExprs {
		for _, s2 := range d2.SecExprs {
			if s1 != s2 {
				return false
			}
		}
	}
	return true
}

//--------------------------------
// local functions to map replicas
//--------------------------------

// manage load statistics.
type loadHeuristics struct {
	avgLoad float64
	count   uint64
}

// pick an optimal replica for the index `defnID` under least load.
func (b *metadataClient) pickOptimal(
	defnID common.IndexDefnId) common.IndexDefnId {

	b.rw.Lock()
	defer b.rw.Unlock()

	optimalID, currLoad := defnID, 0.0
	if load, ok := b.loads[defnID]; ok {
		currLoad = load.avgLoad
	}
	for _, replicaID := range b.replicas[defnID] {
		load, ok := b.loads[replicaID]
		if !ok { // no load for this replica
			return replicaID
		}
		if currLoad == 0.0 || load.avgLoad < currLoad {
			// found an index under less load
			optimalID, currLoad = replicaID, load.avgLoad
		}
	}
	return optimalID
}

//----------------
// local functions
//----------------

// getNodes return the set of nodes hosting the specified set
// of indexes
func (b *metadataClient) getNodes(
	defnIDs []common.IndexDefnId) (adminport []string, ok bool) {

	m := make(map[string]bool)
	for _, defnID := range defnIDs {
		adminport, ok := b.getNode(defnID)
		if !ok {
			return nil, false
		}
		m[adminport] = true
	}

	adminports := make([]string, 0)
	for adminport := range m {
		adminports = append(adminports, adminport)
	}
	return adminports, true
}

// getNode hosting index with `defnID`.
func (b *metadataClient) getNode(
	defnID common.IndexDefnId) (adminport string, ok bool) {

	b.rw.RLock()
	defer b.rw.RUnlock()

	for addr, indexes := range b.topology {
		for _, index := range indexes {
			if defnID == index.Definition.DefnId {
				return addr, true
			}
		}
	}
	return "", false
}

// given queryport fetch the corresponding adminport for the indexer node.
func (b *metadataClient) queryport2adminport(queryport string) string {
	for adminport, qport := range b.queryports {
		if qport == queryport {
			return adminport
		}
	}
	panic(fmt.Errorf("cannot find adminport for %v", queryport))
}

// return adminports for all known indexers.
func getIndexerAdminports(
	cinfo *common.ClusterInfoCache) ([]string, error) {

	iAdminports := make([]string, 0)
	for _, node := range cinfo.GetNodesByServiceType("indexAdmin") {
		adminport, err := cinfo.GetServiceAddress(node, "indexAdmin")
		if err != nil {
			return nil, err
		}
		iAdminports = append(iAdminports, adminport)
	}
	return iAdminports, nil
}

// return queryports for all known indexers.
func getIndexerQueryports(
	cinfo *common.ClusterInfoCache) (map[string]string, error) {

	iQueryports := make(map[string]string)
	for _, node := range cinfo.GetNodesByServiceType("indexAdmin") {
		adminport, err := cinfo.GetServiceAddress(node, "indexAdmin")
		if err != nil {
			return nil, err
		}
		queryport, err := cinfo.GetServiceAddress(node, "indexScan")
		if err != nil {
			return nil, err
		}
		iQueryports[adminport] = queryport
	}
	return iQueryports, nil
}
<commit_msg>IndexState() state to be interpreted by caller.<commit_after>package client

import "sync"
import "fmt"
import "encoding/json"
import "math/rand"

import common "github.com/couchbase/indexing/secondary/common"
import mclient "github.com/couchbase/indexing/secondary/manager/client"

type metadataClient struct {
	clusterURL  string
	serviceAddr string
	mdClient    *mclient.MetadataProvider
	rw          sync.RWMutex      // protects all fields listed below
	adminports  []string          // list of nodes represented by its adminport.
	queryports  map[string]string // adminport -> queryport
	// sherlock topology management, multi-node & single-partition.
	topology map[string][]*mclient.IndexMetadata // adminport -> indexes
	// shelock load replicas.
	replicas map[common.IndexDefnId][]common.IndexDefnId
	// shelock load balancing.
	loads map[common.IndexDefnId]*loadHeuristics // adminport -> loadHeuristics
}

func newMetaBridgeClient(
	cluster, serviceAddr string) (c *metadataClient, err error) {

	cinfo, err := common.NewClusterInfoCache(common.ClusterUrl(cluster), "default" /*pooln*/)
	if err != nil {
		return nil, err
	}
	if err = cinfo.Fetch(); err != nil {
		return nil, err
	}
	b := &metadataClient{
		clusterURL:  cluster,
		serviceAddr: serviceAddr,
		adminports:  make([]string, 0),
		queryports:  make(map[string]string, 0),
		loads:       make(map[common.IndexDefnId]*loadHeuristics),
	}
	// initialize meta-data-provide.
	b.mdClient, err = mclient.NewMetadataProvider(b.serviceAddr)
	if err != nil {
		return nil, err
	}
	// populate indexers' adminport and queryport
	if b.adminports, err = getIndexerAdminports(cinfo); err != nil {
		return nil, err
	}
	if b.queryports, err = getIndexerQueryports(cinfo); err != nil {
		return nil, err
	}
	// watch all indexers
	for _, adminport := range b.adminports {
		b.mdClient.WatchMetadata(adminport)
	}
	b.Refresh()
	return b, nil
}

// Refresh implement BridgeAccessor{} interface.
func (b *metadataClient) Refresh() ([]*mclient.IndexMetadata, error) {
	mindexes := b.mdClient.ListIndex()
	indexes := make([]*mclient.IndexMetadata, 0, len(mindexes))
	for _, mindex := range mindexes {
		indexes = append(indexes, mindex)
	}
	b.rw.Lock()
	defer b.rw.Unlock()

	b.topology = make(map[string][]*mclient.IndexMetadata)
	for _, adminport := range b.adminports {
		b.topology[adminport] = make([]*mclient.IndexMetadata, 0)
	}
	// gather topology of each index.
	for _, index := range indexes {
		for _, instance := range index.Instances {
			for _, queryport := range instance.Endpts {
				adminport := b.queryport2adminport(string(queryport))
				b.topology[adminport] = append(b.topology[adminport], index)
			}
		}
	}
	// compute replicas
	b.replicas = b.computeReplicas()
	return indexes, nil
}

// Nodes implement BridgeAccessor{} interface.
func (b *metadataClient) Nodes() (map[string]string, error) {
	b.rw.Lock()
	defer b.rw.Unlock()

	nodes := make(map[string]string)
	for adminport, queryport := range b.queryports {
		nodes[adminport] = queryport
	}
	return nodes, nil
}

// CreateIndex implements BridgeAccessor{} interface.
func (b *metadataClient) CreateIndex(
	indexName, bucket, using, exprType, partnExpr, whereExpr string,
	secExprs []string, isPrimary bool,
	planJSON []byte) (common.IndexDefnId, error) {

	var adminport string

	if planJSON != nil && len(planJSON) > 0 {
		plan := make(map[string]interface{})
		err := json.Unmarshal(planJSON, &plan)
		if err != nil {
			return common.IndexDefnId(0), err
		}
		nodes, ok := plan["nodes"].([]interface{})
		if ok {
			if len(nodes) < 1 {
				return common.IndexDefnId(0), ErrorEmptyDeployment
			} else if len(nodes) > 1 {
				return common.IndexDefnId(0), ErrorManyDeployment
			}
			adminport, ok = nodes[0].(string) //topology
			if !ok {
				return common.IndexDefnId(0), ErrorInvalidDeploymentNode
			}
		}
	}
	if adminport == "" {
		// if plan is not provided pick a random indexer node
		n := rand.Intn(len(b.adminports))
		adminport = b.adminports[n]
	}

	defnID, err := b.mdClient.CreateIndex(
		indexName, bucket, using, exprType, partnExpr, whereExpr,
		adminport, secExprs, isPrimary)
	b.Refresh() // refresh so that we too have IndexMetadata table.
	return defnID, err
}

// BuildIndexes implements BridgeAccessor{} interface.
func (b *metadataClient) BuildIndexes(defnIDs []common.IndexDefnId) error {
	_, ok := b.getNodes(defnIDs)
	if !ok {
		return ErrorIndexNotFound
	}

	dispatch := make(map[string][]common.IndexDefnId) // adminport -> []indexes
	for _, defnID := range defnIDs {
		adminport, ok := b.getNode(defnID)
		if !ok {
			return ErrorIndexNotFound
		}
		if _, ok := dispatch[adminport]; !ok {
			dispatch[adminport] = make([]common.IndexDefnId, 0)
		}
		dispatch[adminport] = append(dispatch[adminport], defnID)
	}

	// TODO: This needs to be added at the metadata-client.
	//for adminport, defnIDs := range dispatch {
	//  b.mdClient.BuildIndexes(adminport, defnIDs)
	//}
	return nil
}

// DropIndex implements BridgeAccessor{} interface.
func (b *metadataClient) DropIndex(defnID common.IndexDefnId) error {
	adminport, ok := b.getNode(defnID)
	if !ok {
		return ErrorIndexNotFound
	}
	return b.mdClient.DropIndex(defnID, adminport)
}

// GetScanports implements BridgeAccessor{} interface.
func (b *metadataClient) GetScanports() (queryports []string) {
	b.rw.Lock()
	defer b.rw.Unlock()

	if len(b.queryports) == 0 {
		return nil
	}
	queryports = make([]string, 0, len(b.queryports))
	for _, queryport := range b.queryports {
		queryports = append(queryports, queryport)
	}
	return queryports
}

// GetScanport implements BridgeAccessor{} interface.
func (b *metadataClient) GetScanport(
	defnID common.IndexDefnId) (queryport string, ok bool) {

	defnID = b.pickOptimal(defnID) // defnID (aka index) under least load
	adminport, ok := b.getNode(defnID)
	if !ok {
		return "", false
	}

	b.rw.Lock()
	defer b.rw.Unlock()
	queryport, ok = b.queryports[adminport]
	return queryport, ok
}

// Timeit implement BridgeAccessor{} interface.
func (b *metadataClient) Timeit(defnID uint64, value float64) {
	b.rw.Lock()
	defer b.rw.Unlock()

	id := common.IndexDefnId(defnID)
	if load, ok := b.loads[id]; !ok {
		b.loads[id] = &loadHeuristics{avgLoad: value, count: 1}
	} else {
		// compute incremental average.
		avg, n := load.avgLoad, load.count
		load.avgLoad = (float64(n)*avg + float64(value)) / float64(n+1)
		load.count = n + 1
	}
}

// IndexState implement BridgeAccessor{} interface.
func (b *metadataClient) IndexState(defnID uint64) (common.IndexState, error) {
	b.Refresh()

	b.rw.RLock()
	defer b.rw.RUnlock()

	for _, indexes := range b.topology {
		for _, index := range indexes {
			if index.Definition.DefnId == common.IndexDefnId(defnID) {
				if index.Instances != nil && len(index.Instances) > 0 {
					return index.Instances[0].State, nil
				}
				return common.INDEX_STATE_ERROR, ErrorInstanceNotFound
			}
		}
	}
	return common.INDEX_STATE_ERROR, ErrorIndexNotFound
}

// close this bridge, to be called when a new indexer is added or
// an active indexer leaves the cluster or during system shutdown.
func (b *metadataClient) Close() {
	b.mdClient.Close()
}

//--------------------------------
// local functions to map replicas
//--------------------------------

// compute a map of replicas for each index in 2i.
func (b *metadataClient) computeReplicas() map[common.IndexDefnId][]common.IndexDefnId {
	replicaMap := make(map[common.IndexDefnId][]common.IndexDefnId, 0)
	for adminport1, indexes1 := range b.topology {
		for _, index1 := range indexes1 {
			replicas := make([]common.IndexDefnId, 0)
			replicas = append(replicas, index1.Definition.DefnId) // add itself
			for adminport2, indexes2 := range b.topology {
				if adminport1 == adminport2 { // skip colocated indexes
					continue
				}
				for _, index2 := range indexes2 {
					if b.equivalentIndex(index1, index2) { // pick equivalents
						replicas = append(replicas, index2.Definition.DefnId)
					}
				}
			}
			replicaMap[index1.Definition.DefnId] = replicas // map it
		}
	}
	return replicaMap
}

// compare whether two index are equivalent.
func (b *metadataClient) equivalentIndex(
	index1, index2 *mclient.IndexMetadata) bool {

	d1, d2 := index1.Definition, index2.Definition
	if d1.Using != d1.Using ||
		d1.Bucket != d2.Bucket ||
		d1.IsPrimary != d2.IsPrimary ||
		d1.ExprType != d2.ExprType ||
		d1.PartitionScheme != d2.PartitionScheme ||
		d1.PartitionKey != d2.PartitionKey {

		return false
	}

	for _, s1 := range d1.SecExprs {
		for _, s2 := range d2.SecExprs {
			if s1 != s2 {
				return false
			}
		}
	}
	return true
}

//--------------------------------
// local functions to map replicas
//--------------------------------

// manage load statistics.
type loadHeuristics struct {
	avgLoad float64
	count   uint64
}

// pick an optimal replica for the index `defnID` under least load.
func (b *metadataClient) pickOptimal(
	defnID common.IndexDefnId) common.IndexDefnId {

	b.rw.Lock()
	defer b.rw.Unlock()

	optimalID, currLoad := defnID, 0.0
	if load, ok := b.loads[defnID]; ok {
		currLoad = load.avgLoad
	}
	for _, replicaID := range b.replicas[defnID] {
		load, ok := b.loads[replicaID]
		if !ok { // no load for this replica
			return replicaID
		}
		if currLoad == 0.0 || load.avgLoad < currLoad {
			// found an index under less load
			optimalID, currLoad = replicaID, load.avgLoad
		}
	}
	return optimalID
}

//----------------
// local functions
//----------------

// getNodes return the set of nodes hosting the specified set
// of indexes
func (b *metadataClient) getNodes(
	defnIDs []common.IndexDefnId) (adminport []string, ok bool) {

	m := make(map[string]bool)
	for _, defnID := range defnIDs {
		adminport, ok := b.getNode(defnID)
		if !ok {
			return nil, false
		}
		m[adminport] = true
	}

	adminports := make([]string, 0)
	for adminport := range m {
		adminports = append(adminports, adminport)
	}
	return adminports, true
}

// getNode hosting index with `defnID`.
func (b *metadataClient) getNode(
	defnID common.IndexDefnId) (adminport string, ok bool) {

	b.rw.RLock()
	defer b.rw.RUnlock()

	for addr, indexes := range b.topology {
		for _, index := range indexes {
			if defnID == index.Definition.DefnId {
				return addr, true
			}
		}
	}
	return "", false
}

// given queryport fetch the corresponding adminport for the indexer node.
func (b *metadataClient) queryport2adminport(queryport string) string {
	for adminport, qport := range b.queryports {
		if qport == queryport {
			return adminport
		}
	}
	panic(fmt.Errorf("cannot find adminport for %v", queryport))
}

// return adminports for all known indexers.
func getIndexerAdminports(
	cinfo *common.ClusterInfoCache) ([]string, error) {

	iAdminports := make([]string, 0)
	for _, node := range cinfo.GetNodesByServiceType("indexAdmin") {
		adminport, err := cinfo.GetServiceAddress(node, "indexAdmin")
		if err != nil {
			return nil, err
		}
		iAdminports = append(iAdminports, adminport)
	}
	return iAdminports, nil
}

// return queryports for all known indexers.
func getIndexerQueryports(
	cinfo *common.ClusterInfoCache) (map[string]string, error) {

	iQueryports := make(map[string]string)
	for _, node := range cinfo.GetNodesByServiceType("indexAdmin") {
		adminport, err := cinfo.GetServiceAddress(node, "indexAdmin")
		if err != nil {
			return nil, err
		}
		queryport, err := cinfo.GetServiceAddress(node, "indexScan")
		if err != nil {
			return nil, err
		}
		iQueryports[adminport] = queryport
	}
	return iQueryports, nil
}
<|endoftext|>