<commit_before>/*
RpcClient for Go RPC Servers
Copyright (C) ITsysCOM GmbH

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

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

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

package rpcclient

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	//"log/syslog"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"
	"reflect"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	JSON_RPC       = "json"
	JSON_HTTP      = "http_jsonrpc"
	GOB_RPC        = "gob"
	INTERNAL_RPC   = "*internal"
	POOL_FIRST     = "first"
	POOL_RANDOM    = "random"
	POOL_NEXT      = "next"
	POOL_BROADCAST = "broadcast"
)

var (
	ErrReqUnsynchronized       = errors.New("REQ_UNSYNCHRONIZED")
	ErrUnsupporteServiceMethod = errors.New("UNSUPPORTED_SERVICE_METHOD")
	ErrWrongArgsType           = errors.New("WRONG_ARGS_TYPE")
	ErrWrongReplyType          = errors.New("WRONG_REPLY_TYPE")
	ErrDisconnected            = errors.New("DISCONNECTED")
	ErrReplyTimeout            = errors.New("REPLY_TIMEOUT")
	//logger                     *syslog.Writer
)

func init() {
	//logger, _ = syslog.New(syslog.LOG_INFO, "RPCClient") // If we need to report anything to syslog
}

// successive Fibonacci numbers.
func Fib() func() time.Duration {
	a, b := 0, 1
	return func() time.Duration {
		a, b = b, a+b
		return time.Duration(a*10) * time.Millisecond
	}
}

func NewRpcClient(transport, addr string, connectAttempts, reconnects int, connTimeout, replyTimeout time.Duration, codec string, internalConn RpcClientConnection) (*RpcClient, error) {
	var err error
	rpcClient := &RpcClient{transport: transport, address: addr, reconnects: reconnects,
		connTimeout: connTimeout, replyTimeout: replyTimeout, codec: codec, connection: internalConn, connMux: new(sync.Mutex)}
	delay := Fib()
	for i := 0; i < connectAttempts; i++ {
		err = rpcClient.connect()
		if err == nil { //Connected so no need to reiterate
			break
		}
		time.Sleep(delay())
	}
	return rpcClient, err
}

type RpcClient struct {
	transport    string
	address      string
	reconnects   int
	connTimeout  time.Duration
	replyTimeout time.Duration
	codec        string // JSON_RPC or GOB_RPC
	connection   RpcClientConnection
	connMux      *sync.Mutex
}

func (self *RpcClient) connect() (err error) {
	self.connMux.Lock()
	defer self.connMux.Unlock()
	if self.codec == INTERNAL_RPC {
		return nil
	} else if self.codec == JSON_HTTP {
		self.connection = &HttpJsonRpcClient{httpClient: new(http.Client), url: self.address}
		return
	}
	// RPC compliant connections here, manually create connection to timeout
	netconn, err := net.DialTimeout(self.transport, self.address, self.connTimeout)
	if err != nil {
		self.connection = nil // So we don't wrap nil into the interface
		return err
	}
	if self.codec == JSON_RPC {
		self.connection = jsonrpc.NewClient(netconn)
	} else {
		self.connection = rpc.NewClient(netconn)
	}
	return
}

func (self *RpcClient) reconnect() (err error) {
	if self.reconnects == 0 {
		return nil // no reconnects
	}
	if self.codec == JSON_HTTP { // http client has automatic reconnects in place
		return self.connect()
	}
	i := 0
	delay := Fib()
	for {
		i++
		if self.reconnects != -1 && i > self.reconnects { // Maximum reconnects reached, -1 for infinite reconnects
			break
		}
		if err = self.connect(); err == nil { // No error on connect, succcess
			return nil
		}
		time.Sleep(delay()) // Cound not reconnect, retry
	}
	return errors.New("RECONNECT_FAIL")
}

func (self *RpcClient) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	if args == nil {
		return fmt.Errorf("nil rpc in argument method: %s in: %v out: %v", serviceMethod, args, reply)
	}
	if self.connection == nil {
		err = ErrDisconnected
	} else {
		errChan := make(chan error, 1)
		go func() {
			errChan <- self.connection.Call(serviceMethod, args, reply)
		}()
		select {
		case err = <-errChan:
		case <-time.After(self.replyTimeout):
			err = ErrReplyTimeout
		}
	}
	if isNetworkError(err) && err != ErrReplyTimeout && self.reconnects != 0 { // ReplyTimeout should not reconnect since it creates loop
		if errReconnect := self.reconnect(); errReconnect != nil {
			return err
		} else { // Run command after reconnect
			return self.Call(serviceMethod, args, reply)
		}
	}
	return err
}

// Connection used in RpcClient, as interface so we can combine the rpc.RpcClient with http one or websocket
type RpcClientConnection interface {
	Call(string, interface{}, interface{}) error
}

// Response received for
type JsonRpcResponse struct {
	Id     uint64
	Result *json.RawMessage
	Error  interface{}
}

type HttpJsonRpcClient struct {
	httpClient *http.Client
	id         uint64
	url        string
}

func (self *HttpJsonRpcClient) Call(serviceMethod string, args interface{}, reply interface{}) error {
	self.id += 1
	id := self.id
	data, err := json.Marshal(map[string]interface{}{
		"method": serviceMethod,
		"id":     self.id,
		"params": [1]interface{}{args},
	})
	if err != nil {
		return err
	}
	resp, err := self.httpClient.Post(self.url, "application/json", ioutil.NopCloser(strings.NewReader(string(data)))) // Closer so we automatically have close after response
	if err != nil {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var jsonRsp JsonRpcResponse
	err = json.Unmarshal(body, &jsonRsp)
	if err != nil {
		return err
	}
	if jsonRsp.Id != id {
		return ErrReqUnsynchronized
	}
	if jsonRsp.Error != nil || jsonRsp.Result == nil {
		x, ok := jsonRsp.Error.(string)
		if !ok {
			return fmt.Errorf("invalid error %v", jsonRsp.Error)
		}
		if x == "" {
			x = "unspecified error"
		}
		return errors.New(x)
	}
	return json.Unmarshal(*jsonRsp.Result, reply)
}

type RpcClientPool struct {
	transmissionType string
	connections      []RpcClientConnection
	counter          int
	replyTimeout     time.Duration
}

func NewRpcClientPool(transmissionType string, replyTimeout time.Duration) *RpcClientPool {
	return &RpcClientPool{transmissionType: transmissionType, replyTimeout: replyTimeout}
}

func (pool *RpcClientPool) AddClient(rcc RpcClientConnection) {
	if rcc != nil {
		pool.connections = append(pool.connections, rcc)
	}
}

func (pool *RpcClientPool) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	switch pool.transmissionType {
	case POOL_BROADCAST:
		replyChan := make(chan *rpcReplyError, len(pool.connections))
		for _, rc := range pool.connections {
			go func(conn RpcClientConnection) {
				// make a new pointer of the same type
				rpl := reflect.New(reflect.TypeOf(reflect.ValueOf(reply).Elem().Interface()))
				err := conn.Call(serviceMethod, args, rpl.Interface())
				if !isNetworkError(err) {
					replyChan <- &rpcReplyError{reply: rpl.Interface(), err: err}
				}
			}(rc)
		}
		//get first response with timeout
		var re *rpcReplyError
		select {
		case re = <-replyChan:
		case <-time.After(pool.replyTimeout):
			return ErrReplyTimeout
		}
		// put received value in the orig reply
		reflect.ValueOf(reply).Elem().Set(reflect.ValueOf(re.reply).Elem())
		return re.err
	case POOL_FIRST:
		for _, rc := range pool.connections {
			err = rc.Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_NEXT:
		ln := len(pool.connections)
		rrIndexes := roundIndex(int(math.Mod(float64(pool.counter), float64(ln))), ln)
		pool.counter++
		for _, index := range rrIndexes {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_RANDOM:
		rand.Seed(time.Now().UnixNano())
		randomIndex := rand.Perm(len(pool.connections))
		for _, index := range randomIndex {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	}
	return
}

type rpcReplyError struct {
	reply interface{}
	err   error
}

// generates round robin indexes for a slice of length max
// starting from index start
func roundIndex(start, max int) []int {
	if start < 0 {
		start = 0
	}
	result := make([]int, max)
	for i := 0; i < max; i++ {
		if start+i < max {
			result[i] = start + i
		} else {
			result[i] = int(math.Abs(float64(max - (start + i))))
		}
	}
	return result
}

func isNetworkError(err error) bool {
	if err == nil {
		return false
	}
	if operr, ok := err.(*net.OpError); ok && strings.HasSuffix(operr.Err.Error(), syscall.ECONNRESET.Error()) { // connection reset
		return true
	}
	return err == rpc.ErrShutdown ||
		err == ErrReqUnsynchronized ||
		err == ErrDisconnected ||
		err == ErrReplyTimeout ||
		strings.HasPrefix(err.Error(), "rpc: can't find service")
}
<commit_msg>Avoid infinite loop on client.Call<commit_after>/*
RpcClient for Go RPC Servers
Copyright (C) ITsysCOM GmbH

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

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

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

package rpcclient

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	//"log/syslog"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"
	"reflect"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	JSON_RPC       = "json"
	JSON_HTTP      = "http_jsonrpc"
	GOB_RPC        = "gob"
	INTERNAL_RPC   = "*internal"
	POOL_FIRST     = "first"
	POOL_RANDOM    = "random"
	POOL_NEXT      = "next"
	POOL_BROADCAST = "broadcast"
)

var (
	ErrReqUnsynchronized       = errors.New("REQ_UNSYNCHRONIZED")
	ErrUnsupporteServiceMethod = errors.New("UNSUPPORTED_SERVICE_METHOD")
	ErrWrongArgsType           = errors.New("WRONG_ARGS_TYPE")
	ErrWrongReplyType          = errors.New("WRONG_REPLY_TYPE")
	ErrDisconnected            = errors.New("DISCONNECTED")
	ErrReplyTimeout            = errors.New("REPLY_TIMEOUT")
	ErrFailedReconnect         = errors.New("FAILED_RECONNECT")
	//logger                     *syslog.Writer
)

func init() {
	//logger, _ = syslog.New(syslog.LOG_INFO, "RPCClient") // If we need to report anything to syslog
}

// successive Fibonacci numbers.
func Fib() func() time.Duration {
	a, b := 0, 1
	return func() time.Duration {
		a, b = b, a+b
		return time.Duration(a*10) * time.Millisecond
	}
}

func NewRpcClient(transport, addr string, connectAttempts, reconnects int, connTimeout, replyTimeout time.Duration, codec string, internalConn RpcClientConnection) (*RpcClient, error) {
	var err error
	rpcClient := &RpcClient{transport: transport, address: addr, reconnects: reconnects,
		connTimeout: connTimeout, replyTimeout: replyTimeout, codec: codec, connection: internalConn, connMux: new(sync.Mutex)}
	delay := Fib()
	for i := 0; i < connectAttempts; i++ {
		err = rpcClient.connect()
		if err == nil { //Connected so no need to reiterate
			break
		}
		time.Sleep(delay())
	}
	return rpcClient, err
}

type RpcClient struct {
	transport    string
	address      string
	reconnects   int
	connTimeout  time.Duration
	replyTimeout time.Duration
	codec        string // JSON_RPC or GOB_RPC
	connection   RpcClientConnection
	connMux      *sync.Mutex
}

func (self *RpcClient) connect() (err error) {
	self.connMux.Lock()
	defer self.connMux.Unlock()
	if self.codec == INTERNAL_RPC {
		return nil
	} else if self.codec == JSON_HTTP {
		self.connection = &HttpJsonRpcClient{httpClient: new(http.Client), url: self.address}
		return
	}
	// RPC compliant connections here, manually create connection to timeout
	netconn, err := net.DialTimeout(self.transport, self.address, self.connTimeout)
	if err != nil {
		self.connection = nil // So we don't wrap nil into the interface
		return err
	}
	if self.codec == JSON_RPC {
		self.connection = jsonrpc.NewClient(netconn)
	} else {
		self.connection = rpc.NewClient(netconn)
	}
	return
}

func (self *RpcClient) reconnect() (err error) {
	if self.reconnects == 0 {
		return nil // no reconnects
	}
	if self.codec == JSON_HTTP { // http client has automatic reconnects in place
		return self.connect()
	}
	i := 0
	delay := Fib()
	for {
		i++
		if self.reconnects != -1 && i > self.reconnects { // Maximum reconnects reached, -1 for infinite reconnects
			break
		}
		if err = self.connect(); err == nil { // No error on connect, succcess
			return nil
		}
		time.Sleep(delay()) // Cound not reconnect, retry
	}
	return ErrFailedReconnect
}

func (self *RpcClient) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	if args == nil {
		return fmt.Errorf("nil rpc in argument method: %s in: %v out: %v", serviceMethod, args, reply)
	}
	if self.connection == nil {
		err = ErrDisconnected
	} else {
		errChan := make(chan error, 1)
		go func() {
			errChan <- self.connection.Call(serviceMethod, args, reply)
		}()
		select {
		case err = <-errChan:
		case <-time.After(self.replyTimeout):
			err = ErrReplyTimeout
		}
	}
	if isNetworkError(err) && err != ErrReplyTimeout && self.reconnects != 0 { // ReplyTimeout should not reconnect since it creates loop
		if errReconnect := self.reconnect(); errReconnect != nil {
			return err
		} else { // Run command after reconnect
			return self.connection.Call(serviceMethod, args, reply)
		}
	}
	return err
}

// Connection used in RpcClient, as interface so we can combine the rpc.RpcClient with http one or websocket
type RpcClientConnection interface {
	Call(string, interface{}, interface{}) error
}

// Response received for
type JsonRpcResponse struct {
	Id     uint64
	Result *json.RawMessage
	Error  interface{}
}

type HttpJsonRpcClient struct {
	httpClient *http.Client
	id         uint64
	url        string
}

func (self *HttpJsonRpcClient) Call(serviceMethod string, args interface{}, reply interface{}) error {
	self.id += 1
	id := self.id
	data, err := json.Marshal(map[string]interface{}{
		"method": serviceMethod,
		"id":     self.id,
		"params": [1]interface{}{args},
	})
	if err != nil {
		return err
	}
	resp, err := self.httpClient.Post(self.url, "application/json", ioutil.NopCloser(strings.NewReader(string(data)))) // Closer so we automatically have close after response
	if err != nil {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var jsonRsp JsonRpcResponse
	err = json.Unmarshal(body, &jsonRsp)
	if err != nil {
		return err
	}
	if jsonRsp.Id != id {
		return ErrReqUnsynchronized
	}
	if jsonRsp.Error != nil || jsonRsp.Result == nil {
		x, ok := jsonRsp.Error.(string)
		if !ok {
			return fmt.Errorf("invalid error %v", jsonRsp.Error)
		}
		if x == "" {
			x = "unspecified error"
		}
		return errors.New(x)
	}
	return json.Unmarshal(*jsonRsp.Result, reply)
}

type RpcClientPool struct {
	transmissionType string
	connections      []RpcClientConnection
	counter          int
	replyTimeout     time.Duration
}

func NewRpcClientPool(transmissionType string, replyTimeout time.Duration) *RpcClientPool {
	return &RpcClientPool{transmissionType: transmissionType, replyTimeout: replyTimeout}
}

func (pool *RpcClientPool) AddClient(rcc RpcClientConnection) {
	if rcc != nil {
		pool.connections = append(pool.connections, rcc)
	}
}

func (pool *RpcClientPool) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	switch pool.transmissionType {
	case POOL_BROADCAST:
		replyChan := make(chan *rpcReplyError, len(pool.connections))
		for _, rc := range pool.connections {
			go func(conn RpcClientConnection) {
				// make a new pointer of the same type
				rpl := reflect.New(reflect.TypeOf(reflect.ValueOf(reply).Elem().Interface()))
				err := conn.Call(serviceMethod, args, rpl.Interface())
				if !isNetworkError(err) {
					replyChan <- &rpcReplyError{reply: rpl.Interface(), err: err}
				}
			}(rc)
		}
		//get first response with timeout
		var re *rpcReplyError
		select {
		case re = <-replyChan:
		case <-time.After(pool.replyTimeout):
			return ErrReplyTimeout
		}
		// put received value in the orig reply
		reflect.ValueOf(reply).Elem().Set(reflect.ValueOf(re.reply).Elem())
		return re.err
	case POOL_FIRST:
		for _, rc := range pool.connections {
			err = rc.Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_NEXT:
		ln := len(pool.connections)
		rrIndexes := roundIndex(int(math.Mod(float64(pool.counter), float64(ln))), ln)
		pool.counter++
		for _, index := range rrIndexes {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_RANDOM:
		rand.Seed(time.Now().UnixNano())
		randomIndex := rand.Perm(len(pool.connections))
		for _, index := range randomIndex {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	}
	return
}

type rpcReplyError struct {
	reply interface{}
	err   error
}

// generates round robin indexes for a slice of length max
// starting from index start
func roundIndex(start, max int) []int {
	if start < 0 {
		start = 0
	}
	result := make([]int, max)
	for i := 0; i < max; i++ {
		if start+i < max {
			result[i] = start + i
		} else {
			result[i] = int(math.Abs(float64(max - (start + i))))
		}
	}
	return result
}

func isNetworkError(err error) bool {
	if err == nil {
		return false
	}
	if operr, ok := err.(*net.OpError); ok && strings.HasSuffix(operr.Err.Error(), syscall.ECONNRESET.Error()) { // connection reset
		return true
	}
	return err == rpc.ErrShutdown ||
		err == ErrReqUnsynchronized ||
		err == ErrDisconnected ||
		err == ErrReplyTimeout ||
		strings.HasPrefix(err.Error(), "rpc: can't find service")
}
<|endoftext|>


<commit_before>package artnet

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/jsimonetti/go-artnet"
	"github.com/sirupsen/logrus"

	"github.com/StageAutoControl/controller/pkg/cntl"
	"github.com/StageAutoControl/controller/pkg/internal/logging"
)

// Controller is a transport for the ArtNet protocol (DMX over UDP/IP)
type controller struct {
	logger      logging.Logger
	sender      *artnet.Controller
	state       *State
	sendTrigger chan UniverseStateMap
	context     context.Context
}

// NewController returns a artnet Controller as an anonymous interface
func NewController(logger logging.Logger) (Controller, error) {
	ip, err := FindArtNetIP()
	if err != nil {
		return nil, fmt.Errorf("failed to find the art-net IP: %v", err)
	}

	if len(ip) == 0 {
		return nil, errors.New("failed to find the art-net IP: No interface found")
	}

	host, err := os.Hostname()
	if err != nil {
		return nil, fmt.Errorf("failed to resolve hostname: %v", err)
	}

	host = strings.ToLower(strings.Split(host, ".")[0])
	logger.Infof("Using ArtNet IP %s and hostname %s", ip.String(), host)

	senderLogger := artnet.NewLogger(logger.(*logrus.Entry).WithField("module", "artnet"))
	control := &controller{
		logger:      logger,
		sender:      artnet.NewController(host, ip, senderLogger),
		state:       NewState(),
		sendTrigger: make(chan UniverseStateMap, 100),
	}

	return control, nil
}

// Start the controller
func (c *controller) Start(ctx context.Context) error {
	if err := c.sender.Start(); err != nil {
		return fmt.Errorf("failed to start Controller: %v", err)
	}

	c.context = ctx
	go c.sendBackground()
	go c.debugDevices()

	return nil
}

// Stop the controller
func (c *controller) Stop() {
	close(c.sendTrigger)
	c.sender.Stop()
}

func (c *controller) SetDMXChannelValue(value ChannelValue) {
	c.state.SetChannel(value.Universe, value.Channel, value.Value)
	c.triggerSend()
}

func (c *controller) SetDMXChannelValues(values []ChannelValue) {
	c.state.SetChannelValues(values)
	c.triggerSend()
}

// Write implements the playback.TransportWriter interface to compatibility :)
func (c *controller) Write(cmd cntl.Command) error {
	values := make([]ChannelValue, len(cmd.DMXCommands))
	for i, dmxCmd := range cmd.DMXCommands {
		values[i].Universe = uint16(dmxCmd.Universe)
		values[i].Channel = uint16(dmxCmd.Channel)
		values[i].Value = dmxCmd.Value.Uint8()
	}

	c.SetDMXChannelValues(values)

	return nil
}

func (c *controller) triggerSend() {
	c.sendTrigger <- c.state.Get()
}

func (c *controller) sendBackground() {
	for data := range c.sendTrigger {
		for u, dmx := range data {
			c.sender.SendDMXToAddress(dmx.toByteSlice(), c.universeToAddress(u))
		}
	}
}

// universeToAddress converts a dmx universe to a artnet address
// Code stolen from https://play.golang.org/p/pdQPC5u7JX
func (c *controller) universeToAddress(universe uint16) artnet.Address {
	v := make([]uint8, 2)
	binary.BigEndian.PutUint16(v, universe)

	return artnet.Address{
		Net:    v[0],
		SubUni: v[1],
	}
}

func (c *controller) debugDevices() {
	t := time.NewTicker(30 * time.Second)
	for range t.C {
		c.logger.Debugf("Currently %d devices are registered: %+s", len(c.sender.Nodes), ips(c.sender.Nodes))
	}
}

func ips(nodes []*artnet.ControlledNode) (ips []string) {
	for _, n := range nodes {
		ips = append(ips, NodeToString(n))
	}
	return
}
<commit_msg>Set maxFPS<commit_after>package artnet

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/jsimonetti/go-artnet"
	"github.com/sirupsen/logrus"

	"github.com/StageAutoControl/controller/pkg/cntl"
	"github.com/StageAutoControl/controller/pkg/internal/logging"
)

// Controller is a transport for the ArtNet protocol (DMX over UDP/IP)
type controller struct {
	logger      logging.Logger
	sender      *artnet.Controller
	state       *State
	sendTrigger chan UniverseStateMap
	context     context.Context
}

// NewController returns a artnet Controller as an anonymous interface
func NewController(logger logging.Logger) (Controller, error) {
	ip, err := FindArtNetIP()
	if err != nil {
		return nil, fmt.Errorf("failed to find the art-net IP: %v", err)
	}

	if len(ip) == 0 {
		return nil, errors.New("failed to find the art-net IP: No interface found")
	}

	host, err := os.Hostname()
	if err != nil {
		return nil, fmt.Errorf("failed to resolve hostname: %v", err)
	}

	host = strings.ToLower(strings.Split(host, ".")[0])
	logger.Infof("Using ArtNet IP %s and hostname %s", ip.String(), host)

	senderLogger := artnet.NewLogger(logger.(*logrus.Entry).WithField("module", "artnet"))
	control := &controller{
		logger:      logger,
		sender:      artnet.NewController(host, ip, senderLogger, artnet.MaxFPS(40)),
		state:       NewState(),
		sendTrigger: make(chan UniverseStateMap, 100),
	}

	return control, nil
}

// Start the controller
func (c *controller) Start(ctx context.Context) error {
	if err := c.sender.Start(); err != nil {
		return fmt.Errorf("failed to start Controller: %v", err)
	}

	c.context = ctx
	go c.sendBackground()
	go c.debugDevices()

	return nil
}

// Stop the controller
func (c *controller) Stop() {
	close(c.sendTrigger)
	c.sender.Stop()
}

func (c *controller) SetDMXChannelValue(value ChannelValue) {
	c.state.SetChannel(value.Universe, value.Channel, value.Value)
	c.triggerSend()
}

func (c *controller) SetDMXChannelValues(values []ChannelValue) {
	c.state.SetChannelValues(values)
	c.triggerSend()
}

// Write implements the playback.TransportWriter interface to compatibility :)
func (c *controller) Write(cmd cntl.Command) error {
	values := make([]ChannelValue, len(cmd.DMXCommands))
	for i, dmxCmd := range cmd.DMXCommands {
		values[i].Universe = uint16(dmxCmd.Universe)
		values[i].Channel = uint16(dmxCmd.Channel)
		values[i].Value = dmxCmd.Value.Uint8()
	}

	c.SetDMXChannelValues(values)

	return nil
}

func (c *controller) triggerSend() {
	c.sendTrigger <- c.state.Get()
}

func (c *controller) sendBackground() {
	for data := range c.sendTrigger {
		for u, dmx := range data {
			c.sender.SendDMXToAddress(dmx.toByteSlice(), c.universeToAddress(u))
		}
	}
}

// universeToAddress converts a dmx universe to a artnet address
// Code stolen from https://play.golang.org/p/pdQPC5u7JX
func (c *controller) universeToAddress(universe uint16) artnet.Address {
	v := make([]uint8, 2)
	binary.BigEndian.PutUint16(v, universe)

	return artnet.Address{
		Net:    v[0],
		SubUni: v[1],
	}
}

func (c *controller) debugDevices() {
	t := time.NewTicker(30 * time.Second)
	for range t.C {
		c.logger.Debugf("Currently %d devices are registered: %+s", len(c.sender.Nodes), ips(c.sender.Nodes))
	}
}

func ips(nodes []*artnet.ControlledNode) (ips []string) {
	for _, n := range nodes {
		ips = append(ips, NodeToString(n))
	}
	return
}
<|endoftext|>


<commit_before>package datasource

import (
	"strings"
	"time"

	"github.com/alexanderzobnin/grafana-zabbix/pkg/timeseries"

	"github.com/alexanderzobnin/grafana-zabbix/pkg/zabbix"
	"github.com/grafana/grafana-plugin-sdk-go/backend"
	"github.com/grafana/grafana-plugin-sdk-go/data"
	"golang.org/x/net/context"
)

// ZabbixAPIQuery handles query requests to Zabbix API
func (ds *ZabbixDatasourceInstance) ZabbixAPIQuery(ctx context.Context, apiReq *zabbix.ZabbixAPIRequest) (*ZabbixAPIResourceResponse, error) {
	resultJson, err := ds.zabbix.Request(ctx, apiReq)
	if err != nil {
		return nil, err
	}
	result := resultJson.Interface()
	return BuildAPIResponse(&result)
}

func BuildAPIResponse(responseData *interface{}) (*ZabbixAPIResourceResponse, error) {
	return &ZabbixAPIResourceResponse{
		Result: *responseData,
	}, nil
}

// TestConnection checks authentication and version of the Zabbix API and returns that info
func (ds *ZabbixDatasourceInstance) TestConnection(ctx context.Context) (string, error) {
	_, err := ds.zabbix.GetAllGroups(ctx)
	if err != nil {
		return "", err
	}

	response, err := ds.zabbix.Request(ctx, &zabbix.ZabbixAPIRequest{Method: "apiinfo.version"})
	if err != nil {
		return "", err
	}

	resultByte, _ := response.MarshalJSON()
	return string(resultByte), nil
}

func (ds *ZabbixDatasourceInstance) queryNumericItems(ctx context.Context, query *QueryModel) ([]*data.Frame, error) {
	groupFilter := query.Group.Filter
	hostFilter := query.Host.Filter
	appFilter := query.Application.Filter
	itemTagFilter := query.ItemTag.Filter
	itemFilter := query.Item.Filter

	items, err := ds.zabbix.GetItems(ctx, groupFilter, hostFilter, appFilter, itemTagFilter, itemFilter, "num")
	if err != nil {
		return nil, err
	}

	frames, err := ds.queryNumericDataForItems(ctx, query, items)
	if err != nil {
		return nil, err
	}

	return frames, nil
}

func (ds *ZabbixDatasourceInstance) queryItemIdData(ctx context.Context, query *QueryModel) ([]*data.Frame, error) {
	itemids := strings.Split(query.ItemIDs, ",")
	for i, id := range itemids {
		itemids[i] = strings.Trim(id, " ")
	}

	items, err := ds.zabbix.GetItemsByIDs(ctx, itemids)
	if err != nil {
		return nil, err
	}

	frames, err := ds.queryNumericDataForItems(ctx, query, items)
	if err != nil {
		return nil, err
	}

	return frames, nil
}

func (ds *ZabbixDatasourceInstance) queryNumericDataForItems(ctx context.Context, query *QueryModel, items []*zabbix.Item) ([]*data.Frame, error) {
	trendValueType := ds.getTrendValueType(query)
	consolidateBy := ds.getConsolidateBy(query)

	if consolidateBy != "" {
		trendValueType = consolidateBy
	}

	err := applyFunctionsPre(query, items)
	if err != nil {
		return nil, err
	}

	history, err := ds.getHistotyOrTrend(ctx, query, items, trendValueType)
	if err != nil {
		return nil, err
	}

	series := convertHistoryToTimeSeries(history, items)
	return ds.applyDataProcessing(ctx, query, series, false)
}

func (ds *ZabbixDatasourceInstance) applyDataProcessing(ctx context.Context, query *QueryModel, series []*timeseries.TimeSeriesData, DBPostProcessing bool) ([]*data.Frame, error) {
	consolidateBy := ds.getConsolidateBy(query)
	useTrend := ds.isUseTrend(query.TimeRange)

	// Sort trend data (in some cases Zabbix API returns it unsorted)
	if useTrend {
		sortSeriesPoints(series)
	}

	// Align time series data if possible
	disableDataAlignment := query.Options.DisableDataAlignment || ds.Settings.DisableDataAlignment || query.QueryType == MODE_ITSERVICE
	if !disableDataAlignment {
		if useTrend && !DBPostProcessing {
			for _, s := range series {
				// Trend data is already aligned (by 1 hour interval), but null values should be added
				s.TS = s.TS.FillTrendWithNulls()
			}
		} else {
			for _, s := range series {
				if s.Meta.Interval != nil {
					s.TS = s.TS.Align(*s.Meta.Interval)
				}
			}
		}

		if len(series) > 1 {
			series = timeseries.PrepareForStack(series)
		}
	}

	series, err := applyFunctions(series, query.Functions)
	if err != nil {
		return nil, err
	}

	for _, s := range series {
		if int64(s.Len()) > query.MaxDataPoints && query.Interval > 0 {
			downsampleFunc := consolidateBy
			if downsampleFunc == "" {
				downsampleFunc = "avg"
			}
			downsampled, err := applyGroupBy(s.TS, query.Interval.String(), downsampleFunc)
			if err == nil {
				s.TS = downsampled
			} else {
				ds.logger.Debug("Error downsampling series", "error", err)
			}
		}
	}

	valueMaps := make([]zabbix.ValueMap, 0)
	if query.Options.UseZabbixValueMapping {
		valueMaps, err = ds.zabbix.GetValueMappings(ctx)
		if err != nil {
			ds.logger.Error("Error getting value maps", "error", err)
			valueMaps = []zabbix.ValueMap{}
		}
	}
	frames := convertTimeSeriesToDataFrames(series, valueMaps)
	return frames, nil
}

func (ds *ZabbixDatasourceInstance) getTrendValueType(query *QueryModel) string {
	trendValue := "avg"

	for _, fn := range query.Functions {
		if fn.Def.Name == "trendValue" && len(fn.Params) > 0 {
			trendValue = fn.Params[0].(string)
		}
	}

	return trendValue
}

func (ds *ZabbixDatasourceInstance) getConsolidateBy(query *QueryModel) string {
	consolidateBy := ""

	for _, fn := range query.Functions {
		if fn.Def.Name == "consolidateBy" && len(fn.Params) > 0 {
			consolidateBy = fn.Params[0].(string)
		}
	}
	return consolidateBy
}

func (ds *ZabbixDatasourceInstance) getHistotyOrTrend(ctx context.Context, query *QueryModel, items []*zabbix.Item, trendValueType string) (zabbix.History, error) {
	timeRange := query.TimeRange
	useTrend := ds.isUseTrend(timeRange)

	if useTrend {
		result, err := ds.zabbix.GetTrend(ctx, items, timeRange)
		if err != nil {
			return nil, err
		}
		return convertTrendToHistory(result, trendValueType)
	}

	return ds.zabbix.GetHistory(ctx, items, timeRange)
}

func (ds *ZabbixDatasourceInstance) isUseTrend(timeRange backend.TimeRange) bool {
	if !ds.Settings.Trends {
		return false
	}

	trendsFrom := ds.Settings.TrendsFrom
	trendsRange := ds.Settings.TrendsRange
	fromSec := timeRange.From.Unix()
	toSec := timeRange.To.Unix()
	rangeSec := float64(toSec - fromSec)

	if (fromSec < time.Now().Add(-trendsFrom).Unix()) || (rangeSec > trendsRange.Seconds()) {
		return true
	}
	return false
}
<commit_msg>Skip data alignment if item interval less than query interval<commit_after>package datasource

import (
	"strings"
	"time"

	"github.com/alexanderzobnin/grafana-zabbix/pkg/timeseries"

	"github.com/alexanderzobnin/grafana-zabbix/pkg/zabbix"
	"github.com/grafana/grafana-plugin-sdk-go/backend"
	"github.com/grafana/grafana-plugin-sdk-go/data"
	"golang.org/x/net/context"
)

// ZabbixAPIQuery handles query requests to Zabbix API
func (ds *ZabbixDatasourceInstance) ZabbixAPIQuery(ctx context.Context, apiReq *zabbix.ZabbixAPIRequest) (*ZabbixAPIResourceResponse, error) {
	resultJson, err := ds.zabbix.Request(ctx, apiReq)
	if err != nil {
		return nil, err
	}
	result := resultJson.Interface()
	return BuildAPIResponse(&result)
}

func BuildAPIResponse(responseData *interface{}) (*ZabbixAPIResourceResponse, error) {
	return &ZabbixAPIResourceResponse{
		Result: *responseData,
	}, nil
}

// TestConnection checks authentication and version of the Zabbix API and returns that info
func (ds *ZabbixDatasourceInstance) TestConnection(ctx context.Context) (string, error) {
	_, err := ds.zabbix.GetAllGroups(ctx)
	if err != nil {
		return "", err
	}

	response, err := ds.zabbix.Request(ctx, &zabbix.ZabbixAPIRequest{Method: "apiinfo.version"})
	if err != nil {
		return "", err
	}

	resultByte, _ := response.MarshalJSON()
	return string(resultByte), nil
}

func (ds *ZabbixDatasourceInstance) queryNumericItems(ctx context.Context, query *QueryModel) ([]*data.Frame, error) {
	groupFilter := query.Group.Filter
	hostFilter := query.Host.Filter
	appFilter := query.Application.Filter
	itemTagFilter := query.ItemTag.Filter
	itemFilter := query.Item.Filter

	items, err := ds.zabbix.GetItems(ctx, groupFilter, hostFilter, appFilter, itemTagFilter, itemFilter, "num")
	if err != nil {
		return nil, err
	}

	frames, err := ds.queryNumericDataForItems(ctx, query, items)
	if err != nil {
		return nil, err
	}

	return frames, nil
}

func (ds *ZabbixDatasourceInstance) queryItemIdData(ctx context.Context, query *QueryModel) ([]*data.Frame, error) {
	itemids := strings.Split(query.ItemIDs, ",")
	for i, id := range itemids {
		itemids[i] = strings.Trim(id, " ")
	}

	items, err := ds.zabbix.GetItemsByIDs(ctx, itemids)
	if err != nil {
		return nil, err
	}

	frames, err := ds.queryNumericDataForItems(ctx, query, items)
	if err != nil {
		return nil, err
	}

	return frames, nil
}

func (ds *ZabbixDatasourceInstance) queryNumericDataForItems(ctx context.Context, query *QueryModel, items []*zabbix.Item) ([]*data.Frame, error) {
	trendValueType := ds.getTrendValueType(query)
	consolidateBy := ds.getConsolidateBy(query)

	if consolidateBy != "" {
		trendValueType = consolidateBy
	}

	err := applyFunctionsPre(query, items)
	if err != nil {
		return nil, err
	}

	history, err := ds.getHistotyOrTrend(ctx, query, items, trendValueType)
	if err != nil {
		return nil, err
	}

	series := convertHistoryToTimeSeries(history, items)
	return ds.applyDataProcessing(ctx, query, series, false)
}

func (ds *ZabbixDatasourceInstance) applyDataProcessing(ctx context.Context, query *QueryModel, series []*timeseries.TimeSeriesData, DBPostProcessing bool) ([]*data.Frame, error) {
	consolidateBy := ds.getConsolidateBy(query)
	useTrend := ds.isUseTrend(query.TimeRange)

	// Sort trend data (in some cases Zabbix API returns it unsorted)
	if useTrend {
		sortSeriesPoints(series)
	}

	// Align time series data if possible
	disableDataAlignment := query.Options.DisableDataAlignment || ds.Settings.DisableDataAlignment || query.QueryType == MODE_ITSERVICE
	if !disableDataAlignment {
		if useTrend && !DBPostProcessing {
			for _, s := range series {
				// Trend data is already aligned (by 1 hour interval), but null values should be added
				s.TS = s.TS.FillTrendWithNulls()
			}
		} else {
			for _, s := range series {
				// Skip unnecessary data alignment if item interval less than query interval
				// because data will be downsampled in this case
				if s.Meta.Interval != nil && *s.Meta.Interval > query.Interval {
					s.TS = s.TS.Align(*s.Meta.Interval)
				}
			}
		}

		if len(series) > 1 {
			series = timeseries.PrepareForStack(series)
		}
	}

	series, err := applyFunctions(series, query.Functions)
	if err != nil {
		return nil, err
	}

	for _, s := range series {
		if int64(s.Len()) > query.MaxDataPoints && query.Interval > 0 {
			downsampleFunc := consolidateBy
			if downsampleFunc == "" {
				downsampleFunc = "avg"
			}
			downsampled, err := applyGroupBy(s.TS, query.Interval.String(), downsampleFunc)
			if err == nil {
				s.TS = downsampled
			} else {
				ds.logger.Debug("Error downsampling series", "error", err)
			}
		}
	}

	valueMaps := make([]zabbix.ValueMap, 0)
	if query.Options.UseZabbixValueMapping {
		valueMaps, err = ds.zabbix.GetValueMappings(ctx)
		if err != nil {
			ds.logger.Error("Error getting value maps", "error", err)
			valueMaps = []zabbix.ValueMap{}
		}
	}
	frames := convertTimeSeriesToDataFrames(series, valueMaps)
	return frames, nil
}

func (ds *ZabbixDatasourceInstance) getTrendValueType(query *QueryModel) string {
	trendValue := "avg"

	for _, fn := range query.Functions {
		if fn.Def.Name == "trendValue" && len(fn.Params) > 0 {
			trendValue = fn.Params[0].(string)
		}
	}

	return trendValue
}

func (ds *ZabbixDatasourceInstance) getConsolidateBy(query *QueryModel) string {
	consolidateBy := ""

	for _, fn := range query.Functions {
		if fn.Def.Name == "consolidateBy" && len(fn.Params) > 0 {
			consolidateBy = fn.Params[0].(string)
		}
	}
	return consolidateBy
}

func (ds *ZabbixDatasourceInstance) getHistotyOrTrend(ctx context.Context, query *QueryModel, items []*zabbix.Item, trendValueType string) (zabbix.History, error) {
	timeRange := query.TimeRange
	useTrend := ds.isUseTrend(timeRange)

	if useTrend {
		result, err := ds.zabbix.GetTrend(ctx, items, timeRange)
		if err != nil {
			return nil, err
		}
		return convertTrendToHistory(result, trendValueType)
	}

	return ds.zabbix.GetHistory(ctx, items, timeRange)
}

func (ds *ZabbixDatasourceInstance) isUseTrend(timeRange backend.TimeRange) bool {
	if !ds.Settings.Trends {
		return false
	}

	trendsFrom := ds.Settings.TrendsFrom
	trendsRange := ds.Settings.TrendsRange
	fromSec := timeRange.From.Unix()
	toSec := timeRange.To.Unix()
	rangeSec := float64(toSec - fromSec)

	if (fromSec < time.Now().Add(-trendsFrom).Unix()) || (rangeSec > trendsRange.Seconds()) {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package keychain

import (
	"sync"

	"nimona.io/pkg/crypto"
	"nimona.io/pkg/object"
	"nimona.io/pkg/peer"
)

var (
	DefaultKeychain = New()
)

type (
	Keychain interface {
		Put(keytype, crypto.PrivateKey)
		List(keytype) []crypto.PrivateKey
		ListPublicKeys(keytype) []crypto.PublicKey
		GetPrimaryPeerKey() crypto.PrivateKey
		PutCertificate(*peer.Certificate)
		GetCertificates(crypto.PublicKey) []*peer.Certificate
	}
	memorystore struct {
		keyLock        sync.RWMutex
		keys           map[keytype]map[crypto.PrivateKey]struct{}
		certLock       sync.RWMutex
		certs          map[crypto.PublicKey]map[object.Hash]*peer.Certificate
		primaryPeerKey crypto.PrivateKey
	}
)

func New() Keychain {
	return &memorystore{
		keyLock: sync.RWMutex{},
		keys: map[keytype]map[crypto.PrivateKey]struct{}{
			PeerKey:     {},
			IdentityKey: {},
		},
		certLock: sync.RWMutex{},
		certs:    map[crypto.PublicKey]map[object.Hash]*peer.Certificate{},
	}
}

func (s *memorystore) Put(t keytype, k crypto.PrivateKey) {
	s.keyLock.Lock()
	switch t {
	case PrimaryPeerKey:
		s.primaryPeerKey = k
		s.keys[PeerKey][k] = struct{}{}
	default:
		s.keys[t][k] = struct{}{}
	}
	s.keyLock.Unlock()
}

func (s *memorystore) List(t keytype) []crypto.PrivateKey {
	s.keyLock.RLock()
	ks := []crypto.PrivateKey{}
	for k := range s.keys[t] {
		ks = append(ks, k)
	}
	s.keyLock.RUnlock()
	return ks
}

func (s *memorystore) ListPublicKeys(t keytype) []crypto.PublicKey {
	s.keyLock.RLock()
	pks := []crypto.PublicKey{}
	for k := range s.keys[t] {
		pks = append(pks, k.PublicKey())
	}
	s.keyLock.RUnlock()
	return pks
}

func (s *memorystore) GetPrimaryPeerKey() crypto.PrivateKey {
	s.keyLock.RLock()
	defer s.keyLock.RUnlock()
	return s.primaryPeerKey
}

func (s *memorystore) PutCertificate(c *peer.Certificate) {
	s.certLock.Lock()
	defer s.certLock.Unlock()
	if _, ok := s.certs[c.Subject]; !ok {
		s.certs[c.Subject] = map[object.Hash]*peer.Certificate{}
	}
	h := object.NewHash(c.ToObject())
	s.certs[c.Subject][h] = c
}

func (s *memorystore) GetCertificates(
	sub crypto.PublicKey,
) []*peer.Certificate {
	cm, ok := s.certs[sub]
	if !ok {
		return []*peer.Certificate{}
	}
	cs := []*peer.Certificate{}
	for _, c := range cm {
		cs = append(cs, c)
	}
	return cs
}
<commit_msg>feat(keychain): expose default methods<commit_after>package keychain

import (
	"sync"

	"nimona.io/pkg/crypto"
	"nimona.io/pkg/object"
	"nimona.io/pkg/peer"
)

var (
	DefaultKeychain = New()
)

type (
	Keychain interface {
		Put(keytype, crypto.PrivateKey)
		List(keytype) []crypto.PrivateKey
		ListPublicKeys(keytype) []crypto.PublicKey
		GetPrimaryPeerKey() crypto.PrivateKey
		PutCertificate(*peer.Certificate)
		GetCertificates(crypto.PublicKey) []*peer.Certificate
	}
	memorystore struct {
		keyLock        sync.RWMutex
		keys           map[keytype]map[crypto.PrivateKey]struct{}
		certLock       sync.RWMutex
		certs          map[crypto.PublicKey]map[object.Hash]*peer.Certificate
		primaryPeerKey crypto.PrivateKey
	}
)

func Put(kt keytype, pk crypto.PrivateKey) {
	DefaultKeychain.Put(kt, pk)
}

func List(kt keytype) []crypto.PrivateKey {
	return DefaultKeychain.List(kt)
}

func ListPublicKeys(kt keytype) []crypto.PublicKey {
	return DefaultKeychain.ListPublicKeys(kt)
}

func GetPrimaryPeerKey() crypto.PrivateKey {
	return DefaultKeychain.GetPrimaryPeerKey()
}

func PutCertificate(cert *peer.Certificate) {
	DefaultKeychain.PutCertificate(cert)
}

func GetCertificates(pk crypto.PublicKey) []*peer.Certificate {
	return DefaultKeychain.GetCertificates(pk)
}

func New() Keychain {
	return &memorystore{
		keyLock: sync.RWMutex{},
		keys: map[keytype]map[crypto.PrivateKey]struct{}{
			PeerKey:     {},
			IdentityKey: {},
		},
		certLock: sync.RWMutex{},
		certs:    map[crypto.PublicKey]map[object.Hash]*peer.Certificate{},
	}
}

func (s *memorystore) Put(t keytype, k crypto.PrivateKey) {
	s.keyLock.Lock()
	switch t {
	case PrimaryPeerKey:
		s.primaryPeerKey = k
		s.keys[PeerKey][k] = struct{}{}
	default:
		s.keys[t][k] = struct{}{}
	}
	s.keyLock.Unlock()
}

func (s *memorystore) List(t keytype) []crypto.PrivateKey {
	s.keyLock.RLock()
	ks := []crypto.PrivateKey{}
	for k := range s.keys[t] {
		ks = append(ks, k)
	}
	s.keyLock.RUnlock()
	return ks
}

func (s *memorystore) ListPublicKeys(t keytype) []crypto.PublicKey {
	s.keyLock.RLock()
	pks := []crypto.PublicKey{}
	for k := range s.keys[t] {
		pks = append(pks, k.PublicKey())
	}
	s.keyLock.RUnlock()
	return pks
}

func (s *memorystore) GetPrimaryPeerKey() crypto.PrivateKey {
	s.keyLock.RLock()
	defer s.keyLock.RUnlock()
	return s.primaryPeerKey
}

func (s *memorystore) PutCertificate(c *peer.Certificate) {
	s.certLock.Lock()
	defer s.certLock.Unlock()
	if _, ok := s.certs[c.Subject]; !ok {
		s.certs[c.Subject] = map[object.Hash]*peer.Certificate{}
	}
	h := object.NewHash(c.ToObject())
	s.certs[c.Subject][h] = c
}

func (s *memorystore) GetCertificates(
	sub crypto.PublicKey,
) []*peer.Certificate {
	cm, ok := s.certs[sub]
	if !ok {
		return []*peer.Certificate{}
	}
	cs := []*peer.Certificate{}
	for _, c := range cm {
		cs = append(cs, c)
	}
	return cs
}
<|endoftext|>


<commit_before>/*
Copyright 2022 The Kubeflow Authors.

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

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

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

package v1beta1

import (
	"context"
	"encoding/json"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	commonv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/common/v1beta1"
	api_pb_v1beta1 "github.com/kubeflow/katib/pkg/apis/manager/v1beta1"
)

const kfpRunIDAnnotation = "kubeflow-kale.org/kfp-run-uuid"

func (k *KatibUIHandler) FetchHPJobInfo(w http.ResponseWriter, r *http.Request) {
	//enableCors(&w)
	experimentName := r.URL.Query()["experimentName"][0]
	namespace := r.URL.Query()["namespace"][0]

	conn, c := k.connectManager()
	defer conn.Close()

	resultText := "trialName,Status"
	experiment, err := k.katibClient.GetExperiment(experimentName, namespace)
	if err != nil {
		log.Printf("GetExperiment from HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	log.Printf("Got Experiment")
	metricsList := map[string]int{}
	metricsName := experiment.Spec.Objective.ObjectiveMetricName
	resultText += "," + metricsName
	metricsList[metricsName] = 0
	for i, m := range experiment.Spec.Objective.AdditionalMetricNames {
		resultText += "," + m
		metricsList[m] = i + 1
	}
	log.Printf("Got metrics names")
	paramList := map[string]int{}
	for i, p := range experiment.Spec.Parameters {
		resultText += "," + p.Name
		paramList[p.Name] = i + len(metricsList)
	}
	log.Printf("Got Parameters names")

	trialList, err := k.katibClient.GetTrialList(experimentName, namespace)
	if err != nil {
		log.Printf("GetTrialList from HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	log.Printf("Got Trial List")

	// append a column for the Pipeline UID associated with the Trial
	if havePipelineUID(trialList.Items) {
		resultText += ",KFP Run"
	}

	foundPipelineUID := false
	for _, t := range trialList.Items {
		runUid, ok := t.GetAnnotations()[kfpRunIDAnnotation]
		if !ok {
			log.Printf("Trial %s has no pipeline run.", t.Name)
			runUid = ""
		} else {
			foundPipelineUID = true
		}

		var lastTrialCondition string

		// Take only the latest condition
		if len(t.Status.Conditions) > 0 {
			lastTrialCondition = string(t.Status.Conditions[len(t.Status.Conditions)-1].Type)
		}

		trialResText := make([]string, len(metricsList)+len(paramList))

		if t.IsSucceeded() || t.IsEarlyStopped() {
			obsLogResp, err := c.GetObservationLog(
				context.Background(),
				&api_pb_v1beta1.GetObservationLogRequest{
					TrialName: t.Name,
					StartTime: "",
					EndTime:   "",
				},
			)
			if err != nil {
				log.Printf("GetObservationLog from HP job failed: %v", err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			for _, m := range obsLogResp.ObservationLog.MetricLogs {
				if trialResText[metricsList[m.Metric.Name]] == "" {
					trialResText[metricsList[m.Metric.Name]] = m.Metric.Value
				} else {
					currentValue, _ := strconv.ParseFloat(m.Metric.Value, 64)
					bestValue, _ := strconv.ParseFloat(trialResText[metricsList[m.Metric.Name]], 64)
					if t.Spec.Objective.Type == commonv1beta1.ObjectiveTypeMinimize && currentValue < bestValue {
						trialResText[metricsList[m.Metric.Name]] = m.Metric.Value
					} else if t.Spec.Objective.Type == commonv1beta1.ObjectiveTypeMaximize && currentValue > bestValue {
						trialResText[metricsList[m.Metric.Name]] = m.Metric.Value
					}
				}
			}
		}
		for _, trialParam := range t.Spec.ParameterAssignments {
			trialResText[paramList[trialParam.Name]] = trialParam.Value
		}
		resultText += "\n" + t.Name + "," + lastTrialCondition + "," + strings.Join(trialResText, ",")
		if foundPipelineUID {
			resultText += "," + runUid
		}
	}
	log.Printf("Logs parsed, results:\n %v", resultText)
	response, err := json.Marshal(resultText)
	if err != nil {
		log.Printf("Marshal result text for HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err = w.Write(response); err != nil {
		log.Printf("Write result text for HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// FetchHPJobTrialInfo returns all metrics for the HP Job Trial
func (k *KatibUIHandler) FetchHPJobTrialInfo(w http.ResponseWriter, r *http.Request) {
	//enableCors(&w)
	trialName := r.URL.Query()["trialName"][0]
	namespace := r.URL.Query()["namespace"][0]
	conn, c := k.connectManager()
	defer conn.Close()

	trial, err := k.katibClient.GetTrial(trialName, namespace)

	if err != nil {
		log.Printf("GetTrial from HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	objectiveType := trial.Spec.Objective.Type

	// resultArray - array of arrays, where [i][0] - metricName, [i][1] - metricTime, [i][2] - metricValue
	var resultArray [][]string
	resultArray = append(resultArray, strings.Split("metricName,time,value", ","))
	obsLogResp, err := c.GetObservationLog(
		context.Background(),
		&api_pb_v1beta1.GetObservationLogRequest{
			TrialName: trialName,
			StartTime: "",
			EndTime:   "",
		},
	)
	if err != nil {
		log.Printf("GetObservationLog failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// prevMetricTimeValue is the dict, where key = metric name,
	// value = array, where [0] - Last metric time, [1] - Best metric value for this time
	prevMetricTimeValue := make(map[string][]string)
	for _, m := range obsLogResp.ObservationLog.MetricLogs {
		parsedCurrentTime, _ := time.Parse(time.RFC3339Nano, m.TimeStamp)
		formatCurrentTime := parsedCurrentTime.Format("2006-01-02T15:04:05")
		if _, found := prevMetricTimeValue[m.Metric.Name]; !found {
			prevMetricTimeValue[m.Metric.Name] = []string{"", ""}

		}

		newMetricValue, err := strconv.ParseFloat(m.Metric.Value, 64)
		if err != nil {
			log.Printf("ParseFloat for new metric value: %v failed: %v", m.Metric.Value, err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		var prevMetricValue float64
		if prevMetricTimeValue[m.Metric.Name][1] != "" {
			prevMetricValue, err = strconv.ParseFloat(prevMetricTimeValue[m.Metric.Name][1], 64)
			if err != nil {
				log.Printf("ParseFloat for prev metric value: %v failed: %v", prevMetricTimeValue[m.Metric.Name][1], err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
		}

		if formatCurrentTime == prevMetricTimeValue[m.Metric.Name][0] &&
			((objectiveType == commonv1beta1.ObjectiveTypeMinimize &&
				newMetricValue < prevMetricValue) ||
				(objectiveType == commonv1beta1.ObjectiveTypeMaximize &&
					newMetricValue > prevMetricValue)) {

			prevMetricTimeValue[m.Metric.Name][1] = m.Metric.Value
			for i := len(resultArray) - 1; i >= 0; i-- {
				if resultArray[i][0] == m.Metric.Name {
					resultArray[i][2] = m.Metric.Value
					break
				}
			}
		} else if formatCurrentTime != prevMetricTimeValue[m.Metric.Name][0] {
			resultArray = append(resultArray, []string{m.Metric.Name, formatCurrentTime, m.Metric.Value})
			prevMetricTimeValue[m.Metric.Name][0] = formatCurrentTime
			prevMetricTimeValue[m.Metric.Name][1] = m.Metric.Value
		}
	}

	var resultText string
	for _, metric := range resultArray {
		resultText += strings.Join(metric, ",") + "\n"
	}

	response, err := json.Marshal(resultText)
	if err != nil {
		log.Printf("Marshal result text in Trial info failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err = w.Write(response); err != nil {
		log.Printf("Write result text in Trial info failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<commit_msg>Improve UI API/controller logging to ease troubleshooting (#1966)<commit_after>/*
Copyright 2022 The Kubeflow Authors.

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

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

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

package v1beta1

import (
	"context"
	"encoding/json"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	commonv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/common/v1beta1"
	api_pb_v1beta1 "github.com/kubeflow/katib/pkg/apis/manager/v1beta1"
)

const kfpRunIDAnnotation = "kubeflow-kale.org/kfp-run-uuid"

func (k *KatibUIHandler) FetchHPJobInfo(w http.ResponseWriter, r *http.Request) {
	//enableCors(&w)
	experimentName := r.URL.Query()["experimentName"][0]
	namespace := r.URL.Query()["namespace"][0]

	log.Printf("Start FetchHPJobInfo for Experiment: %v in namespace: %v", experimentName, namespace)

	conn, c := k.connectManager()
	defer conn.Close()

	resultText := "trialName,Status"
	experiment, err := k.katibClient.GetExperiment(experimentName, namespace)
	if err != nil {
		log.Printf("GetExperiment from HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	log.Printf("Got Experiment %v", experimentName)
	metricsList := map[string]int{}
	metricsName := experiment.Spec.Objective.ObjectiveMetricName
	resultText += "," + metricsName
	metricsList[metricsName] = 0
	for i, m := range experiment.Spec.Objective.AdditionalMetricNames {
		resultText += "," + m
		metricsList[m] = i + 1
	}
	log.Printf("Got metrics names")
	paramList := map[string]int{}
	for i, p := range experiment.Spec.Parameters {
		resultText += "," + p.Name
		paramList[p.Name] = i + len(metricsList)
	}
	log.Printf("Got Parameters names")

	trialList, err := k.katibClient.GetTrialList(experimentName, namespace)
	if err != nil {
		log.Printf("GetTrialList from HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	log.Printf("Got Trial List - Count: %v", len(trialList.Items))

	// append a column for the Pipeline UID associated with the Trial
	if havePipelineUID(trialList.Items) {
		resultText += ",KFP Run"
	}

	foundPipelineUID := false
	for _, t := range trialList.Items {
		runUid, ok := t.GetAnnotations()[kfpRunIDAnnotation]
		if !ok {
			log.Printf("Trial %s has no pipeline run.", t.Name)
			runUid = ""
		} else {
			foundPipelineUID = true
		}

		var lastTrialCondition string

		// Take only the latest condition
		if len(t.Status.Conditions) > 0 {
			lastTrialCondition = string(t.Status.Conditions[len(t.Status.Conditions)-1].Type)
		}

		trialResText := make([]string, len(metricsList)+len(paramList))

		if t.IsSucceeded() || t.IsEarlyStopped() {
			obsLogResp, err := c.GetObservationLog(
				context.Background(),
				&api_pb_v1beta1.GetObservationLogRequest{
					TrialName: t.Name,
					StartTime: "",
					EndTime:   "",
				},
			)
			if err != nil {
				log.Printf("GetObservationLog from HP job failed: %v", err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			for _, m := range obsLogResp.ObservationLog.MetricLogs {
				if trialResText[metricsList[m.Metric.Name]] == "" {
					trialResText[metricsList[m.Metric.Name]] = m.Metric.Value
				} else {
					currentValue, _ := strconv.ParseFloat(m.Metric.Value, 64)
					bestValue, _ := strconv.ParseFloat(trialResText[metricsList[m.Metric.Name]], 64)
					if t.Spec.Objective.Type == commonv1beta1.ObjectiveTypeMinimize && currentValue < bestValue {
						trialResText[metricsList[m.Metric.Name]] = m.Metric.Value
					} else if t.Spec.Objective.Type == commonv1beta1.ObjectiveTypeMaximize && currentValue > bestValue {
						trialResText[metricsList[m.Metric.Name]] = m.Metric.Value
					}
				}
			}
		}
		for _, trialParam := range t.Spec.ParameterAssignments {
			trialResText[paramList[trialParam.Name]] = trialParam.Value
		}
		resultText += "\n" + t.Name + "," + lastTrialCondition + "," + strings.Join(trialResText, ",")
		if foundPipelineUID {
			resultText += "," + runUid
		}
	}
	log.Printf("Logs parsed, results:\n %v", resultText)
	response, err := json.Marshal(resultText)
	if err != nil {
		log.Printf("Marshal result text for HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err = w.Write(response); err != nil {
		log.Printf("Write result text for HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// FetchHPJobTrialInfo returns all metrics for the HP Job Trial
func (k *KatibUIHandler) FetchHPJobTrialInfo(w http.ResponseWriter, r *http.Request) {
	//enableCors(&w)
	trialName := r.URL.Query()["trialName"][0]
	namespace := r.URL.Query()["namespace"][0]
	conn, c := k.connectManager()
	defer conn.Close()

	trial, err := k.katibClient.GetTrial(trialName, namespace)

	if err != nil {
		log.Printf("GetTrial from HP job failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	objectiveType := trial.Spec.Objective.Type

	// resultArray - array of arrays, where [i][0] - metricName, [i][1] - metricTime, [i][2] - metricValue
	var resultArray [][]string
	resultArray = append(resultArray, strings.Split("metricName,time,value", ","))
	obsLogResp, err := c.GetObservationLog(
		context.Background(),
		&api_pb_v1beta1.GetObservationLogRequest{
			TrialName: trialName,
			StartTime: "",
			EndTime:   "",
		},
	)
	if err != nil {
		log.Printf("GetObservationLog failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// prevMetricTimeValue is the dict, where key = metric name,
	// value = array, where [0] - Last metric time, [1] - Best metric value for this time
	prevMetricTimeValue := make(map[string][]string)
	for _, m := range obsLogResp.ObservationLog.MetricLogs {
		parsedCurrentTime, _ := time.Parse(time.RFC3339Nano, m.TimeStamp)
		formatCurrentTime := parsedCurrentTime.Format("2006-01-02T15:04:05")
		if _, found := prevMetricTimeValue[m.Metric.Name]; !found {
			prevMetricTimeValue[m.Metric.Name] = []string{"", ""}

		}

		newMetricValue, err := strconv.ParseFloat(m.Metric.Value, 64)
		if err != nil {
			log.Printf("ParseFloat for new metric value: %v failed: %v", m.Metric.Value, err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		var prevMetricValue float64
		if prevMetricTimeValue[m.Metric.Name][1] != "" {
			prevMetricValue, err = strconv.ParseFloat(prevMetricTimeValue[m.Metric.Name][1], 64)
			if err != nil {
				log.Printf("ParseFloat for prev metric value: %v failed: %v", prevMetricTimeValue[m.Metric.Name][1], err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
		}

		if formatCurrentTime == prevMetricTimeValue[m.Metric.Name][0] &&
			((objectiveType == commonv1beta1.ObjectiveTypeMinimize &&
				newMetricValue < prevMetricValue) ||
				(objectiveType == commonv1beta1.ObjectiveTypeMaximize &&
					newMetricValue > prevMetricValue)) {

			prevMetricTimeValue[m.Metric.Name][1] = m.Metric.Value
			for i := len(resultArray) - 1; i >= 0; i-- {
				if resultArray[i][0] == m.Metric.Name {
					resultArray[i][2] = m.Metric.Value
					break
				}
			}
		} else if formatCurrentTime != prevMetricTimeValue[m.Metric.Name][0] {
			resultArray = append(resultArray, []string{m.Metric.Name, formatCurrentTime, m.Metric.Value})
			prevMetricTimeValue[m.Metric.Name][0] = formatCurrentTime
			prevMetricTimeValue[m.Metric.Name][1] = m.Metric.Value
		}
	}

	var resultText string
	for _, metric := range resultArray {
		resultText += strings.Join(metric, ",") + "\n"
	}

	response, err := json.Marshal(resultText)
	if err != nil {
		log.Printf("Marshal result text in Trial info failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err = w.Write(response); err != nil {
		log.Printf("Write result text in Trial info failed: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package server

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/errdefs"
	containerdimages "github.com/containerd/containerd/images"
	"github.com/containerd/containerd/reference"
	"github.com/containerd/containerd/remotes"
	"github.com/containerd/containerd/remotes/docker"
	"github.com/containerd/cri/pkg/config"
	distribution "github.com/docker/distribution/reference"
	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
)

// For image management:
// 1) We have an in-memory metadata index to:
//   a. Maintain ImageID -> RepoTags, ImageID -> RepoDigset relationships; ImageID
//   is the digest of image config, which conforms to oci image spec.
//   b. Cache constant and useful information such as image chainID, config etc.
//   c. An image will be added into the in-memory metadata only when it's successfully
//   pulled and unpacked.
//
// 2) We use containerd image metadata store and content store:
//   a. To resolve image reference (digest/tag) locally. During pulling image, we
//   normalize the image reference provided by user, and put it into image metadata
//   store with resolved descriptor. For the other operations, if image id is provided,
//   we'll access the in-memory metadata index directly; if image reference is
//   provided, we'll normalize it, resolve it in containerd image metadata store
//   to get the image id.
//   b. As the backup of in-memory metadata in 1). During startup, the in-memory
//   metadata could be re-constructed from image metadata store + content store.
//
// Several problems with current approach:
// 1) An entry in containerd image metadata store doesn't mean a "READY" (successfully
// pulled and unpacked) image. E.g. during pulling, the client gets killed. In that case,
// if we saw an image without snapshots or with in-complete contents during startup,
// should we re-pull the image? Or should we remove the entry?
//
// yanxuean: We can't delete image directly, because we don't know if the image
// is pulled by us. There are resource leakage.
//
// 2) Containerd suggests user to add entry before pulling the image. However if
// an error occurs during the pulling, should we remove the entry from metadata
// store? Or should we leave it there until next startup (resource leakage)?
//
// 3) The cri plugin only exposes "READY" (successfully pulled and unpacked) images
// to the user, which are maintained in the in-memory metadata index. However, it's
// still possible that someone else removes the content or snapshot by-pass the cri plugin,
// how do we detect that and update the in-memory metadata correspondingly? Always
// check whether corresponding snapshot is ready when reporting image status?
//
// 4) Is the content important if we cached necessary information in-memory
// after we pull the image? How to manage the disk usage of contents? If some
// contents are missing but snapshots are ready, is the image still "READY"?

// PullImage pulls an image with authentication config.
func (c *criService) PullImage(ctx context.Context, r *runtime.PullImageRequest) (*runtime.PullImageResponse, error) {
	imageRef := r.GetImage().GetImage()
	namedRef, err := distribution.ParseDockerRef(imageRef)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to parse image reference %q", imageRef)
	}
	ref := namedRef.String()
	if ref != imageRef {
		logrus.Debugf("PullImage using normalized image ref: %q", ref)
	}
	resolver, desc, err := c.getResolver(ctx, ref, c.credentials(r.GetAuth()))
	if err != nil {
		return nil, errors.Wrapf(err, "failed to resolve image %q", ref)
	}
	// We have to check schema1 here, because after `Pull`, schema1
	// image has already been converted.
	isSchema1 := desc.MediaType == containerdimages.MediaTypeDockerSchema1Manifest

	image, err := c.client.Pull(ctx, ref,
		containerd.WithSchema1Conversion,
		containerd.WithResolver(resolver),
		containerd.WithPullSnapshotter(c.config.ContainerdConfig.Snapshotter),
		containerd.WithPullUnpack,
	)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to pull and unpack image %q", ref)
	}

	configDesc, err := image.Config(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "get image config descriptor")
	}
	imageID := configDesc.Digest.String()

	repoDigest, repoTag := getRepoDigestAndTag(namedRef, image.Target().Digest, isSchema1)
	for _, r := range []string{imageID, repoTag, repoDigest} {
		if r == "" {
			continue
		}
		if err := c.createImageReference(ctx, r, image.Target()); err != nil {
			return nil, errors.Wrapf(err, "failed to create image reference %q", r)
		}
		// Update image store to reflect the newest state in containerd.
		// No need to use `updateImage`, because the image reference must
		// have been managed by the cri plugin.
		if err := c.imageStore.Update(ctx, r); err != nil {
			return nil, errors.Wrapf(err, "failed to update image store %q", r)
		}
	}

	logrus.Debugf("Pulled image %q with image id %q, repo tag %q, repo digest %q", imageRef, imageID,
		repoTag, repoDigest)
	// NOTE(random-liu): the actual state in containerd is the source of truth, even we maintain
	// in-memory image store, it's only for in-memory indexing. The image could be removed
	// by someone else anytime, before/during/after we create the metadata. We should always
	// check the actual state in containerd before using the image or returning status of the
	// image.
	return &runtime.PullImageResponse{ImageRef: imageID}, nil
}

// ParseAuth parses AuthConfig and returns username and password/secret required by containerd.
func ParseAuth(auth *runtime.AuthConfig) (string, string, error) {
	if auth == nil {
		return "", "", nil
	}
	if auth.Username != "" {
		return auth.Username, auth.Password, nil
	}
	if auth.IdentityToken != "" {
		return "", auth.IdentityToken, nil
	}
	if auth.Auth != "" {
		decLen := base64.StdEncoding.DecodedLen(len(auth.Auth))
		decoded := make([]byte, decLen)
		_, err := base64.StdEncoding.Decode(decoded, []byte(auth.Auth))
		if err != nil {
			return "", "", err
		}
		fields := strings.SplitN(string(decoded), ":", 2)
		if len(fields) != 2 {
			return "", "", errors.Errorf("invalid decoded auth: %q", decoded)
		}
		user, passwd := fields[0], fields[1]
		return user, strings.Trim(passwd, "\x00"), nil
	}
	// TODO(random-liu): Support RegistryToken.
	return "", "", errors.New("invalid auth config")
}

// createImageReference creates image reference inside containerd image store.
// Note that because create and update are not finished in one transaction, there could be race. E.g.
// the image reference is deleted by someone else after create returns already exists, but before update
// happens.
func (c *criService) createImageReference(ctx context.Context, name string, desc imagespec.Descriptor) error {
	img := containerdimages.Image{
		Name:   name,
		Target: desc,
		// Add a label to indicate that the image is managed by the cri plugin.
		Labels: map[string]string{imageLabelKey: imageLabelValue},
	}
	// TODO(random-liu): Figure out which is the more performant sequence create then update or
	// update then create.
	oldImg, err := c.client.ImageService().Create(ctx, img)
	if err == nil || !errdefs.IsAlreadyExists(err) {
		return err
	}
	if oldImg.Target.Digest == img.Target.Digest && oldImg.Labels[imageLabelKey] == imageLabelValue {
		return nil
	}
	_, err = c.client.ImageService().Update(ctx, img, "target", "labels")
	return err
}

// updateImage updates image store to reflect the newest state of an image reference
// in containerd. If the reference is not managed by the cri plugin, the function also
// generates necessary metadata for the image and make it managed.
func (c *criService) updateImage(ctx context.Context, r string) error {
	img, err := c.client.GetImage(ctx, r)
	if err != nil && !errdefs.IsNotFound(err) {
		return errors.Wrap(err, "get image by reference")
	}
	if err == nil && img.Labels()[imageLabelKey] != imageLabelValue {
		// Make sure the image has the image id as its unique
		// identifier that references the image in its lifetime.
		configDesc, err := img.Config(ctx)
		if err != nil {
			return errors.Wrap(err, "get image id")
		}
		id := configDesc.Digest.String()
		if err := c.createImageReference(ctx, id, img.Target()); err != nil {
			return errors.Wrapf(err, "create image id reference %q", id)
		}
		if err := c.imageStore.Update(ctx, id); err != nil {
			return errors.Wrapf(err, "update image store for %q", id)
		}
		// The image id is ready, add the label to mark the image as managed.
		if err := c.createImageReference(ctx, r, img.Target()); err != nil {
			return errors.Wrap(err, "create managed label")
		}
	}
	// If the image is not found, we should continue updating the cache,
	// so that the image can be removed from the cache.
	if err := c.imageStore.Update(ctx, r); err != nil {
		return errors.Wrapf(err, "update image store for %q", r)
	}
	return nil
}

// credentials returns a credential function for docker resolver to use.
func (c *criService) credentials(auth *runtime.AuthConfig) func(string) (string, string, error) {
	return func(host string) (string, string, error) {
		if auth == nil {
			// Get default auth from config.
			for h, ac := range c.config.Registry.Auths {
				u, err := url.Parse(h)
				if err != nil {
					return "", "", errors.Wrapf(err, "parse auth host %q", h)
				}
				if u.Host == host {
					auth = toRuntimeAuthConfig(ac)
					break
				}
			}
		}
		return ParseAuth(auth)
	}
}

// getRegistryTLSTransport returns a http.Transport configured with a CA/Cert/Key specified by registryTLSConfig
func (c *criService) getRegistryTLSTransport(ctx context.Context, registryTLSConfig config.TLSConfig) (*http.Transport, error) {
	cert, err := tls.LoadX509KeyPair(registryTLSConfig.CertFile, registryTLSConfig.KeyFile)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load cert file")
	}

	caCertPool, err := x509.SystemCertPool()
	if err != nil {
		return nil, errors.Wrap(err, "failed to get system cert pool")
	}
	caCert, err := ioutil.ReadFile(registryTLSConfig.CAFile)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load CA file")
	}
	caCertPool.AppendCertsFromPEM(caCert)

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs:      caCertPool,
	}
	tlsConfig.BuildNameToCertificate()
	transport := &http.Transport{TLSClientConfig: tlsConfig}

	return transport, nil
}

// getResolver tries registry mirrors and the default registry, and returns the resolver and descriptor
// from the first working registry.
func (c *criService) getResolver(ctx context.Context, ref string, cred func(string) (string, string, error)) (remotes.Resolver, imagespec.Descriptor, error) {
	refspec, err := reference.Parse(ref)
	if err != nil {
		return nil, imagespec.Descriptor{}, errors.Wrap(err, "parse image reference")
	}

	httpClient := http.DefaultClient

	// Try mirrors in order first, and then try default host name.
	for _, e := range c.config.Registry.Mirrors[refspec.Hostname()].Endpoints {
		u, err := url.Parse(e)
		if err != nil {
			return nil, imagespec.Descriptor{}, errors.Wrapf(err, "parse registry endpoint %q", e)
		}

		if registryTLSConfig, ok := c.config.Registry.TLSConfigs[u.Host]; ok {
			tlsTransport, err := c.getRegistryTLSTransport(ctx, registryTLSConfig)
			if err != nil {
				return nil, imagespec.Descriptor{}, errors.Wrapf(err, "get tlsTransport for registry %q", refspec.Hostname())
			}
			httpClient = &http.Client{Transport: tlsTransport}
		}

		resolver := docker.NewResolver(docker.ResolverOptions{
			Authorizer: docker.NewAuthorizer(http.DefaultClient, cred),
			Client:     httpClient,
			Host:       func(string) (string, error) { return u.Host, nil },
			// By default use "https".
			PlainHTTP: u.Scheme == "http",
		})
		_, desc, err := resolver.Resolve(ctx, ref)
		if err == nil {
			return resolver, desc, nil
		}
		logrus.WithError(err).Debugf("Tried registry mirror %q but failed", e)
		// Continue to try next endpoint
	}

	hostname, err := docker.DefaultHost(refspec.Hostname())
	if err != nil {
		return nil, imagespec.Descriptor{}, errors.Wrapf(err, "get host for refspec %q", refspec.Hostname())
	}
	if registryTLSConfig, ok := c.config.Registry.TLSConfigs[hostname]; ok {
		tlsTransport, err := c.getRegistryTLSTransport(ctx, registryTLSConfig)
		if err != nil {
			return nil, imagespec.Descriptor{}, errors.Wrapf(err, "get tlsTransport for registry %q", refspec.Hostname())
		}
		httpClient = &http.Client{Transport: tlsTransport}
	}

	resolver := docker.NewResolver(docker.ResolverOptions{
		Credentials: cred,
		Client:      httpClient,
	})
	_, desc, err := resolver.Resolve(ctx, ref)
	if err != nil {
		return nil, imagespec.Descriptor{}, errors.Wrap(err, "no available registry endpoint")
	}
	return resolver, desc, nil
}
<commit_msg>Add cri managed image label when pulling the image.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package server

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/errdefs"
	containerdimages "github.com/containerd/containerd/images"
	"github.com/containerd/containerd/reference"
	"github.com/containerd/containerd/remotes"
	"github.com/containerd/containerd/remotes/docker"
	"github.com/containerd/cri/pkg/config"
	distribution "github.com/docker/distribution/reference"
	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
)

// For image management:
// 1) We have an in-memory metadata index to:
//   a. Maintain ImageID -> RepoTags, ImageID -> RepoDigset relationships; ImageID
//   is the digest of image config, which conforms to oci image spec.
//   b. Cache constant and useful information such as image chainID, config etc.
//   c. An image will be added into the in-memory metadata only when it's successfully
//   pulled and unpacked.
//
// 2) We use containerd image metadata store and content store:
//   a. To resolve image reference (digest/tag) locally. During pulling image, we
//   normalize the image reference provided by user, and put it into image metadata
//   store with resolved descriptor. For the other operations, if image id is provided,
//   we'll access the in-memory metadata index directly; if image reference is
//   provided, we'll normalize it, resolve it in containerd image metadata store
//   to get the image id.
//   b. As the backup of in-memory metadata in 1). During startup, the in-memory
//   metadata could be re-constructed from image metadata store + content store.
//
// Several problems with current approach:
// 1) An entry in containerd image metadata store doesn't mean a "READY" (successfully
// pulled and unpacked) image. E.g. during pulling, the client gets killed. In that case,
// if we saw an image without snapshots or with in-complete contents during startup,
// should we re-pull the image? Or should we remove the entry?
//
// yanxuean: We can't delete image directly, because we don't know if the image
// is pulled by us. There are resource leakage.
//
// 2) Containerd suggests user to add entry before pulling the image. However if
// an error occurs during the pulling, should we remove the entry from metadata
// store? Or should we leave it there until next startup (resource leakage)?
//
// 3) The cri plugin only exposes "READY" (successfully pulled and unpacked) images
// to the user, which are maintained in the in-memory metadata index. However, it's
// still possible that someone else removes the content or snapshot by-pass the cri plugin,
// how do we detect that and update the in-memory metadata correspondingly? Always
// check whether corresponding snapshot is ready when reporting image status?
//
// 4) Is the content important if we cached necessary information in-memory
// after we pull the image? How to manage the disk usage of contents? If some
// contents are missing but snapshots are ready, is the image still "READY"?

// PullImage pulls an image with authentication config.
func (c *criService) PullImage(ctx context.Context, r *runtime.PullImageRequest) (*runtime.PullImageResponse, error) {
	imageRef := r.GetImage().GetImage()
	namedRef, err := distribution.ParseDockerRef(imageRef)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to parse image reference %q", imageRef)
	}
	ref := namedRef.String()
	if ref != imageRef {
		logrus.Debugf("PullImage using normalized image ref: %q", ref)
	}
	resolver, desc, err := c.getResolver(ctx, ref, c.credentials(r.GetAuth()))
	if err != nil {
		return nil, errors.Wrapf(err, "failed to resolve image %q", ref)
	}
	// We have to check schema1 here, because after `Pull`, schema1
	// image has already been converted.
	isSchema1 := desc.MediaType == containerdimages.MediaTypeDockerSchema1Manifest

	image, err := c.client.Pull(ctx, ref,
		containerd.WithSchema1Conversion,
		containerd.WithResolver(resolver),
		containerd.WithPullSnapshotter(c.config.ContainerdConfig.Snapshotter),
		containerd.WithPullUnpack,
		containerd.WithPullLabel(imageLabelKey, imageLabelValue),
	)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to pull and unpack image %q", ref)
	}

	configDesc, err := image.Config(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "get image config descriptor")
	}
	imageID := configDesc.Digest.String()

	repoDigest, repoTag := getRepoDigestAndTag(namedRef, image.Target().Digest, isSchema1)
	for _, r := range []string{imageID, repoTag, repoDigest} {
		if r == "" {
			continue
		}
		if err := c.createImageReference(ctx, r, image.Target()); err != nil {
			return nil, errors.Wrapf(err, "failed to create image reference %q", r)
		}
		// Update image store to reflect the newest state in containerd.
		// No need to use `updateImage`, because the image reference must
		// have been managed by the cri plugin.
		if err := c.imageStore.Update(ctx, r); err != nil {
			return nil, errors.Wrapf(err, "failed to update image store %q", r)
		}
	}

	logrus.Debugf("Pulled image %q with image id %q, repo tag %q, repo digest %q", imageRef, imageID,
		repoTag, repoDigest)
	// NOTE(random-liu): the actual state in containerd is the source of truth, even we maintain
	// in-memory image store, it's only for in-memory indexing. The image could be removed
	// by someone else anytime, before/during/after we create the metadata. We should always
	// check the actual state in containerd before using the image or returning status of the
	// image.
	return &runtime.PullImageResponse{ImageRef: imageID}, nil
}

// ParseAuth parses AuthConfig and returns username and password/secret required by containerd.
func ParseAuth(auth *runtime.AuthConfig) (string, string, error) {
	if auth == nil {
		return "", "", nil
	}
	if auth.Username != "" {
		return auth.Username, auth.Password, nil
	}
	if auth.IdentityToken != "" {
		return "", auth.IdentityToken, nil
	}
	if auth.Auth != "" {
		decLen := base64.StdEncoding.DecodedLen(len(auth.Auth))
		decoded := make([]byte, decLen)
		_, err := base64.StdEncoding.Decode(decoded, []byte(auth.Auth))
		if err != nil {
			return "", "", err
		}
		fields := strings.SplitN(string(decoded), ":", 2)
		if len(fields) != 2 {
			return "", "", errors.Errorf("invalid decoded auth: %q", decoded)
		}
		user, passwd := fields[0], fields[1]
		return user, strings.Trim(passwd, "\x00"), nil
	}
	// TODO(random-liu): Support RegistryToken.
	return "", "", errors.New("invalid auth config")
}

// createImageReference creates image reference inside containerd image store.
// Note that because create and update are not finished in one transaction, there could be race. E.g.
// the image reference is deleted by someone else after create returns already exists, but before update
// happens.
func (c *criService) createImageReference(ctx context.Context, name string, desc imagespec.Descriptor) error {
	img := containerdimages.Image{
		Name:   name,
		Target: desc,
		// Add a label to indicate that the image is managed by the cri plugin.
		Labels: map[string]string{imageLabelKey: imageLabelValue},
	}
	// TODO(random-liu): Figure out which is the more performant sequence create then update or
	// update then create.
	oldImg, err := c.client.ImageService().Create(ctx, img)
	if err == nil || !errdefs.IsAlreadyExists(err) {
		return err
	}
	if oldImg.Target.Digest == img.Target.Digest && oldImg.Labels[imageLabelKey] == imageLabelValue {
		return nil
	}
	_, err = c.client.ImageService().Update(ctx, img, "target", "labels")
	return err
}

// updateImage updates image store to reflect the newest state of an image reference
// in containerd. If the reference is not managed by the cri plugin, the function also
// generates necessary metadata for the image and make it managed.
func (c *criService) updateImage(ctx context.Context, r string) error {
	img, err := c.client.GetImage(ctx, r)
	if err != nil && !errdefs.IsNotFound(err) {
		return errors.Wrap(err, "get image by reference")
	}
	if err == nil && img.Labels()[imageLabelKey] != imageLabelValue {
		// Make sure the image has the image id as its unique
		// identifier that references the image in its lifetime.
		configDesc, err := img.Config(ctx)
		if err != nil {
			return errors.Wrap(err, "get image id")
		}
		id := configDesc.Digest.String()
		if err := c.createImageReference(ctx, id, img.Target()); err != nil {
			return errors.Wrapf(err, "create image id reference %q", id)
		}
		if err := c.imageStore.Update(ctx, id); err != nil {
			return errors.Wrapf(err, "update image store for %q", id)
		}
		// The image id is ready, add the label to mark the image as managed.
		if err := c.createImageReference(ctx, r, img.Target()); err != nil {
			return errors.Wrap(err, "create managed label")
		}
	}
	// If the image is not found, we should continue updating the cache,
	// so that the image can be removed from the cache.
	if err := c.imageStore.Update(ctx, r); err != nil {
		return errors.Wrapf(err, "update image store for %q", r)
	}
	return nil
}

// credentials returns a credential function for docker resolver to use.
func (c *criService) credentials(auth *runtime.AuthConfig) func(string) (string, string, error) {
	return func(host string) (string, string, error) {
		if auth == nil {
			// Get default auth from config.
			for h, ac := range c.config.Registry.Auths {
				u, err := url.Parse(h)
				if err != nil {
					return "", "", errors.Wrapf(err, "parse auth host %q", h)
				}
				if u.Host == host {
					auth = toRuntimeAuthConfig(ac)
					break
				}
			}
		}
		return ParseAuth(auth)
	}
}

// getRegistryTLSTransport returns a http.Transport configured with a CA/Cert/Key specified by registryTLSConfig
func (c *criService) getRegistryTLSTransport(ctx context.Context, registryTLSConfig config.TLSConfig) (*http.Transport, error) {
	cert, err := tls.LoadX509KeyPair(registryTLSConfig.CertFile, registryTLSConfig.KeyFile)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load cert file")
	}

	caCertPool, err := x509.SystemCertPool()
	if err != nil {
		return nil, errors.Wrap(err, "failed to get system cert pool")
	}
	caCert, err := ioutil.ReadFile(registryTLSConfig.CAFile)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load CA file")
	}
	caCertPool.AppendCertsFromPEM(caCert)

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs:      caCertPool,
	}
	tlsConfig.BuildNameToCertificate()
	transport := &http.Transport{TLSClientConfig: tlsConfig}

	return transport, nil
}

// getResolver tries registry mirrors and the default registry, and returns the resolver and descriptor
// from the first working registry.
func (c *criService) getResolver(ctx context.Context, ref string, cred func(string) (string, string, error)) (remotes.Resolver, imagespec.Descriptor, error) {
	refspec, err := reference.Parse(ref)
	if err != nil {
		return nil, imagespec.Descriptor{}, errors.Wrap(err, "parse image reference")
	}

	httpClient := http.DefaultClient

	// Try mirrors in order first, and then try default host name.
	for _, e := range c.config.Registry.Mirrors[refspec.Hostname()].Endpoints {
		u, err := url.Parse(e)
		if err != nil {
			return nil, imagespec.Descriptor{}, errors.Wrapf(err, "parse registry endpoint %q", e)
		}

		if registryTLSConfig, ok := c.config.Registry.TLSConfigs[u.Host]; ok {
			tlsTransport, err := c.getRegistryTLSTransport(ctx, registryTLSConfig)
			if err != nil {
				return nil, imagespec.Descriptor{}, errors.Wrapf(err, "get tlsTransport for registry %q", refspec.Hostname())
			}
			httpClient = &http.Client{Transport: tlsTransport}
		}

		resolver := docker.NewResolver(docker.ResolverOptions{
			Authorizer: docker.NewAuthorizer(http.DefaultClient, cred),
			Client:     httpClient,
			Host:       func(string) (string, error) { return u.Host, nil },
			// By default use "https".
			PlainHTTP: u.Scheme == "http",
		})
		_, desc, err := resolver.Resolve(ctx, ref)
		if err == nil {
			return resolver, desc, nil
		}
		logrus.WithError(err).Debugf("Tried registry mirror %q but failed", e)
		// Continue to try next endpoint
	}

	hostname, err := docker.DefaultHost(refspec.Hostname())
	if err != nil {
		return nil, imagespec.Descriptor{}, errors.Wrapf(err, "get host for refspec %q", refspec.Hostname())
	}
	if registryTLSConfig, ok := c.config.Registry.TLSConfigs[hostname]; ok {
		tlsTransport, err := c.getRegistryTLSTransport(ctx, registryTLSConfig)
		if err != nil {
			return nil, imagespec.Descriptor{}, errors.Wrapf(err, "get tlsTransport for registry %q", refspec.Hostname())
		}
		httpClient = &http.Client{Transport: tlsTransport}
	}

	resolver := docker.NewResolver(docker.ResolverOptions{
		Credentials: cred,
		Client:      httpClient,
	})
	_, desc, err := resolver.Resolve(ctx, ref)
	if err != nil {
		return nil, imagespec.Descriptor{}, errors.Wrap(err, "no available registry endpoint")
	}
	return resolver, desc, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Mirantis

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

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

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

package stream

import (
	"io/ioutil"
	"net"
	"os"
	"runtime"
	"testing"
	"time"

	testutils "github.com/Mirantis/virtlet/pkg/utils/testing"
)

func TestGetPidFromConnection(t *testing.T) {
	if runtime.GOOS != "linux" {
		t.Skip("getPidFromConnection only works on Linux")
	}
	socket, err := ioutil.TempFile(os.TempDir(), "utilTest")
	defer os.Remove(socket.Name())
	socketPath := socket.Name()

	tc := testutils.RunProcess(t, "nc", []string{"-l", "-U", socketPath}, nil)
	defer tc.Stop()

	// wait for nc to start
	time.Sleep(2 * time.Second)

	conn, err := net.DialUnix("unix", nil, &net.UnixAddr{Name: socketPath, Net: "unix"})
	if err != nil {
		t.Fatal("Error when connecting to socket:", err)
	}

	pid, err := getPidFromConnection(conn)

	if err != nil {
		t.Errorf("Couldn't get pid from Unix socket: %v", err)
	}
	if pid != int32(tc.Pid()) {
		t.Errorf("Wrong pid from getPidFromConnection. Expected: %d, got %d", tc.Pid(), pid)
	}
}

func TestGetProcessEnvironment(t *testing.T) {
	if runtime.GOOS != "linux" {
		t.Skip("getProcessEnvironment only works on Linux")
	}

	tc := testutils.RunProcess(t, "sleep", []string{"10"}, []string{
		"FOO=1",
		"BAR=asd",
	})
	defer tc.Stop()
	env, err := getProcessEnvironment(int32(tc.Pid()))
	if err != nil {
		t.Error(err)
	}
	for k, v := range map[string]string{"FOO": "1", "BAR": "asd"} {
		envVal, ok := env[k]
		if !ok {
			t.Errorf("%s variable not found in env", k)
		}
		if envVal != v {
			t.Errorf("%s variable value not equal to %s", k, v)
		}
	}
}
<commit_msg>Fix flake in streamer utils test<commit_after>/*
Copyright 2017 Mirantis

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

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

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

package stream

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"runtime"
	"testing"
	"time"

	testutils "github.com/Mirantis/virtlet/pkg/utils/testing"
)

func TestGetPidFromConnection(t *testing.T) {
	if runtime.GOOS != "linux" {
		t.Skip("getPidFromConnection only works on Linux")
	}
	socket, err := ioutil.TempFile(os.TempDir(), "utilTest")
	defer os.Remove(socket.Name())
	socketPath := socket.Name()

	tc := testutils.RunProcess(t, "nc", []string{"-l", "-U", socketPath}, nil)
	defer tc.Stop()

	// wait for nc to start
	time.Sleep(2 * time.Second)

	conn, err := net.DialUnix("unix", nil, &net.UnixAddr{Name: socketPath, Net: "unix"})
	if err != nil {
		t.Fatal("Error when connecting to socket:", err)
	}

	pid, err := getPidFromConnection(conn)

	if err != nil {
		t.Errorf("Couldn't get pid from Unix socket: %v", err)
	}
	if pid != int32(tc.Pid()) {
		t.Errorf("Wrong pid from getPidFromConnection. Expected: %d, got %d", tc.Pid(), pid)
	}
}

func TestGetProcessEnvironment(t *testing.T) {
	if runtime.GOOS != "linux" {
		t.Skip("getProcessEnvironment only works on Linux")
	}

	tc := testutils.RunProcess(t, "sleep", []string{"10"}, []string{
		"FOO=1",
		"BAR=asd",
	})
	defer tc.Stop()
	waitForChildExec(t, tc.Pid())
	env, err := getProcessEnvironment(int32(tc.Pid()))
	if err != nil {
		t.Error(err)
	}
	for k, v := range map[string]string{"FOO": "1", "BAR": "asd"} {
		envVal, ok := env[k]
		if !ok {
			t.Errorf("%s variable not found in env", k)
		}
		if envVal != v {
			t.Errorf("%s variable value not equal to %s", k, v)
		}
	}
}

// waitForChildExec waits for a child process to perform exec()
// by means of waiting when its `/proc/PID/cmdline` becomes different from the
// cmdline of the current process
func waitForChildExec(t *testing.T, pid int) {
	ownPid := os.Getpid()
	ownCmdline, err := ioutil.ReadFile(fmt.Sprintf("/proc/%d/cmdline", ownPid))
	if err != nil {
		t.Fatalf("cannot read own cmdline: %v", err)
	}
	processCmdlineLocation := fmt.Sprintf("/proc/%d/cmdline", pid)
	for {
		processCmdline, err := ioutil.ReadFile(processCmdlineLocation)
		if err != nil {
			t.Fatalf("cannot read cmdline for pid %q: %v", pid, err)
		}
		if string(processCmdline) != string(ownCmdline) {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}
}
<|endoftext|>