<commit_before>package api

import (
	"github.com/zenoss/glog"
	"github.com/zenoss/serviced"
	"github.com/zenoss/serviced/dao"
	"github.com/zenoss/serviced/shell"

	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var proxyOptions ProxyOptions

// ProxyOptions are options to be run when starting a new proxy server
type ProxyOptions struct {
	MuxPort          int
	Mux              bool
	TLS              bool
	KeyPEMFile       string
	CertPEMFile      string
	ServicedEndpoint string
	Autorestart      bool
	Logstash         bool
}

// LoadProxyOptions loads the proxy option information
func LoadProxyOptions(ops ProxyOptions) {
	proxyOptions = ops
}

// ProxyConfig is the config object for starting a proxy server
type ProxyConfig struct {
	ServiceID string
	Command   []string
}

func getNextIP() string {
	s := fmt.Sprintf("10.3.0.%d", nextip)
	nextip += 1
	return s
}

// Start a service proxy
func (a *api) StartProxy(cfg ProxyConfig) error {
	config := serviced.MuxConfig{}
	config.TCPMux.Port = proxyOptions.MuxPort
	config.TCPMux.Enabled = proxyOptions.Mux
	config.TCPMux.UseTLS = proxyOptions.TLS
	config.ServiceId = cfg.ServiceID
	config.Command = strings.Join(cfg.Command, " ")

	if config.TCPMux.Enabled {
		go config.TCPMux.ListenAndMux()
	}

	sio := shell.NewProcessForwarderServer(proxyOptions.ServicedEndpoint)
	sio.Handle("/", http.FileServer(http.Dir("/serviced/www/")))
	go http.ListenAndServe(":50000", sio)

	procexit := make(chan int)
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	// continually execute subprocess
	go func(cmdString string) {
		defer func() { procexit <- 0 }()
		for {
			glog.V(0).Info("About to execute: ", cmdString)
			cmd := exec.Command("bash", "-c", cmdString)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			cmd.Stdin = os.Stdin
			serviceExit := make(chan error)
			go func() {
				serviceExit <- cmd.Run()
			}()
			select {
			case sig := <-sigc:
				glog.V(1).Infof("Caught signal %d", sig)
				cmd.Process.Signal(sig)
				cmd.Wait()
				exitCode := 0
				if cmd.ProcessState != nil {
					exitCode, _ = cmd.ProcessState.Sys().(int)
				}
				procexit <- exitCode
			case cmderr := <-serviceExit:
				if cmderr != nil {
					client, err := serviced.NewLBClient(proxyOptions.ServicedEndpoint)
					message := fmt.Sprintf("Service returned a non-zero exit code: %v. Command: \"%v\" Message: %v", config.ServiceId, config.Command, err)
					if err == nil {
						defer client.Close()
						glog.Errorf(message)

						// send the log message to the master
						client.SendLogMessage(serviced.ServiceLogInfo{config.ServiceId, message}, nil)
					} else {
						glog.Errorf("Failed to create a client to endpoint %s: %s", proxyOptions.ServicedEndpoint, err)
					}

					glog.Infof("%s", err)
					glog.Flush()
					if exiterr, ok := cmderr.(*exec.ExitError); ok && !proxyOptions.Autorestart {
						if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
							procexit <- status.ExitStatus()
						}
					}
				}
			}

			if !proxyOptions.Autorestart {
				break
			}
			glog.V(0).Info("service exited, sleeping...")
			time.Sleep(time.Minute)
		}
	}(config.Command)

	if proxyOptions.Logstash {
		go func() {
			// make sure we pick up any logfile that was modified within the
			// last three years
			// TODO: Either expose the 3 years a configurable or get rid of it
			logstashCmd := serviced.LOGSTASH_CONTAINER_DIRECTORY + "/logstash-forwarder"
			args := []string{"-old-files-hours=26280", "-config", serviced.LOGSTASH_CONTAINER_CONFIG}
			glog.V(0).Info("About to execute: %s %s", logstashCmd, args)

			myCmd := exec.Command(logstashCmd, args...)
			myCmd.Stdout = os.Stdout
			myCmd.Stderr = os.Stderr
			myErr := myCmd.Run()
			if myErr != nil {
				glog.Errorf("Problem running service: %v", myErr)
				glog.Flush()
			}
		}()
	}

	//monitor application endpoints to mux ports
	go func() {
		for {
			func() {
				client, err := serviced.NewLBClient(proxyOptions.ServicedEndpoint)
				if err != nil {
					glog.Errorf("Could not create a client to endpoint %s: %s", proxyOptions.ServicedEndpoint, err)
					return
				}
				defer client.Close()

				var endpoints map[string][]*dao.ApplicationEndpoint
				err = client.GetServiceEndpoints(config.ServiceId, &endpoints)
				if err != nil {
					glog.Errorf("Error getting application endpoints for service %s: %s", config.ServiceId, err)
					return
				}

				for key, endpointList := range endpoints {
					if len(endpointList) <= 0 {
						if proxy, ok := proxies[key]; ok {
							emptyAddressList := make([]string, 0)
							proxy.SetNewAddresses(emptyAddressList)
						}
						continue
					}

					addresses := make([]string, len(endpointList))
					for i, endpoint := range endpointList {
						addresses[i] = fmt.Sprintf("%s:%d", endpoint.HostIp, endpoint.HostPort)
					}
					sort.Strings(addresses)

					var (
						proxy *serviced.Proxy
						ok    bool
					)

					if proxy, ok = proxies[key]; !ok {
						glog.Infof("Attempting port map for: %s -> %+v", key, *endpointList[0])

						// setup a new proxy
						listener, err := net.Listen("tcp4", fmt.Sprintf(":%d", endpointList[0].ContainerPort))
						if err != nil {
							glog.Errorf("Could not bind to port: %s", err)
							continue
						}
						proxy, err = serviced.NewProxy(
							fmt.Sprintf("%v", endpointList[0]),
							uint16(config.TCPMux.Port),
							config.TCPMux.UseTLS,
							listener)
						if err != nil {
							glog.Errorf("Could not build proxy %s", err)
							continue
						}

						glog.Infof("Success binding port: %s -> %+v", key, proxy)
						proxies[key] = proxy

						if ep := endpointList[0]; ep.VirtualAddress != "" {
							p := strconv.FormatUint(uint64(ep.ContainerPort), 10)
							err := vifs.RegisterVirtualAddress(ep.VirtualAddress, p, ep.Protocol)
							if err != nil {
								glog.Errorf("Error creating virtual address: %+v", err)
							}
						}
					}
					proxy.SetNewAddresses(addresses)
				}
			}()

			time.Sleep(time.Second * 10)
		}
	}()

	//setup container metric forwarder
	go func() {
		//loop until successfully identifying this container's tenant id
		var tenantID string
		for {
			client, err := serviced.NewLBClient(proxyOptions.ServicedEndpoint)
			if err == nil {
				defer client.Close()
				if err = client.GetTenantId(config.ServiceId, &tenantID); err != nil {
					glog.Errorf("Failed to get tenant id: %s", err)
				} else {
					//success
					break
				}
			} else {
				glog.Errorf("Failed to create a client to endpoint %s: %s", proxyOptions.ServicedEndpoint, err)
			}
		}

		//build metric redirect url -- assumes 8444 is port mapped
		metricRedirect := "http://localhost:8444/api/metrics/store"
		metricRedirect += "?controlplane_tenant_id=" + tenantID
		metricRedirect += "&controlplane_service_id=" + config.ServiceId

		//build and serve the container metric forwarder
		forwarder, _ := serviced.NewMetricForwarder(":22350", metricRedirect)
		forwarder.Serve()
	}()

	exitcode := <-procexit // Wait for proc goroutine to exit

	glog.Flush()
	os.Exit(exitcode)
	return nil
}

var (
	proxies map[string]*serviced.Proxy
	vifs    *VIFRegistry
	nextip  int
)

func init() {
	proxies = make(map[string]*serviced.Proxy)
	vifs = NewVIFRegistry()
	nextip = 1
}
<commit_msg>Remove unused func<commit_after>package api

import (
	"github.com/zenoss/glog"
	"github.com/zenoss/serviced"
	"github.com/zenoss/serviced/dao"
	"github.com/zenoss/serviced/shell"

	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var proxyOptions ProxyOptions

// ProxyOptions are options to be run when starting a new proxy server
type ProxyOptions struct {
	MuxPort          int
	Mux              bool
	TLS              bool
	KeyPEMFile       string
	CertPEMFile      string
	ServicedEndpoint string
	Autorestart      bool
	Logstash         bool
}

// LoadProxyOptions loads the proxy option information
func LoadProxyOptions(ops ProxyOptions) {
	proxyOptions = ops
}

// ProxyConfig is the config object for starting a proxy server
type ProxyConfig struct {
	ServiceID string
	Command   []string
}

// Start a service proxy
func (a *api) StartProxy(cfg ProxyConfig) error {
	config := serviced.MuxConfig{}
	config.TCPMux.Port = proxyOptions.MuxPort
	config.TCPMux.Enabled = proxyOptions.Mux
	config.TCPMux.UseTLS = proxyOptions.TLS
	config.ServiceId = cfg.ServiceID
	config.Command = strings.Join(cfg.Command, " ")

	if config.TCPMux.Enabled {
		go config.TCPMux.ListenAndMux()
	}

	sio := shell.NewProcessForwarderServer(proxyOptions.ServicedEndpoint)
	sio.Handle("/", http.FileServer(http.Dir("/serviced/www/")))
	go http.ListenAndServe(":50000", sio)

	procexit := make(chan int)
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	// continually execute subprocess
	go func(cmdString string) {
		defer func() { procexit <- 0 }()
		for {
			glog.V(0).Info("About to execute: ", cmdString)
			cmd := exec.Command("bash", "-c", cmdString)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			cmd.Stdin = os.Stdin
			serviceExit := make(chan error)
			go func() {
				serviceExit <- cmd.Run()
			}()
			select {
			case sig := <-sigc:
				glog.V(1).Infof("Caught signal %d", sig)
				cmd.Process.Signal(sig)
				cmd.Wait()
				exitCode := 0
				if cmd.ProcessState != nil {
					exitCode, _ = cmd.ProcessState.Sys().(int)
				}
				procexit <- exitCode
			case cmderr := <-serviceExit:
				if cmderr != nil {
					client, err := serviced.NewLBClient(proxyOptions.ServicedEndpoint)
					message := fmt.Sprintf("Service returned a non-zero exit code: %v. Command: \"%v\" Message: %v", config.ServiceId, config.Command, err)
					if err == nil {
						defer client.Close()
						glog.Errorf(message)

						// send the log message to the master
						client.SendLogMessage(serviced.ServiceLogInfo{config.ServiceId, message}, nil)
					} else {
						glog.Errorf("Failed to create a client to endpoint %s: %s", proxyOptions.ServicedEndpoint, err)
					}

					glog.Infof("%s", err)
					glog.Flush()
					if exiterr, ok := cmderr.(*exec.ExitError); ok && !proxyOptions.Autorestart {
						if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
							procexit <- status.ExitStatus()
						}
					}
				}
			}

			if !proxyOptions.Autorestart {
				break
			}
			glog.V(0).Info("service exited, sleeping...")
			time.Sleep(time.Minute)
		}
	}(config.Command)

	if proxyOptions.Logstash {
		go func() {
			// make sure we pick up any logfile that was modified within the
			// last three years
			// TODO: Either expose the 3 years a configurable or get rid of it
			logstashCmd := serviced.LOGSTASH_CONTAINER_DIRECTORY + "/logstash-forwarder"
			args := []string{"-old-files-hours=26280", "-config", serviced.LOGSTASH_CONTAINER_CONFIG}
			glog.V(0).Info("About to execute: %s %s", logstashCmd, args)

			myCmd := exec.Command(logstashCmd, args...)
			myCmd.Stdout = os.Stdout
			myCmd.Stderr = os.Stderr
			myErr := myCmd.Run()
			if myErr != nil {
				glog.Errorf("Problem running service: %v", myErr)
				glog.Flush()
			}
		}()
	}

	//monitor application endpoints to mux ports
	go func() {
		for {
			func() {
				client, err := serviced.NewLBClient(proxyOptions.ServicedEndpoint)
				if err != nil {
					glog.Errorf("Could not create a client to endpoint %s: %s", proxyOptions.ServicedEndpoint, err)
					return
				}
				defer client.Close()

				var endpoints map[string][]*dao.ApplicationEndpoint
				err = client.GetServiceEndpoints(config.ServiceId, &endpoints)
				if err != nil {
					glog.Errorf("Error getting application endpoints for service %s: %s", config.ServiceId, err)
					return
				}

				for key, endpointList := range endpoints {
					if len(endpointList) <= 0 {
						if proxy, ok := proxies[key]; ok {
							emptyAddressList := make([]string, 0)
							proxy.SetNewAddresses(emptyAddressList)
						}
						continue
					}

					addresses := make([]string, len(endpointList))
					for i, endpoint := range endpointList {
						addresses[i] = fmt.Sprintf("%s:%d", endpoint.HostIp, endpoint.HostPort)
					}
					sort.Strings(addresses)

					var (
						proxy *serviced.Proxy
						ok    bool
					)

					if proxy, ok = proxies[key]; !ok {
						glog.Infof("Attempting port map for: %s -> %+v", key, *endpointList[0])

						// setup a new proxy
						listener, err := net.Listen("tcp4", fmt.Sprintf(":%d", endpointList[0].ContainerPort))
						if err != nil {
							glog.Errorf("Could not bind to port: %s", err)
							continue
						}
						proxy, err = serviced.NewProxy(
							fmt.Sprintf("%v", endpointList[0]),
							uint16(config.TCPMux.Port),
							config.TCPMux.UseTLS,
							listener)
						if err != nil {
							glog.Errorf("Could not build proxy %s", err)
							continue
						}

						glog.Infof("Success binding port: %s -> %+v", key, proxy)
						proxies[key] = proxy

						if ep := endpointList[0]; ep.VirtualAddress != "" {
							p := strconv.FormatUint(uint64(ep.ContainerPort), 10)
							err := vifs.RegisterVirtualAddress(ep.VirtualAddress, p, ep.Protocol)
							if err != nil {
								glog.Errorf("Error creating virtual address: %+v", err)
							}
						}
					}
					proxy.SetNewAddresses(addresses)
				}
			}()

			time.Sleep(time.Second * 10)
		}
	}()

	//setup container metric forwarder
	go func() {
		//loop until successfully identifying this container's tenant id
		var tenantID string
		for {
			client, err := serviced.NewLBClient(proxyOptions.ServicedEndpoint)
			if err == nil {
				defer client.Close()
				if err = client.GetTenantId(config.ServiceId, &tenantID); err != nil {
					glog.Errorf("Failed to get tenant id: %s", err)
				} else {
					//success
					break
				}
			} else {
				glog.Errorf("Failed to create a client to endpoint %s: %s", proxyOptions.ServicedEndpoint, err)
			}
		}

		//build metric redirect url -- assumes 8444 is port mapped
		metricRedirect := "http://localhost:8444/api/metrics/store"
		metricRedirect += "?controlplane_tenant_id=" + tenantID
		metricRedirect += "&controlplane_service_id=" + config.ServiceId

		//build and serve the container metric forwarder
		forwarder, _ := serviced.NewMetricForwarder(":22350", metricRedirect)
		forwarder.Serve()
	}()

	exitcode := <-procexit // Wait for proc goroutine to exit

	glog.Flush()
	os.Exit(exitcode)
	return nil
}

var (
	proxies map[string]*serviced.Proxy
	vifs    *VIFRegistry
	nextip  int
)

func init() {
	proxies = make(map[string]*serviced.Proxy)
	vifs = NewVIFRegistry()
	nextip = 1
}
<|endoftext|>


<commit_before>// Copyright 2014 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

package state

import (
	"bytes"
	"fmt"
	"io"
	"math/big"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/trie"
)

var emptyCodeHash = crypto.Keccak256(nil)

type Code []byte

func (self Code) String() string {
	return string(self) //strings.Join(Disassemble(self), " ")
}

type Storage map[common.Hash]common.Hash

func (self Storage) String() (str string) {
	for key, value := range self {
		str += fmt.Sprintf("%X : %X\n", key, value)
	}

	return
}

func (self Storage) Copy() Storage {
	cpy := make(Storage)
	for key, value := range self {
		cpy[key] = value
	}

	return cpy
}

// StateObject represents an Ethereum account which is being modified.
//
// The usage pattern is as follows:
// First you need to obtain a state object.
// Account values can be accessed and modified through the object.
// Finally, call CommitTrie to write the modified storage trie into a database.
type StateObject struct {
	address common.Address // Ethereum address of this account
	data    Account
	db      *StateDB

	// DB error.
	// State objects are used by the consensus core and VM which are
	// unable to deal with database-level errors. Any error that occurs
	// during a database read is memoized here and will eventually be returned
	// by StateDB.Commit.
	dbErr error

	// Write caches.
	trie *trie.SecureTrie // storage trie, which becomes non-nil on first access
	code Code             // contract bytecode, which gets set when code is loaded

	cachedStorage Storage // Storage entry cache to avoid duplicate reads
	dirtyStorage  Storage // Storage entries that need to be flushed to disk

	// Cache flags.
	// When an object is marked suicided it will be delete from the trie
	// during the "update" phase of the state transition.
	dirtyCode bool // true if the code was updated
	suicided  bool
	deleted   bool
	onDirty   func(addr common.Address) // Callback method to mark a state object newly dirty
}

// Account is the Ethereum consensus representation of accounts.
// These objects are stored in the main account trie.
type Account struct {
	Nonce    uint64
	Balance  *big.Int
	Root     common.Hash // merkle root of the storage trie
	CodeHash []byte
}

// newObject creates a state object.
func newObject(db *StateDB, address common.Address, data Account, onDirty func(addr common.Address)) *StateObject {
	if data.Balance == nil {
		data.Balance = new(big.Int)
	}
	if data.CodeHash == nil {
		data.CodeHash = emptyCodeHash
	}
	return &StateObject{db: db, address: address, data: data, cachedStorage: make(Storage), dirtyStorage: make(Storage), onDirty: onDirty}
}

// EncodeRLP implements rlp.Encoder.
func (c *StateObject) EncodeRLP(w io.Writer) error {
	return rlp.Encode(w, c.data)
}

// setError remembers the first non-nil error it is called with.
func (self *StateObject) setError(err error) {
	if self.dbErr == nil {
		self.dbErr = err
	}
}

func (self *StateObject) markSuicided() {
	self.suicided = true
	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
	if glog.V(logger.Core) {
		glog.Infof("%x: #%d %v X\n", self.Address(), self.Nonce(), self.Balance())
	}
}

func (c *StateObject) getTrie(db trie.Database) *trie.SecureTrie {
	if c.trie == nil {
		var err error
		c.trie, err = trie.NewSecure(c.data.Root, db)
		if err != nil {
			c.trie, _ = trie.NewSecure(common.Hash{}, db)
			c.setError(fmt.Errorf("can't create storage trie: %v", err))
		}
	}
	return c.trie
}

// GetState returns a value in account storage.
func (self *StateObject) GetState(db trie.Database, key common.Hash) common.Hash {
	value, exists := self.cachedStorage[key]
	if exists {
		return value
	}
	// Load from DB in case it is missing.
	tr := self.getTrie(db)
	var ret []byte
	rlp.DecodeBytes(tr.Get(key[:]), &ret)
	value = common.BytesToHash(ret)
	if (value != common.Hash{}) {
		self.cachedStorage[key] = value
	}
	return value
}

// SetState updates a value in account storage.
func (self *StateObject) SetState(db trie.Database, key, value common.Hash) {
	self.db.journal = append(self.db.journal, storageChange{
		account:  &self.address,
		key:      key,
		prevalue: self.GetState(db, key),
	})
	self.setState(key, value)
}

func (self *StateObject) setState(key, value common.Hash) {
	self.cachedStorage[key] = value
	self.dirtyStorage[key] = value

	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
}

// updateTrie writes cached storage modifications into the object's storage trie.
func (self *StateObject) updateTrie(db trie.Database) {
	tr := self.getTrie(db)
	for key, value := range self.dirtyStorage {
		delete(self.dirtyStorage, key)
		if (value == common.Hash{}) {
			tr.Delete(key[:])
			continue
		}
		// Encoding []byte cannot fail, ok to ignore the error.
		v, _ := rlp.EncodeToBytes(bytes.TrimLeft(value[:], "\x00"))
		tr.Update(key[:], v)
	}
}

// UpdateRoot sets the trie root to the current root hash of
func (self *StateObject) updateRoot(db trie.Database) {
	self.updateTrie(db)
	self.data.Root = self.trie.Hash()
}

// CommitTrie the storage trie of the object to dwb.
// This updates the trie root.
func (self *StateObject) CommitTrie(db trie.Database, dbw trie.DatabaseWriter) error {
	self.updateTrie(db)
	if self.dbErr != nil {
		fmt.Println("dbErr:", self.dbErr)
		return self.dbErr
	}
	root, err := self.trie.CommitTo(dbw)
	if err == nil {
		self.data.Root = root
	}
	return err
}

func (c *StateObject) AddBalance(amount *big.Int) {
	if amount.Cmp(common.Big0) == 0 {
		return
	}
	c.SetBalance(new(big.Int).Add(c.Balance(), amount))

	if glog.V(logger.Core) {
		glog.Infof("%x: #%d %v (+ %v)\n", c.Address(), c.Nonce(), c.Balance(), amount)
	}
}

func (c *StateObject) SubBalance(amount *big.Int) {
	if amount.Cmp(common.Big0) == 0 {
		return
	}
	c.SetBalance(new(big.Int).Sub(c.Balance(), amount))

	if glog.V(logger.Core) {
		glog.Infof("%x: #%d %v (- %v)\n", c.Address(), c.Nonce(), c.Balance(), amount)
	}
}

func (self *StateObject) SetBalance(amount *big.Int) {
	self.db.journal = append(self.db.journal, balanceChange{
		account: &self.address,
		prev:    new(big.Int).Set(self.data.Balance),
	})
	self.setBalance(amount)
}

func (self *StateObject) setBalance(amount *big.Int) {
	self.data.Balance = amount
	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
}

// Return the gas back to the origin. Used by the Virtual machine or Closures
func (c *StateObject) ReturnGas(gas, price *big.Int) {}

func (self *StateObject) deepCopy(db *StateDB, onDirty func(addr common.Address)) *StateObject {
	stateObject := newObject(db, self.address, self.data, onDirty)
	stateObject.trie = self.trie
	stateObject.code = self.code
	stateObject.dirtyStorage = self.dirtyStorage.Copy()
	stateObject.cachedStorage = self.dirtyStorage.Copy()
	stateObject.suicided = self.suicided
	stateObject.dirtyCode = self.dirtyCode
	stateObject.deleted = self.deleted
	return stateObject
}

//
// Attribute accessors
//

// Returns the address of the contract/account
func (c *StateObject) Address() common.Address {
	return c.address
}

// Code returns the contract code associated with this object, if any.
func (self *StateObject) Code(db trie.Database) []byte {
	if self.code != nil {
		return self.code
	}
	if bytes.Equal(self.CodeHash(), emptyCodeHash) {
		return nil
	}
	code, err := db.Get(self.CodeHash())
	if err != nil {
		self.setError(fmt.Errorf("can't load code hash %x: %v", self.CodeHash(), err))
	}
	self.code = code
	return code
}

func (self *StateObject) SetCode(codeHash common.Hash, code []byte) {
	prevcode := self.Code(self.db.db)
	self.db.journal = append(self.db.journal, codeChange{
		account:  &self.address,
		prevhash: self.CodeHash(),
		prevcode: prevcode,
	})
	self.setCode(codeHash, code)
}

func (self *StateObject) setCode(codeHash common.Hash, code []byte) {
	self.code = code
	self.data.CodeHash = codeHash[:]
	self.dirtyCode = true
	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
}

func (self *StateObject) SetNonce(nonce uint64) {
	self.db.journal = append(self.db.journal, nonceChange{
		account: &self.address,
		prev:    self.data.Nonce,
	})
	self.setNonce(nonce)
}

func (self *StateObject) setNonce(nonce uint64) {
	self.data.Nonce = nonce
	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
}

func (self *StateObject) CodeHash() []byte {
	return self.data.CodeHash
}

func (self *StateObject) Balance() *big.Int {
	return self.data.Balance
}

func (self *StateObject) Nonce() uint64 {
	return self.data.Nonce
}

// Never called, but must be present to allow StateObject to be used
// as a vm.Account interface that also satisfies the vm.ContractRef
// interface. Interfaces are awesome.
func (self *StateObject) Value() *big.Int {
	panic("Value on StateObject should never be called")
}

func (self *StateObject) ForEachStorage(cb func(key, value common.Hash) bool) {
	// When iterating over the storage check the cache first
	for h, value := range self.cachedStorage {
		cb(h, value)
	}

	it := self.getTrie(self.db.db).Iterator()
	for it.Next() {
		// ignore cached values
		key := common.BytesToHash(self.trie.GetKey(it.Key))
		if _, ok := self.cachedStorage[key]; !ok {
			cb(key, common.BytesToHash(it.Value))
		}
	}
}
<commit_msg>core/state: optimize GetState<commit_after>// Copyright 2014 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

package state

import (
	"bytes"
	"fmt"
	"io"
	"math/big"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/trie"
)

var emptyCodeHash = crypto.Keccak256(nil)

type Code []byte

func (self Code) String() string {
	return string(self) //strings.Join(Disassemble(self), " ")
}

type Storage map[common.Hash]common.Hash

func (self Storage) String() (str string) {
	for key, value := range self {
		str += fmt.Sprintf("%X : %X\n", key, value)
	}

	return
}

func (self Storage) Copy() Storage {
	cpy := make(Storage)
	for key, value := range self {
		cpy[key] = value
	}

	return cpy
}

// StateObject represents an Ethereum account which is being modified.
//
// The usage pattern is as follows:
// First you need to obtain a state object.
// Account values can be accessed and modified through the object.
// Finally, call CommitTrie to write the modified storage trie into a database.
type StateObject struct {
	address common.Address // Ethereum address of this account
	data    Account
	db      *StateDB

	// DB error.
	// State objects are used by the consensus core and VM which are
	// unable to deal with database-level errors. Any error that occurs
	// during a database read is memoized here and will eventually be returned
	// by StateDB.Commit.
	dbErr error

	// Write caches.
	trie *trie.SecureTrie // storage trie, which becomes non-nil on first access
	code Code             // contract bytecode, which gets set when code is loaded

	cachedStorage Storage // Storage entry cache to avoid duplicate reads
	dirtyStorage  Storage // Storage entries that need to be flushed to disk

	// Cache flags.
	// When an object is marked suicided it will be delete from the trie
	// during the "update" phase of the state transition.
	dirtyCode bool // true if the code was updated
	suicided  bool
	deleted   bool
	onDirty   func(addr common.Address) // Callback method to mark a state object newly dirty
}

// Account is the Ethereum consensus representation of accounts.
// These objects are stored in the main account trie.
type Account struct {
	Nonce    uint64
	Balance  *big.Int
	Root     common.Hash // merkle root of the storage trie
	CodeHash []byte
}

// newObject creates a state object.
func newObject(db *StateDB, address common.Address, data Account, onDirty func(addr common.Address)) *StateObject {
	if data.Balance == nil {
		data.Balance = new(big.Int)
	}
	if data.CodeHash == nil {
		data.CodeHash = emptyCodeHash
	}
	return &StateObject{db: db, address: address, data: data, cachedStorage: make(Storage), dirtyStorage: make(Storage), onDirty: onDirty}
}

// EncodeRLP implements rlp.Encoder.
func (c *StateObject) EncodeRLP(w io.Writer) error {
	return rlp.Encode(w, c.data)
}

// setError remembers the first non-nil error it is called with.
func (self *StateObject) setError(err error) {
	if self.dbErr == nil {
		self.dbErr = err
	}
}

func (self *StateObject) markSuicided() {
	self.suicided = true
	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
	if glog.V(logger.Core) {
		glog.Infof("%x: #%d %v X\n", self.Address(), self.Nonce(), self.Balance())
	}
}

func (c *StateObject) getTrie(db trie.Database) *trie.SecureTrie {
	if c.trie == nil {
		var err error
		c.trie, err = trie.NewSecure(c.data.Root, db)
		if err != nil {
			c.trie, _ = trie.NewSecure(common.Hash{}, db)
			c.setError(fmt.Errorf("can't create storage trie: %v", err))
		}
	}
	return c.trie
}

// GetState returns a value in account storage.
func (self *StateObject) GetState(db trie.Database, key common.Hash) common.Hash {
	value, exists := self.cachedStorage[key]
	if exists {
		return value
	}
	// Load from DB in case it is missing.
	if enc := self.getTrie(db).Get(key[:]); len(enc) > 0 {
		_, content, _, err := rlp.Split(enc)
		if err != nil {
			self.setError(err)
		}
		value.SetBytes(content)
	}
	if (value != common.Hash{}) {
		self.cachedStorage[key] = value
	}
	return value
}

// SetState updates a value in account storage.
func (self *StateObject) SetState(db trie.Database, key, value common.Hash) {
	self.db.journal = append(self.db.journal, storageChange{
		account:  &self.address,
		key:      key,
		prevalue: self.GetState(db, key),
	})
	self.setState(key, value)
}

func (self *StateObject) setState(key, value common.Hash) {
	self.cachedStorage[key] = value
	self.dirtyStorage[key] = value

	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
}

// updateTrie writes cached storage modifications into the object's storage trie.
func (self *StateObject) updateTrie(db trie.Database) {
	tr := self.getTrie(db)
	for key, value := range self.dirtyStorage {
		delete(self.dirtyStorage, key)
		if (value == common.Hash{}) {
			tr.Delete(key[:])
			continue
		}
		// Encoding []byte cannot fail, ok to ignore the error.
		v, _ := rlp.EncodeToBytes(bytes.TrimLeft(value[:], "\x00"))
		tr.Update(key[:], v)
	}
}

// UpdateRoot sets the trie root to the current root hash of
func (self *StateObject) updateRoot(db trie.Database) {
	self.updateTrie(db)
	self.data.Root = self.trie.Hash()
}

// CommitTrie the storage trie of the object to dwb.
// This updates the trie root.
func (self *StateObject) CommitTrie(db trie.Database, dbw trie.DatabaseWriter) error {
	self.updateTrie(db)
	if self.dbErr != nil {
		return self.dbErr
	}
	root, err := self.trie.CommitTo(dbw)
	if err == nil {
		self.data.Root = root
	}
	return err
}

func (c *StateObject) AddBalance(amount *big.Int) {
	if amount.Cmp(common.Big0) == 0 {
		return
	}
	c.SetBalance(new(big.Int).Add(c.Balance(), amount))

	if glog.V(logger.Core) {
		glog.Infof("%x: #%d %v (+ %v)\n", c.Address(), c.Nonce(), c.Balance(), amount)
	}
}

func (c *StateObject) SubBalance(amount *big.Int) {
	if amount.Cmp(common.Big0) == 0 {
		return
	}
	c.SetBalance(new(big.Int).Sub(c.Balance(), amount))

	if glog.V(logger.Core) {
		glog.Infof("%x: #%d %v (- %v)\n", c.Address(), c.Nonce(), c.Balance(), amount)
	}
}

func (self *StateObject) SetBalance(amount *big.Int) {
	self.db.journal = append(self.db.journal, balanceChange{
		account: &self.address,
		prev:    new(big.Int).Set(self.data.Balance),
	})
	self.setBalance(amount)
}

func (self *StateObject) setBalance(amount *big.Int) {
	self.data.Balance = amount
	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
}

// Return the gas back to the origin. Used by the Virtual machine or Closures
func (c *StateObject) ReturnGas(gas, price *big.Int) {}

func (self *StateObject) deepCopy(db *StateDB, onDirty func(addr common.Address)) *StateObject {
	stateObject := newObject(db, self.address, self.data, onDirty)
	stateObject.trie = self.trie
	stateObject.code = self.code
	stateObject.dirtyStorage = self.dirtyStorage.Copy()
	stateObject.cachedStorage = self.dirtyStorage.Copy()
	stateObject.suicided = self.suicided
	stateObject.dirtyCode = self.dirtyCode
	stateObject.deleted = self.deleted
	return stateObject
}

//
// Attribute accessors
//

// Returns the address of the contract/account
func (c *StateObject) Address() common.Address {
	return c.address
}

// Code returns the contract code associated with this object, if any.
func (self *StateObject) Code(db trie.Database) []byte {
	if self.code != nil {
		return self.code
	}
	if bytes.Equal(self.CodeHash(), emptyCodeHash) {
		return nil
	}
	code, err := db.Get(self.CodeHash())
	if err != nil {
		self.setError(fmt.Errorf("can't load code hash %x: %v", self.CodeHash(), err))
	}
	self.code = code
	return code
}

func (self *StateObject) SetCode(codeHash common.Hash, code []byte) {
	prevcode := self.Code(self.db.db)
	self.db.journal = append(self.db.journal, codeChange{
		account:  &self.address,
		prevhash: self.CodeHash(),
		prevcode: prevcode,
	})
	self.setCode(codeHash, code)
}

func (self *StateObject) setCode(codeHash common.Hash, code []byte) {
	self.code = code
	self.data.CodeHash = codeHash[:]
	self.dirtyCode = true
	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
}

func (self *StateObject) SetNonce(nonce uint64) {
	self.db.journal = append(self.db.journal, nonceChange{
		account: &self.address,
		prev:    self.data.Nonce,
	})
	self.setNonce(nonce)
}

func (self *StateObject) setNonce(nonce uint64) {
	self.data.Nonce = nonce
	if self.onDirty != nil {
		self.onDirty(self.Address())
		self.onDirty = nil
	}
}

func (self *StateObject) CodeHash() []byte {
	return self.data.CodeHash
}

func (self *StateObject) Balance() *big.Int {
	return self.data.Balance
}

func (self *StateObject) Nonce() uint64 {
	return self.data.Nonce
}

// Never called, but must be present to allow StateObject to be used
// as a vm.Account interface that also satisfies the vm.ContractRef
// interface. Interfaces are awesome.
func (self *StateObject) Value() *big.Int {
	panic("Value on StateObject should never be called")
}

func (self *StateObject) ForEachStorage(cb func(key, value common.Hash) bool) {
	// When iterating over the storage check the cache first
	for h, value := range self.cachedStorage {
		cb(h, value)
	}

	it := self.getTrie(self.db.db).Iterator()
	for it.Next() {
		// ignore cached values
		key := common.BytesToHash(self.trie.GetKey(it.Key))
		if _, ok := self.cachedStorage[key]; !ok {
			cb(key, common.BytesToHash(it.Value))
		}
	}
}
<|endoftext|>


<commit_before>package cluster

import (
	"context"
	"sync"
	"time"

	log "gopkg.in/inconshreveable/log15.v2"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/endpoints"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/logger"
)

// eventHubMinHosts is the minimum number of members that must have the event-hub role to trigger switching into
// event-hub mode (where cluster members will only connect to event-hub members rather than all members when
// operating in the normal full-mesh mode).
const eventHubMinHosts = 2

var listeners = map[string]*lxd.EventListener{}
var listenersNotify = map[chan struct{}][]string{}
var listenersLock sync.Mutex
var listenersUpdateLock sync.Mutex

// EventListenerWait waits for there to be listener connected to the specified address, or one of the event hubs
// if operating in event hub mode.
func EventListenerWait(ctx context.Context, address string) error {
	// Check if there is already a listener.
	listenersLock.Lock()
	listener, found := listeners[address]
	if found && listener.IsActive() {
		listenersLock.Unlock()
		return nil
	}

	listenAddresses := []string{address}

	// If not setup a notification for when the desired address or any of the event hubs connect.
	connected := make(chan struct{})
	listenersNotify[connected] = listenAddresses
	listenersLock.Unlock()

	defer func() {
		listenersLock.Lock()
		delete(listenersNotify, connected)
		listenersLock.Unlock()
	}()

	// Wait for the connected channel to be closed (indicating a new listener has been connected), and return.
	select {
	case <-connected:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// EventsUpdateListeners refreshes the cluster event listener connections.
func EventsUpdateListeners(endpoints *endpoints.Endpoints, cluster *db.Cluster, serverCert func() *shared.CertInfo, members map[int64]APIHeartbeatMember, f func(int64, api.Event)) {
	listenersUpdateLock.Lock()
	defer listenersUpdateLock.Unlock()

	// If no heartbeat members provided, populate from global database.
	if members == nil {
		var dbMembers []db.NodeInfo
		var offlineThreshold time.Duration

		err := cluster.Transaction(func(tx *db.ClusterTx) error {
			var err error

			dbMembers, err = tx.GetNodes()
			if err != nil {
				return err
			}

			offlineThreshold, err = tx.GetNodeOfflineThreshold()
			if err != nil {
				return err
			}

			return nil
		})
		if err != nil {
			logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
			return
		}

		members = make(map[int64]APIHeartbeatMember, len(dbMembers))
		for _, dbMember := range dbMembers {
			members[dbMember.ID] = APIHeartbeatMember{
				ID:            dbMember.ID,
				Name:          dbMember.Name,
				Address:       dbMember.Address,
				LastHeartbeat: dbMember.Heartbeat,
				Online:        !dbMember.IsOffline(offlineThreshold),
				Roles:         dbMember.Roles,
			}
		}
	}

	networkAddress := endpoints.NetworkAddress()

	keepListeners := make(map[string]struct{})
	wg := sync.WaitGroup{}
	for _, member := range members {
		// Don't bother trying to connect to ourselves or offline members.
		if member.Address == networkAddress || !member.Online {
			continue
		}

		listenersLock.Lock()
		listener, ok := listeners[member.Address]

		// If the member already has a listener associated to it, check that the listener is still active.
		// If it is, just move on to next member, but if not then we'll try to connect again.
		if ok {
			if listener.IsActive() {
				keepListeners[member.Address] = struct{}{} // Add to current listeners list.
				listenersLock.Unlock()
				continue
			}

			// Disconnect and delete listener, but don't delete any listenersNotify entry as there
			// might be something waiting for a future connection.
			listener.Disconnect()
			delete(listeners, member.Address)
			logger.Info("Removed inactive member event listener client", log.Ctx{"local": networkAddress, "remote": member.Address})
		}
		listenersLock.Unlock()

		keepListeners[member.Address] = struct{}{} // Add to current listeners list.

		// Connect to remote concurrently and add to active listeners if successful.
		wg.Add(1)
		go func(m APIHeartbeatMember) {
			defer wg.Done()
			listener, err := eventsConnect(m.Address, endpoints.NetworkCert(), serverCert())
			if err != nil {
				logger.Warn("Failed adding member event listener client", log.Ctx{"local": networkAddress, "remote": m.Address, "err": err})
				return
			}

			listener.AddHandler(nil, func(event api.Event) { f(m.ID, event) })

			listenersLock.Lock()
			listeners[m.Address] = listener

			// Indicate to any notifiers waiting for this member's address that it is connected.
			for connected, notifyAddresses := range listenersNotify {
				if shared.StringInSlice(m.Address, notifyAddresses) {
					close(connected)
					delete(listenersNotify, connected)
				}
			}

			logger.Info("Added member event listener client", log.Ctx{"local": networkAddress, "remote": m.Address})
			listenersLock.Unlock()
		}(member)
	}

	wg.Wait()

	// Disconnect and delete any out of date listeners and their notifiers.
	listenersLock.Lock()
	for address, listener := range listeners {
		if _, found := keepListeners[address]; !found {
			listener.Disconnect()
			delete(listeners, address)
			logger.Info("Removed old member event listener client", log.Ctx{"local": networkAddress, "remote": address})
		}
	}
	listenersLock.Unlock()
}

// Establish a client connection to get events from the given node.
func eventsConnect(address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo) (*lxd.EventListener, error) {
	client, err := Connect(address, networkCert, serverCert, nil, true)
	if err != nil {
		return nil, err
	}

	revert := revert.New()
	revert.Add(func() {
		client.Disconnect()
	})

	listener, err := client.GetEventsAllProjects()
	if err != nil {
		return nil, err
	}

	revert.Success()
	return listener, nil
}
<commit_msg>lxd/cluster/events: Adds EventMode type and constants<commit_after>package cluster

import (
	"context"
	"sync"
	"time"

	log "gopkg.in/inconshreveable/log15.v2"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/endpoints"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/logger"
)

// eventHubMinHosts is the minimum number of members that must have the event-hub role to trigger switching into
// event-hub mode (where cluster members will only connect to event-hub members rather than all members when
// operating in the normal full-mesh mode).
const eventHubMinHosts = 2

// EventMode indicates the event distribution mode.
type EventMode string

// EventModeFullMesh is when every cluster member connects to every other cluster member to pull events.
const EventModeFullMesh EventMode = "full-mesh"

// EventModeHubServer is when the cluster is operating in event-hub mode and this server is designated as a hub
// server, meaning that it will only connect to the other event-hub members and not other members.
const EventModeHubServer EventMode = "hub-server"

// EventModeHubClient is when the cluster is operating in event-hub mode and this member is designated as a hub
// client, meaning that it is expected to connect to the event-hub members.
const EventModeHubClient EventMode = "hub-client"

var listeners = map[string]*lxd.EventListener{}
var listenersNotify = map[chan struct{}][]string{}
var listenersLock sync.Mutex
var listenersUpdateLock sync.Mutex

// EventListenerWait waits for there to be listener connected to the specified address, or one of the event hubs
// if operating in event hub mode.
func EventListenerWait(ctx context.Context, address string) error {
	// Check if there is already a listener.
	listenersLock.Lock()
	listener, found := listeners[address]
	if found && listener.IsActive() {
		listenersLock.Unlock()
		return nil
	}

	listenAddresses := []string{address}

	// If not setup a notification for when the desired address or any of the event hubs connect.
	connected := make(chan struct{})
	listenersNotify[connected] = listenAddresses
	listenersLock.Unlock()

	defer func() {
		listenersLock.Lock()
		delete(listenersNotify, connected)
		listenersLock.Unlock()
	}()

	// Wait for the connected channel to be closed (indicating a new listener has been connected), and return.
	select {
	case <-connected:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// EventsUpdateListeners refreshes the cluster event listener connections.
func EventsUpdateListeners(endpoints *endpoints.Endpoints, cluster *db.Cluster, serverCert func() *shared.CertInfo, members map[int64]APIHeartbeatMember, f func(int64, api.Event)) {
	listenersUpdateLock.Lock()
	defer listenersUpdateLock.Unlock()

	// If no heartbeat members provided, populate from global database.
	if members == nil {
		var dbMembers []db.NodeInfo
		var offlineThreshold time.Duration

		err := cluster.Transaction(func(tx *db.ClusterTx) error {
			var err error

			dbMembers, err = tx.GetNodes()
			if err != nil {
				return err
			}

			offlineThreshold, err = tx.GetNodeOfflineThreshold()
			if err != nil {
				return err
			}

			return nil
		})
		if err != nil {
			logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
			return
		}

		members = make(map[int64]APIHeartbeatMember, len(dbMembers))
		for _, dbMember := range dbMembers {
			members[dbMember.ID] = APIHeartbeatMember{
				ID:            dbMember.ID,
				Name:          dbMember.Name,
				Address:       dbMember.Address,
				LastHeartbeat: dbMember.Heartbeat,
				Online:        !dbMember.IsOffline(offlineThreshold),
				Roles:         dbMember.Roles,
			}
		}
	}

	networkAddress := endpoints.NetworkAddress()

	keepListeners := make(map[string]struct{})
	wg := sync.WaitGroup{}
	for _, member := range members {
		// Don't bother trying to connect to ourselves or offline members.
		if member.Address == networkAddress || !member.Online {
			continue
		}

		listenersLock.Lock()
		listener, ok := listeners[member.Address]

		// If the member already has a listener associated to it, check that the listener is still active.
		// If it is, just move on to next member, but if not then we'll try to connect again.
		if ok {
			if listener.IsActive() {
				keepListeners[member.Address] = struct{}{} // Add to current listeners list.
				listenersLock.Unlock()
				continue
			}

			// Disconnect and delete listener, but don't delete any listenersNotify entry as there
			// might be something waiting for a future connection.
			listener.Disconnect()
			delete(listeners, member.Address)
			logger.Info("Removed inactive member event listener client", log.Ctx{"local": networkAddress, "remote": member.Address})
		}
		listenersLock.Unlock()

		keepListeners[member.Address] = struct{}{} // Add to current listeners list.

		// Connect to remote concurrently and add to active listeners if successful.
		wg.Add(1)
		go func(m APIHeartbeatMember) {
			defer wg.Done()
			listener, err := eventsConnect(m.Address, endpoints.NetworkCert(), serverCert())
			if err != nil {
				logger.Warn("Failed adding member event listener client", log.Ctx{"local": networkAddress, "remote": m.Address, "err": err})
				return
			}

			listener.AddHandler(nil, func(event api.Event) { f(m.ID, event) })

			listenersLock.Lock()
			listeners[m.Address] = listener

			// Indicate to any notifiers waiting for this member's address that it is connected.
			for connected, notifyAddresses := range listenersNotify {
				if shared.StringInSlice(m.Address, notifyAddresses) {
					close(connected)
					delete(listenersNotify, connected)
				}
			}

			logger.Info("Added member event listener client", log.Ctx{"local": networkAddress, "remote": m.Address})
			listenersLock.Unlock()
		}(member)
	}

	wg.Wait()

	// Disconnect and delete any out of date listeners and their notifiers.
	listenersLock.Lock()
	for address, listener := range listeners {
		if _, found := keepListeners[address]; !found {
			listener.Disconnect()
			delete(listeners, address)
			logger.Info("Removed old member event listener client", log.Ctx{"local": networkAddress, "remote": address})
		}
	}
	listenersLock.Unlock()
}

// Establish a client connection to get events from the given node.
func eventsConnect(address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo) (*lxd.EventListener, error) {
	client, err := Connect(address, networkCert, serverCert, nil, true)
	if err != nil {
		return nil, err
	}

	revert := revert.New()
	revert.Add(func() {
		client.Disconnect()
	})

	listener, err := client.GetEventsAllProjects()
	if err != nil {
		return nil, err
	}

	revert.Success()
	return listener, nil
}
<|endoftext|>


<commit_before>package project

import (
	"fmt"
	"strconv"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/units"
	"github.com/pkg/errors"
)

// CheckLimitsUponInstanceCreation returns an error if any project-specific
// limit is violated when creating a new instance.
func CheckLimitsUponInstanceCreation(tx *db.ClusterTx, projectName string, req api.InstancesPost) error {
	project, profiles, instances, err := fetchProject(tx, projectName)
	if err != nil {
		return err
	}

	err = checkInstanceCountLimit(project, len(instances), req.Type)
	if err != nil {
		return err
	}

	// Add the instance being created.
	instances = append(instances, db.Instance{
		Profiles: req.Profiles,
		Config:   req.Config,
	})

	err = checkAggregateInstanceLimits(tx, project, instances, profiles)
	if err != nil {
		return err
	}

	return nil
}

// Check that we have not reached the maximum number of instances for
// this type.
func checkInstanceCountLimit(project *api.Project, instanceCount int, instanceType api.InstanceType) error {
	var key string
	switch instanceType {
	case api.InstanceTypeContainer:
		key = "limits.containers"
	case api.InstanceTypeVM:
		key = "limits.virtual-machines"
	default:
		return fmt.Errorf("Unexpected instance type '%s'", instanceType)
	}
	value, ok := project.Config[key]
	if ok {
		limit, err := strconv.Atoi(value)
		if err != nil || limit < 0 {
			return fmt.Errorf("Unexpected '%s' value: '%s'", key, value)
		}
		if instanceCount >= limit {
			return fmt.Errorf(
				"Reached maximum number of instances of type %s in project %s",
				instanceType, project.Name)
		}
	}

	return nil
}

// Check that we would not violate the project limits if we were to commit the
// given instances and profiles.
func checkAggregateInstanceLimits(tx *db.ClusterTx, project *api.Project, instances []db.Instance, profiles []db.Profile) error {
	// List of config keys for which we need to check aggregate values
	// across all project instances.
	aggregateKeys := []string{}
	for key := range project.Config {
		if shared.StringInSlice(key, []string{"limits.memory"}) {
			aggregateKeys = append(aggregateKeys, key)
		}
	}
	if len(aggregateKeys) == 0 {
		return nil
	}

	instances = expandInstancesConfig(instances, profiles)

	totals, err := getTotalsAcrossInstances(instances, aggregateKeys)
	if err != nil {
		return err
	}

	for _, key := range aggregateKeys {
		parser := aggregateLimitConfigValueParsers[key]
		max, err := parser(project.Config[key])
		if err != nil {
			return err
		}
		if totals[key] > max {
			return fmt.Errorf(
				"Reached maximum aggregate value %s for %q in project %s",
				project.Config[key], key, project.Name)
		}
	}

	return nil
}

// CheckLimitsUponInstanceUpdate returns an error if any project-specific limit
// is violated when updating an existing instance.
func CheckLimitsUponInstanceUpdate(tx *db.ClusterTx, projectName, instanceName string, req api.InstancePut) error {
	project, profiles, instances, err := fetchProject(tx, projectName)
	if err != nil {
		return err
	}

	// Change the instance being updated.
	for i, instance := range instances {
		if instance.Name != instanceName {
			continue
		}
		instances[i].Profiles = req.Profiles
		instances[i].Config = req.Config
	}

	err = checkAggregateInstanceLimits(tx, project, instances, profiles)
	if err != nil {
		return err
	}

	return nil
}

// CheckLimitsUponProfileUpdate checks that project limits are not violated
// when changing a profile.
func CheckLimitsUponProfileUpdate(tx *db.ClusterTx, projectName, profileName string, req api.ProfilePut) error {
	project, profiles, instances, err := fetchProject(tx, projectName)
	if err != nil {
		return err
	}

	// Change the profile being updated.
	for i, profile := range profiles {
		if profile.Name != profileName {
			continue
		}
		profiles[i].Config = req.Config
	}

	err = checkAggregateInstanceLimits(tx, project, instances, profiles)
	if err != nil {
		return err
	}

	return nil
}

// ValidateLimitsUponProjectUpdate checks the new limits to be set on a project
// are valid.
func ValidateLimitsUponProjectUpdate(tx *db.ClusterTx, projectName string, config map[string]string, changed []string) error {
	_, profiles, instances, err := fetchProject(tx, projectName)
	if err != nil {
		return err
	}

	instances = expandInstancesConfig(instances, profiles)

	// List of keys that need to check aggregate values across all project
	// instances.
	aggregateKeys := []string{}

	for _, key := range changed {
		switch key {
		case "limits.containers":
			fallthrough
		case "limits.virtual-machines":
			err := validateInstanceCountLimit(instances, key, config[key], projectName)
			if err != nil {
				return err
			}
		case "limits.memory":
			aggregateKeys = append(aggregateKeys, key)

		}
	}

	if len(aggregateKeys) > 0 {
		totals, err := getTotalsAcrossInstances(instances, aggregateKeys)
		if err != nil {
			return err
		}
		for _, key := range aggregateKeys {
			err := validateAggregateLimit(totals, key, config[key])
			if err != nil {
				return err
			}
		}

	}

	return nil
}

// Check that limits.containers or limits.virtual-machines is equal or above
// the current count.
func validateInstanceCountLimit(instances []db.Instance, key, value, project string) error {
	instanceType := countConfigInstanceType[key]
	limit, err := strconv.Atoi(value)
	if err != nil {
		return err
	}
	dbType, err := instancetype.New(string(instanceType))
	if err != nil {
		return err
	}
	count := 0
	for _, instance := range instances {
		if instance.Type == dbType {
			count++
		}
	}
	if limit < count {
		return fmt.Errorf(
			"'%s' is too low: there currently are %d instances of type %s in project %s",
			key, count, instanceType, project)
	}
	return nil
}

var countConfigInstanceType = map[string]api.InstanceType{
	"limits.containers":       api.InstanceTypeContainer,
	"limits.virtual-machines": api.InstanceTypeVM,
}

// Validates an aggregate limit, checking that the new value is not below the
// current total amount.
func validateAggregateLimit(totals map[string]int64, key, value string) error {
	parser := aggregateLimitConfigValueParsers[key]
	limit, err := parser(value)
	if err != nil {
		errors.Wrapf(err, "Invalid value '%s' for limit %s", value, key)
	}

	total := totals[key]
	if limit < total {
		return fmt.Errorf("'%s' is too low: current total is %s", key, units.GetByteSizeString(total, 1))
	}

	return nil
}

// Fetch the given project from the database along with its profiles and instances.
func fetchProject(tx *db.ClusterTx, projectName string) (*api.Project, []db.Profile, []db.Instance, error) {
	project, err := tx.ProjectGet(projectName)
	if err != nil {
		return nil, nil, nil, errors.Wrap(err, "Fetch project database object")
	}

	profilesFilter := db.ProfileFilter{}

	// If the project has the profiles feature enabled, we use its own
	// profiles to expand the instances configs, otherwise we use the
	// profiles from the default project.
	if projectName == "default" || project.Config["features.profiles"] == "true" {
		profilesFilter.Project = projectName
	} else {
		profilesFilter.Project = "default"
	}

	profiles, err := tx.ProfileList(profilesFilter)
	if err != nil {
		return nil, nil, nil, errors.Wrap(err, "Fetch profiles from database")
	}

	instances, err := tx.InstanceList(db.InstanceFilter{Project: projectName})
	if err != nil {
		return nil, nil, nil, errors.Wrap(err, "Fetch project instances from database")
	}

	return project, profiles, instances, nil
}

// Expand the configuration of the given instances, taking the give project
// profiles into account.
func expandInstancesConfig(instances []db.Instance, profiles []db.Profile) []db.Instance {
	expandedInstances := make([]db.Instance, len(instances))

	// Index of all profiles by name.
	profilesByName := map[string]db.Profile{}
	for _, profile := range profiles {
		profilesByName[profile.Name] = profile
	}

	for i, instance := range instances {
		profiles := make([]api.Profile, len(instance.Profiles))

		for j, name := range instance.Profiles {
			profile := profilesByName[name]
			profiles[j] = *db.ProfileToAPI(&profile)
		}

		expandedInstances[i] = instance
		expandedInstances[i].Config = db.ProfilesExpandConfig(instance.Config, profiles)
	}

	return expandedInstances
}

// Sum of the effective instance-level value for the given limits across all
// project instances. If excludeInstance is not the empty string, exclude the
// instance with that name.
func getTotalsAcrossInstances(instances []db.Instance, keys []string) (map[string]int64, error) {
	totals := map[string]int64{}

	for _, key := range keys {
		totals[key] = 0
	}

	for _, instance := range instances {
		limits, err := getInstanceLimits(instance, keys)
		if err != nil {
			return nil, err
		}
		for _, key := range keys {
			totals[key] += limits[key]
		}
	}

	return totals, nil
}

// Return the effective instance-level values for the limits with the given
// keys.
func getInstanceLimits(instance db.Instance, keys []string) (map[string]int64, error) {
	limits := map[string]int64{}

	for _, key := range keys {
		value, ok := instance.Config[key]
		if !ok || value == "" {
			return nil, fmt.Errorf(
				"Instance %s in project %s has no '%s' config, either directly or via a profile",
				instance.Name, instance.Project, key)
		}
		parser := aggregateLimitConfigValueParsers[key]
		limit, err := parser(value)
		if err != nil {
			return nil, errors.Wrapf(
				err, "Parse '%s' for instance %s in project %s",
				key, instance.Name, instance.Project)
		}
		limits[key] = limit
	}

	return limits, nil
}

var aggregateLimitConfigValueParsers = map[string]func(string) (int64, error){
	"limits.memory": units.ParseByteSizeString,
}
<commit_msg>lxd/project: Check that the project's "limits.processes" config is honored<commit_after>package project

import (
	"fmt"
	"strconv"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/units"
	"github.com/pkg/errors"
)

// CheckLimitsUponInstanceCreation returns an error if any project-specific
// limit is violated when creating a new instance.
func CheckLimitsUponInstanceCreation(tx *db.ClusterTx, projectName string, req api.InstancesPost) error {
	project, profiles, instances, err := fetchProject(tx, projectName)
	if err != nil {
		return err
	}

	err = checkInstanceCountLimit(project, len(instances), req.Type)
	if err != nil {
		return err
	}

	// Add the instance being created.
	instances = append(instances, db.Instance{
		Profiles: req.Profiles,
		Config:   req.Config,
	})

	err = checkAggregateInstanceLimits(tx, project, instances, profiles)
	if err != nil {
		return err
	}

	return nil
}

// Check that we have not reached the maximum number of instances for
// this type.
func checkInstanceCountLimit(project *api.Project, instanceCount int, instanceType api.InstanceType) error {
	var key string
	switch instanceType {
	case api.InstanceTypeContainer:
		key = "limits.containers"
	case api.InstanceTypeVM:
		key = "limits.virtual-machines"
	default:
		return fmt.Errorf("Unexpected instance type '%s'", instanceType)
	}
	value, ok := project.Config[key]
	if ok {
		limit, err := strconv.Atoi(value)
		if err != nil || limit < 0 {
			return fmt.Errorf("Unexpected '%s' value: '%s'", key, value)
		}
		if instanceCount >= limit {
			return fmt.Errorf(
				"Reached maximum number of instances of type %s in project %s",
				instanceType, project.Name)
		}
	}

	return nil
}

// Check that we would not violate the project limits if we were to commit the
// given instances and profiles.
func checkAggregateInstanceLimits(tx *db.ClusterTx, project *api.Project, instances []db.Instance, profiles []db.Profile) error {
	// List of config keys for which we need to check aggregate values
	// across all project instances.
	aggregateKeys := []string{}
	for key := range project.Config {
		if shared.StringInSlice(key, []string{"limits.memory", "limits.processes"}) {
			aggregateKeys = append(aggregateKeys, key)
		}
	}
	if len(aggregateKeys) == 0 {
		return nil
	}

	instances = expandInstancesConfig(instances, profiles)

	totals, err := getTotalsAcrossInstances(instances, aggregateKeys)
	if err != nil {
		return err
	}

	for _, key := range aggregateKeys {
		parser := aggregateLimitConfigValueParsers[key]
		max, err := parser(project.Config[key])
		if err != nil {
			return err
		}
		if totals[key] > max {
			return fmt.Errorf(
				"Reached maximum aggregate value %s for %q in project %s",
				project.Config[key], key, project.Name)
		}
	}

	return nil
}

// CheckLimitsUponInstanceUpdate returns an error if any project-specific limit
// is violated when updating an existing instance.
func CheckLimitsUponInstanceUpdate(tx *db.ClusterTx, projectName, instanceName string, req api.InstancePut) error {
	project, profiles, instances, err := fetchProject(tx, projectName)
	if err != nil {
		return err
	}

	// Change the instance being updated.
	for i, instance := range instances {
		if instance.Name != instanceName {
			continue
		}
		instances[i].Profiles = req.Profiles
		instances[i].Config = req.Config
	}

	err = checkAggregateInstanceLimits(tx, project, instances, profiles)
	if err != nil {
		return err
	}

	return nil
}

// CheckLimitsUponProfileUpdate checks that project limits are not violated
// when changing a profile.
func CheckLimitsUponProfileUpdate(tx *db.ClusterTx, projectName, profileName string, req api.ProfilePut) error {
	project, profiles, instances, err := fetchProject(tx, projectName)
	if err != nil {
		return err
	}

	// Change the profile being updated.
	for i, profile := range profiles {
		if profile.Name != profileName {
			continue
		}
		profiles[i].Config = req.Config
	}

	err = checkAggregateInstanceLimits(tx, project, instances, profiles)
	if err != nil {
		return err
	}

	return nil
}

// ValidateLimitsUponProjectUpdate checks the new limits to be set on a project
// are valid.
func ValidateLimitsUponProjectUpdate(tx *db.ClusterTx, projectName string, config map[string]string, changed []string) error {
	_, profiles, instances, err := fetchProject(tx, projectName)
	if err != nil {
		return err
	}

	instances = expandInstancesConfig(instances, profiles)

	// List of keys that need to check aggregate values across all project
	// instances.
	aggregateKeys := []string{}

	for _, key := range changed {
		switch key {
		case "limits.containers":
			fallthrough
		case "limits.virtual-machines":
			err := validateInstanceCountLimit(instances, key, config[key], projectName)
			if err != nil {
				return err
			}
		case "limits.processes":
			fallthrough
		case "limits.memory":
			aggregateKeys = append(aggregateKeys, key)

		}
	}

	if len(aggregateKeys) > 0 {
		totals, err := getTotalsAcrossInstances(instances, aggregateKeys)
		if err != nil {
			return err
		}
		for _, key := range aggregateKeys {
			err := validateAggregateLimit(totals, key, config[key])
			if err != nil {
				return err
			}
		}

	}

	return nil
}

// Check that limits.containers or limits.virtual-machines is equal or above
// the current count.
func validateInstanceCountLimit(instances []db.Instance, key, value, project string) error {
	instanceType := countConfigInstanceType[key]
	limit, err := strconv.Atoi(value)
	if err != nil {
		return err
	}
	dbType, err := instancetype.New(string(instanceType))
	if err != nil {
		return err
	}
	count := 0
	for _, instance := range instances {
		if instance.Type == dbType {
			count++
		}
	}
	if limit < count {
		return fmt.Errorf(
			"'%s' is too low: there currently are %d instances of type %s in project %s",
			key, count, instanceType, project)
	}
	return nil
}

var countConfigInstanceType = map[string]api.InstanceType{
	"limits.containers":       api.InstanceTypeContainer,
	"limits.virtual-machines": api.InstanceTypeVM,
}

// Validates an aggregate limit, checking that the new value is not below the
// current total amount.
func validateAggregateLimit(totals map[string]int64, key, value string) error {
	parser := aggregateLimitConfigValueParsers[key]
	limit, err := parser(value)
	if err != nil {
		errors.Wrapf(err, "Invalid value '%s' for limit %s", value, key)
	}

	total := totals[key]
	if limit < total {
		printer := aggregateLimitConfigValuePrinters[key]
		return fmt.Errorf("'%s' is too low: current total is %s", key, printer(total))
	}

	return nil
}

// Fetch the given project from the database along with its profiles and instances.
func fetchProject(tx *db.ClusterTx, projectName string) (*api.Project, []db.Profile, []db.Instance, error) {
	project, err := tx.ProjectGet(projectName)
	if err != nil {
		return nil, nil, nil, errors.Wrap(err, "Fetch project database object")
	}

	profilesFilter := db.ProfileFilter{}

	// If the project has the profiles feature enabled, we use its own
	// profiles to expand the instances configs, otherwise we use the
	// profiles from the default project.
	if projectName == "default" || project.Config["features.profiles"] == "true" {
		profilesFilter.Project = projectName
	} else {
		profilesFilter.Project = "default"
	}

	profiles, err := tx.ProfileList(profilesFilter)
	if err != nil {
		return nil, nil, nil, errors.Wrap(err, "Fetch profiles from database")
	}

	instances, err := tx.InstanceList(db.InstanceFilter{Project: projectName})
	if err != nil {
		return nil, nil, nil, errors.Wrap(err, "Fetch project instances from database")
	}

	return project, profiles, instances, nil
}

// Expand the configuration of the given instances, taking the give project
// profiles into account.
func expandInstancesConfig(instances []db.Instance, profiles []db.Profile) []db.Instance {
	expandedInstances := make([]db.Instance, len(instances))

	// Index of all profiles by name.
	profilesByName := map[string]db.Profile{}
	for _, profile := range profiles {
		profilesByName[profile.Name] = profile
	}

	for i, instance := range instances {
		profiles := make([]api.Profile, len(instance.Profiles))

		for j, name := range instance.Profiles {
			profile := profilesByName[name]
			profiles[j] = *db.ProfileToAPI(&profile)
		}

		expandedInstances[i] = instance
		expandedInstances[i].Config = db.ProfilesExpandConfig(instance.Config, profiles)
	}

	return expandedInstances
}

// Sum of the effective instance-level value for the given limits across all
// project instances. If excludeInstance is not the empty string, exclude the
// instance with that name.
func getTotalsAcrossInstances(instances []db.Instance, keys []string) (map[string]int64, error) {
	totals := map[string]int64{}

	for _, key := range keys {
		totals[key] = 0
	}

	for _, instance := range instances {
		limits, err := getInstanceLimits(instance, keys)
		if err != nil {
			return nil, err
		}
		for _, key := range keys {
			totals[key] += limits[key]
		}
	}

	return totals, nil
}

// Return the effective instance-level values for the limits with the given
// keys.
func getInstanceLimits(instance db.Instance, keys []string) (map[string]int64, error) {
	limits := map[string]int64{}

	for _, key := range keys {
		value, ok := instance.Config[key]
		if !ok || value == "" {
			return nil, fmt.Errorf(
				"Instance %s in project %s has no '%s' config, either directly or via a profile",
				instance.Name, instance.Project, key)
		}
		parser := aggregateLimitConfigValueParsers[key]
		limit, err := parser(value)
		if err != nil {
			return nil, errors.Wrapf(
				err, "Parse '%s' for instance %s in project %s",
				key, instance.Name, instance.Project)
		}
		limits[key] = limit
	}

	return limits, nil
}

var aggregateLimitConfigValueParsers = map[string]func(string) (int64, error){
	"limits.memory": units.ParseByteSizeString,
	"limits.processes": func(value string) (int64, error) {
		limit, err := strconv.Atoi(value)
		if err != nil {
			return -1, err
		}
		return int64(limit), nil
	},
}

var aggregateLimitConfigValuePrinters = map[string]func(int64) string{
	"limits.memory": func(limit int64) string {
		return units.GetByteSizeString(limit, 1)
	},
	"limits.processes": func(limit int64) string {
		return fmt.Sprintf("%d", limit)
	},
}
<|endoftext|>


<commit_before>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// syz-ci is a continuous fuzzing system for syzkaller.
// It runs several syz-manager's, polls and rebuilds images for managers
// and polls and rebuilds syzkaller binaries.
// For usage instructions see: docs/ci.md
package main

// Implementation details:
//
// 2 main components:
//  - SyzUpdater: handles syzkaller updates
//  - Manager: handles kernel build and syz-manager process (one per manager)
// Both operate in a similar way and keep 2 builds:
//  - latest: latest known good build (i.e. we tested it)
//    preserved across restarts/reboots, i.e. we can start fuzzing even when
//    current syzkaller/kernel git head is broken, or git is down, or anything else
//  - current: currently used build (a copy of one of the latest builds)
// Other important points:
//  - syz-ci is always built on the same revision as the rest of syzkaller binaries,
//    this allows us to handle e.g. changes in manager config format.
//  - consequently, syzkaller binaries are never updated on-the-fly,
//    instead we re-exec and then update
//  - we understand when the latest build is fresh even after reboot,
//    i.e. we store enough information to identify it (git hash, compiler identity, etc),
//    so we don't rebuild unnecessary (kernel builds take time)
//  - we generally avoid crashing the process and handle all errors gracefully
//    (this is a continuous system), except for some severe/user errors during start
//    (e.g. bad config file, or can't create necessary dirs)
//
// Directory/file structure:
// syz-ci			: current executable
// syz-ci.tag			: tag of the current executable (syzkaller git hash)
// syzkaller/
//	latest/			: latest good syzkaller build
//	current/		: syzkaller build currently in use
// managers/
//	manager1/		: one dir per manager
//		kernel/		: kernel checkout
//		workdir/	: manager workdir (never deleted)
//		latest/		: latest good kernel image build
//		current/	: kernel image currently in use
//
// Current executable, syzkaller and kernel builds are marked with tag files.
// Tag files uniquely identify the build (git hash, compiler identity, kernel config, etc).
// For tag files both contents and modification time are important,
// modification time allows us to understand if we need to rebuild after a restart.

import (
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"sync"

	"github.com/google/syzkaller/pkg/config"
	. "github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/syz-manager/mgrconfig"
)

var flagConfig = flag.String("config", "", "config file")

type Config struct {
	Name                   string
	Http                   string
	Dashboard_Addr         string // Optional.
	Hub_Addr               string // Optional.
	Hub_Key                string // Optional.
	Goroot                 string // Go 1.8+ toolchain dir.
	Syzkaller_Repo         string
	Syzkaller_Branch       string
	Syzkaller_Descriptions string // Dir with additional syscall descriptions (.txt and .const files).
	Managers               []*ManagerConfig
}

type ManagerConfig struct {
	Name             string
	Dashboard_Client string
	Dashboard_Key    string
	Repo             string
	Branch           string
	Compiler         string
	Userspace        string
	Kernel_Config    string
	Kernel_Cmdline   string // File with kernel cmdline values (optional).
	Kernel_Sysctl    string // File with sysctl values (e.g. output of sysctl -a, optional).
	Manager_Config   json.RawMessage
}

func main() {
	flag.Parse()
	EnableLogCaching(1000, 1<<20)
	cfg, err := loadConfig(*flagConfig)
	if err != nil {
		Fatalf("failed to load config: %v", err)
	}

	shutdownPending := make(chan struct{})
	osutil.HandleInterrupts(shutdownPending)

	updater := NewSyzUpdater(cfg)
	updater.UpdateOnStart(shutdownPending)
	updatePending := make(chan struct{})
	go func() {
		updater.WaitForUpdate()
		close(updatePending)
	}()

	stop := make(chan struct{})
	go func() {
		select {
		case <-shutdownPending:
		case <-updatePending:
		}
		close(stop)
	}()

	var wg sync.WaitGroup
	wg.Add(len(cfg.Managers))
	managers := make([]*Manager, len(cfg.Managers))
	for i, mgrcfg := range cfg.Managers {
		managers[i] = createManager(cfg, mgrcfg, stop)
	}
	for _, mgr := range managers {
		mgr := mgr
		go func() {
			defer wg.Done()
			mgr.loop()
		}()
	}

	<-stop
	wg.Wait()

	select {
	case <-shutdownPending:
	case <-updatePending:
		updater.UpdateAndRestart()
	}
}

func loadConfig(filename string) (*Config, error) {
	cfg := &Config{
		Syzkaller_Repo:   "https://github.com/google/syzkaller.git",
		Syzkaller_Branch: "master",
		Goroot:           os.Getenv("GOROOT"),
	}
	if err := config.LoadFile(filename, cfg); err != nil {
		return nil, err
	}
	if cfg.Name == "" {
		return nil, fmt.Errorf("param 'name' is empty")
	}
	if cfg.Http == "" {
		return nil, fmt.Errorf("param 'http' is empty")
	}
	if len(cfg.Managers) == 0 {
		return nil, fmt.Errorf("no managers specified")
	}
	for i, mgr := range cfg.Managers {
		if mgr.Name == "" {
			return nil, fmt.Errorf("param 'managers[%v].name' is empty", i)
		}
		mgrcfg := new(mgrconfig.Config)
		if err := config.LoadData(mgr.Manager_Config, mgrcfg); err != nil {
			return nil, fmt.Errorf("manager %v: %v", mgr.Name, err)
		}
	}
	return cfg, nil
}
<commit_msg>syz-ci: add global dashboard client<commit_after>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// syz-ci is a continuous fuzzing system for syzkaller.
// It runs several syz-manager's, polls and rebuilds images for managers
// and polls and rebuilds syzkaller binaries.
// For usage instructions see: docs/ci.md
package main

// Implementation details:
//
// 2 main components:
//  - SyzUpdater: handles syzkaller updates
//  - Manager: handles kernel build and syz-manager process (one per manager)
// Both operate in a similar way and keep 2 builds:
//  - latest: latest known good build (i.e. we tested it)
//    preserved across restarts/reboots, i.e. we can start fuzzing even when
//    current syzkaller/kernel git head is broken, or git is down, or anything else
//  - current: currently used build (a copy of one of the latest builds)
// Other important points:
//  - syz-ci is always built on the same revision as the rest of syzkaller binaries,
//    this allows us to handle e.g. changes in manager config format.
//  - consequently, syzkaller binaries are never updated on-the-fly,
//    instead we re-exec and then update
//  - we understand when the latest build is fresh even after reboot,
//    i.e. we store enough information to identify it (git hash, compiler identity, etc),
//    so we don't rebuild unnecessary (kernel builds take time)
//  - we generally avoid crashing the process and handle all errors gracefully
//    (this is a continuous system), except for some severe/user errors during start
//    (e.g. bad config file, or can't create necessary dirs)
//
// Directory/file structure:
// syz-ci			: current executable
// syz-ci.tag			: tag of the current executable (syzkaller git hash)
// syzkaller/
//	latest/			: latest good syzkaller build
//	current/		: syzkaller build currently in use
// managers/
//	manager1/		: one dir per manager
//		kernel/		: kernel checkout
//		workdir/	: manager workdir (never deleted)
//		latest/		: latest good kernel image build
//		current/	: kernel image currently in use
//
// Current executable, syzkaller and kernel builds are marked with tag files.
// Tag files uniquely identify the build (git hash, compiler identity, kernel config, etc).
// For tag files both contents and modification time are important,
// modification time allows us to understand if we need to rebuild after a restart.

import (
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"sync"

	"github.com/google/syzkaller/pkg/config"
	. "github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/syz-manager/mgrconfig"
)

var flagConfig = flag.String("config", "", "config file")

type Config struct {
	Name                   string
	Http                   string
	Dashboard_Addr         string // Optional.
	Dashboard_Client       string // Optional.
	Dashboard_Key          string // Optional.
	Hub_Addr               string // Optional.
	Hub_Key                string // Optional.
	Goroot                 string // Go 1.8+ toolchain dir.
	Syzkaller_Repo         string
	Syzkaller_Branch       string
	Syzkaller_Descriptions string // Dir with additional syscall descriptions (.txt and .const files).
	Managers               []*ManagerConfig
}

type ManagerConfig struct {
	Name             string
	Dashboard_Client string
	Dashboard_Key    string
	Repo             string
	Branch           string
	Compiler         string
	Userspace        string
	Kernel_Config    string
	Kernel_Cmdline   string // File with kernel cmdline values (optional).
	Kernel_Sysctl    string // File with sysctl values (e.g. output of sysctl -a, optional).
	Manager_Config   json.RawMessage
}

func main() {
	flag.Parse()
	EnableLogCaching(1000, 1<<20)
	cfg, err := loadConfig(*flagConfig)
	if err != nil {
		Fatalf("failed to load config: %v", err)
	}

	shutdownPending := make(chan struct{})
	osutil.HandleInterrupts(shutdownPending)

	updater := NewSyzUpdater(cfg)
	updater.UpdateOnStart(shutdownPending)
	updatePending := make(chan struct{})
	go func() {
		updater.WaitForUpdate()
		close(updatePending)
	}()

	stop := make(chan struct{})
	go func() {
		select {
		case <-shutdownPending:
		case <-updatePending:
		}
		close(stop)
	}()

	var wg sync.WaitGroup
	wg.Add(len(cfg.Managers))
	managers := make([]*Manager, len(cfg.Managers))
	for i, mgrcfg := range cfg.Managers {
		managers[i] = createManager(cfg, mgrcfg, stop)
	}
	for _, mgr := range managers {
		mgr := mgr
		go func() {
			defer wg.Done()
			mgr.loop()
		}()
	}

	<-stop
	wg.Wait()

	select {
	case <-shutdownPending:
	case <-updatePending:
		updater.UpdateAndRestart()
	}
}

func loadConfig(filename string) (*Config, error) {
	cfg := &Config{
		Syzkaller_Repo:   "https://github.com/google/syzkaller.git",
		Syzkaller_Branch: "master",
		Goroot:           os.Getenv("GOROOT"),
	}
	if err := config.LoadFile(filename, cfg); err != nil {
		return nil, err
	}
	if cfg.Name == "" {
		return nil, fmt.Errorf("param 'name' is empty")
	}
	if cfg.Http == "" {
		return nil, fmt.Errorf("param 'http' is empty")
	}
	if len(cfg.Managers) == 0 {
		return nil, fmt.Errorf("no managers specified")
	}
	for i, mgr := range cfg.Managers {
		if mgr.Name == "" {
			return nil, fmt.Errorf("param 'managers[%v].name' is empty", i)
		}
		mgrcfg := new(mgrconfig.Config)
		if err := config.LoadData(mgr.Manager_Config, mgrcfg); err != nil {
			return nil, fmt.Errorf("manager %v: %v", mgr.Name, err)
		}
	}
	return cfg, nil
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package utils

import (
	"bytes"
	"compress/gzip"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"unicode"

	"launchpad.net/goyaml"
)

// WriteYaml marshals obj as yaml and then writes it to a file, atomically,
// by first writing a sibling with the suffix ".preparing" and then moving
// the sibling to the real path.
func WriteYaml(path string, obj interface{}) error {
	data, err := goyaml.Marshal(obj)
	if err != nil {
		return err
	}
	prep := path + ".preparing"
	f, err := os.OpenFile(prep, os.O_WRONLY|os.O_CREATE|os.O_SYNC, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	if _, err = f.Write(data); err != nil {
		return err
	}
	return ReplaceFile(prep, path)
}

// ReadYaml unmarshals the yaml contained in the file at path into obj. See
// goyaml.Unmarshal.
func ReadYaml(path string, obj interface{}) error {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}
	return goyaml.Unmarshal(data, obj)
}

// ErrorContextf prefixes any error stored in err with text formatted
// according to the format specifier. If err does not contain an error,
// ErrorContextf does nothing.
func ErrorContextf(err *error, format string, args ...interface{}) {
	if *err != nil {
		*err = errors.New(fmt.Sprintf(format, args...) + ": " + (*err).Error())
	}
}

// ShQuote quotes s so that when read by bash, no metacharacters
// within s will be interpreted as such.
func ShQuote(s string) string {
	// single-quote becomes single-quote, double-quote, single-quote, double-quote, single-quote
	return `'` + strings.Replace(s, `'`, `'"'"'`, -1) + `'`
}

// CommandString flattens a sequence of command arguments into a
// string suitable for executing in a shell, escaping slashes,
// variables and quotes as necessary; each argument is double-quoted
// if and only if necessary.
func CommandString(args ...string) string {
	var buf bytes.Buffer
	for i, arg := range args {
		needsQuotes := false
		var argBuf bytes.Buffer
		for _, r := range arg {
			if unicode.IsSpace(r) {
				needsQuotes = true
			} else if r == '"' || r == '$' || r == '\\' {
				needsQuotes = true
				argBuf.WriteByte('\\')
			}
			argBuf.WriteRune(r)
		}
		if i > 0 {
			buf.WriteByte(' ')
		}
		if needsQuotes {
			buf.WriteByte('"')
			argBuf.WriteTo(&buf)
			buf.WriteByte('"')
		} else {
			argBuf.WriteTo(&buf)
		}
	}
	return buf.String()
}

// Gzip compresses the given data.
func Gzip(data []byte) []byte {
	var buf bytes.Buffer
	w := gzip.NewWriter(&buf)
	if _, err := w.Write(data); err != nil {
		// Compression should never fail unless it fails
		// to write to the underlying writer, which is a bytes.Buffer
		// that never fails.
		panic(err)
	}
	if err := w.Close(); err != nil {
		panic(err)
	}
	return buf.Bytes()
}

// Gunzip uncompresses the given data.
func Gunzip(data []byte) ([]byte, error) {
	r, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	return ioutil.ReadAll(r)
}
<commit_msg>Changes after review<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package utils

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"
	"unicode"

	"launchpad.net/goyaml"
)

// WriteYaml marshals obj as yaml and then writes it to a file, atomically,
// by first writing a sibling with the suffix ".preparing" and then moving
// the sibling to the real path.
func WriteYaml(path string, obj interface{}) error {
	data, err := goyaml.Marshal(obj)
	if err != nil {
		return err
	}
	prep := path + ".preparing"
	f, err := os.OpenFile(prep, os.O_WRONLY|os.O_CREATE|os.O_SYNC, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	if _, err = f.Write(data); err != nil {
		return err
	}
	return ReplaceFile(prep, path)
}

// ReadYaml unmarshals the yaml contained in the file at path into obj. See
// goyaml.Unmarshal.
func ReadYaml(path string, obj interface{}) error {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}
	return goyaml.Unmarshal(data, obj)
}

// ErrorContextf prefixes any error stored in err with text formatted
// according to the format specifier. If err does not contain an error,
// ErrorContextf does nothing.
func ErrorContextf(err *error, format string, args ...interface{}) {
	if *err != nil {
		*err = errors.New(fmt.Sprintf(format, args...) + ": " + (*err).Error())
	}
}

// ShQuote quotes s so that when read by bash, no metacharacters
// within s will be interpreted as such.
func ShQuote(s string) string {
	// single-quote becomes single-quote, double-quote, single-quote, double-quote, single-quote
	return `'` + strings.Replace(s, `'`, `'"'"'`, -1) + `'`
}

// CommandString flattens a sequence of command arguments into a
// string suitable for executing in a shell, escaping slashes,
// variables and quotes as necessary; each argument is double-quoted
// if and only if necessary.
func CommandString(args ...string) string {
	var buf bytes.Buffer
	for i, arg := range args {
		needsQuotes := false
		var argBuf bytes.Buffer
		for _, r := range arg {
			if unicode.IsSpace(r) {
				needsQuotes = true
			} else if r == '"' || r == '$' || r == '\\' {
				needsQuotes = true
				argBuf.WriteByte('\\')
			}
			argBuf.WriteRune(r)
		}
		if i > 0 {
			buf.WriteByte(' ')
		}
		if needsQuotes {
			buf.WriteByte('"')
			argBuf.WriteTo(&buf)
			buf.WriteByte('"')
		} else {
			argBuf.WriteTo(&buf)
		}
	}
	return buf.String()
}

// Gzip compresses the given data.
func Gzip(data []byte) []byte {
	var buf bytes.Buffer
	w := gzip.NewWriter(&buf)
	if _, err := w.Write(data); err != nil {
		// Compression should never fail unless it fails
		// to write to the underlying writer, which is a bytes.Buffer
		// that never fails.
		panic(err)
	}
	if err := w.Close(); err != nil {
		panic(err)
	}
	return buf.Bytes()
}

// Gunzip uncompresses the given data.
func Gunzip(data []byte) ([]byte, error) {
	r, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	return ioutil.ReadAll(r)
}

// GetSHA256 returns the SHA256 hash of the contents read from source
// (hex encoded) and the size of the source in bytes.
func GetSHA256(source io.Reader) (string, int64, error) {
	hash := sha256.New()
	size, err := io.Copy(hash, source)
	if err != nil {
		return "", 0, err
	}
	digest := hex.EncodeToString(hash.Sum(nil))
	return digest, size, nil
}

// GetFileSHA256 is like GetSHA256 but reads the contents of the given
// file.
func GetFileSHA256(filename string) (string, int64, error) {
	f, err := os.Open(filename)
	if err != nil {
		return "", 0, err
	}
	defer f.Close()
	return GetSHA256(f)
}
<|endoftext|>


<commit_before>package talker

import (
	"fmt"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
)

// Notifaction new the notifaction object
type Notifaction struct {
	channelName     string
	projectName     string
	isTestSucceed   bool
	isBuildSucceed  bool
	isDepolySucceed bool
	commitNotes     string
}

const (
	testStage   = "test"
	buildStage  = "build"
	deployStage = "deploy"
)

const (
	// ChannelChatID defines the bot name
	// Hacked this on the web client
	// Regex format: *c(.*)_*
	ChannelChatID = 1378084890
	// ChannelName defines build/deploy channel name
	ChannelName = "@AstralServerNotifaction"
)

// NewNotifaction init the Notifaction instance
func NewNotifaction(repo string, stage map[string]bool, commit string) *Notifaction {
	if _, ok := stage[testStage]; !ok {
		return &Notifaction{}
	}
	if _, ok := stage[buildStage]; !ok {
		return &Notifaction{}
	}
	return &Notifaction{
		channelName:     "",
		projectName:     repo,
		isTestSucceed:   stage[testStage],
		isBuildSucceed:  stage[buildStage],
		isDepolySucceed: stage[deployStage],
		commitNotes:     commit,
	}
}

// Notify sends the msg to the tg channel
func (n *Notifaction) Notify() tgbotapi.MessageConfig {
	text := fmt.Sprintf("%s:%v", n.commitNotes, n.isBuildSucceed)
	return tgbotapi.NewMessageToChannel(ChannelName, text)
}
<commit_msg>talker:Format build status<commit_after>package talker

import (
	"fmt"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
)

// Notifaction new the notifaction object
type Notifaction struct {
	channelName     string
	projectName     string
	isTestSucceed   string
	isBuildSucceed  string
	isDepolySucceed string
	commitNotes     string
}

const (
	testStage   = "test"
	buildStage  = "build"
	deployStage = "deploy"
)

const (
	// ChannelChatID defines the bot name
	// Hacked this on the web client
	// Regex format: *c(.*)_*
	ChannelChatID = 1378084890
	// ChannelName defines build/deploy channel name
	ChannelName = "@AstralServerNotifaction"

	successIcon = "✅"

	failedIcon = "❌"
)

func convert2Icon(isSuccess bool) string {
	if isSuccess {
		return successIcon
	}
	return failedIcon
}

// NewNotifaction init the Notifaction instance
func NewNotifaction(repo string, stage map[string]bool, commit string) *Notifaction {
	if _, ok := stage[testStage]; !ok {
		return &Notifaction{}
	}
	if _, ok := stage[buildStage]; !ok {
		return &Notifaction{}
	}
	testStatus := convert2Icon(stage[testStage])
	buildStatus := convert2Icon(stage[buildStage])
	deployStatus := convert2Icon(stage[deployStage])

	return &Notifaction{
		channelName:     "",
		projectName:     repo,
		isTestSucceed:   testStatus,
		isBuildSucceed:  buildStatus,
		isDepolySucceed: deployStatus,
		commitNotes:     commit,
	}
}

// Notify sends the msg to the tg channel
func (n *Notifaction) Notify() tgbotapi.MessageConfig {
	text := fmt.Sprintf("**Commit Note**: `%s` ", n.commitNotes)
	text = fmt.Sprintf("%s\n **Test Status**:%v", text, n.isTestSucceed)
	text = fmt.Sprintf("%s\n **Build Status**:%v", text, n.isBuildSucceed)
	text = fmt.Sprintf("%s\n **Deploy Status**:%v", text, n.isDepolySucceed)
	return tgbotapi.MessageConfig{
		BaseChat: tgbotapi.BaseChat{
			ChannelUsername: ChannelName,
		},
		Text:      text,
		ParseMode: tgbotapi.ModeMarkdown,
	}
}
<|endoftext|>


<commit_before>package ping

import (
	"bytes"
	"network/icmpp"
	"time"

	"network/ipv4p"

	"github.com/hsheth2/logs"
)

func (pm *Ping_Manager) ping_response_dealer() {
	for {
		ping := <-pm.reply
		identNum := uint16(ping.Header.Opt >> 16)
		if _, ok := pm.identifiers[identNum]; !ok {
			logs.Info.Println("Dropped something from response dealer, identnum=", identNum, "options=", ping.Header.Opt)
			continue
		}
		pm.identifiers[identNum] <- ping
	}
}

func sendSinglePing(writer *ipv4p.IP_Writer, id, seq uint16, timeout time.Duration, reply chan *icmpp.ICMP_In) {
	// prepare packet
	packet := &icmpp.ICMP_Header{
		TypeF: PING_ECHO_REQUEST_TYPE,
		Code:  PING_ICMP_CODE,
		Opt:   uint32(id)<<16 | uint32(seq),
		Data:  []byte("abcdefg"), // TODO make legit by sending 56 bytes of Data and putting the timestamp in the data
	}

	// make data
	data, err := packet.MarshalICMPHeader()
	if err != nil {
		logs.Error.Println(err)
		return
	}

	// send
	err = writer.WriteTo(data)
	if err != nil {
		logs.Error.Println(err)
		return
	}
	time1 := time.Now()
	timeoutTimer := time.NewTimer(timeout)
	go func(seqChan chan *icmpp.ICMP_In, header *icmpp.ICMP_Header, time1 *time.Time, timer *time.Timer) {
		for {
			select {
			case pingResonse := <-seqChan:
				if !bytes.Equal(pingResonse.Header.Data, header.Data) {
					logs.Info.Println("Dropped packet cuz data !=")
					continue
				}
				time2 := time.Now()
				logs.Info.Printf("%d bytes from %s: icmp_seq=%d time=%f ms",
					len(header.Data)+icmpp.ICMP_Header_MinSize,
					pingResonse.RIP,
					uint16(header.Opt),
					time2.Sub(*time1).Seconds()*1000) // put ttl
				return
			case <-timer.C:
				logs.Info.Println("Seq num of", uint16(header.Opt), "timed out")
				return
			}
		}
	}(reply, packet, &time1, timeoutTimer)
}

func (pm *Ping_Manager) initIdentifier(terminate chan bool) (id uint16, seqChannel map[uint16](chan *icmpp.ICMP_In), err error) {
	// get identifier
	pm.currentIdentifier++
	id = pm.currentIdentifier

	// setup sequence number dealer
	pm.identifiers[id] = make(chan *icmpp.ICMP_In)
	seqChannel = make(map[uint16](chan *icmpp.ICMP_In))

	// create go routine function to deal packets
	go sequenceDealer(pm.identifiers[id], seqChannel, terminate)

	return id, seqChannel, nil
}

func sequenceDealer(idInput chan *icmpp.ICMP_In, seqChan map[uint16](chan *icmpp.ICMP_In), terminate chan bool) {
	// TODO verify IPs
	for {
		select {
		case <-terminate:
			logs.Info.Println("Terminating seq dealer")
			return
		case packet := <-idInput:
			// logs.Info.Println("icmp in =", packet.Header.Opt)
			seqNum := uint16(packet.Header.Opt)
			if _, ok := seqChan[seqNum]; ok {
				seqChan[seqNum] <- packet
			} else {
				logs.Info.Println("Dropping bad seq num packet with existing identifier")
			}
		}
	}
}

func (pm *Ping_Manager) SendPing(ip string, interval, timeout time.Duration, numPings uint16) error {
	terminate := make(chan bool)
	id, seqChannel, err := pm.initIdentifier(terminate)
	if err != nil {
		logs.Error.Println(err)
		return err
	}
	defer func(){ terminate <- true }()

	// get ip writer
	writer, err := pm.getIP_Writer(ip)
	if err != nil {
		logs.Error.Println(err)
		return err
	}

	for i := uint16(1); i <= numPings; i++ {
		seqChannel[i] = make(chan *icmpp.ICMP_In)

		sendSinglePing(writer, id, i, timeout, seqChannel[i]) // function is non-blocking

		// not last
		if i != numPings {
			time.Sleep(interval)
		}
	}

	time.Sleep(timeout)
	return nil
}
<commit_msg>Made timestamp better<commit_after>package ping

import (
	"bytes"
	"network/icmpp"
	"time"

	"network/ipv4p"

	"github.com/hsheth2/logs"
)

func (pm *Ping_Manager) ping_response_dealer() {
	for {
		ping := <-pm.reply
		identNum := uint16(ping.Header.Opt >> 16)
		if _, ok := pm.identifiers[identNum]; !ok {
			logs.Info.Println("Dropped something from response dealer, identnum=", identNum, "options=", ping.Header.Opt)
			continue
		}
		pm.identifiers[identNum] <- ping
	}
}

func sendSinglePing(writer *ipv4p.IP_Writer, id, seq uint16, timeout time.Duration, reply chan *icmpp.ICMP_In) {
	// prepare packet
	packet := &icmpp.ICMP_Header{
		TypeF: PING_ECHO_REQUEST_TYPE,
		Code:  PING_ICMP_CODE,
		Opt:   uint32(id)<<16 | uint32(seq),
		Data:  []byte("abcdefg"), // TODO make legit by sending 56 bytes of Data and putting the timestamp in the data
	}

	// make data
	data, err := packet.MarshalICMPHeader()
	if err != nil {
		logs.Error.Println(err)
		return
	}

	// send
	err = writer.WriteTo(data)
	if err != nil {
		logs.Error.Println(err)
		return
	}
	time1 := time.Now()
	timeoutTimer := time.NewTimer(timeout)
	go func(seqChan chan *icmpp.ICMP_In, header *icmpp.ICMP_Header, time1 *time.Time, timer *time.Timer) {
		for {
			select {
			case pingResonse := <-seqChan:
				if !bytes.Equal(pingResonse.Header.Data, header.Data) {
					logs.Info.Println("Dropped packet cuz data !=")
					continue
				}
				time2 := time.Now()
				logs.Info.Printf("%d bytes from %s: icmp_seq=%d time=%f ms",
					len(header.Data)+icmpp.ICMP_Header_MinSize,
					pingResonse.RIP,
					uint16(header.Opt),
					float32(time2.Sub(*time1).Nanoseconds())/1000000) // put ttl
				return
			case <-timer.C:
				logs.Info.Println("Seq num of", uint16(header.Opt), "timed out")
				return
			}
		}
	}(reply, packet, &time1, timeoutTimer)
}

func (pm *Ping_Manager) initIdentifier(terminate chan bool) (id uint16, seqChannel map[uint16](chan *icmpp.ICMP_In), err error) {
	// get identifier
	pm.currentIdentifier++
	id = pm.currentIdentifier

	// setup sequence number dealer
	pm.identifiers[id] = make(chan *icmpp.ICMP_In)
	seqChannel = make(map[uint16](chan *icmpp.ICMP_In))

	// create go routine function to deal packets
	go sequenceDealer(pm.identifiers[id], seqChannel, terminate)

	return id, seqChannel, nil
}

func sequenceDealer(idInput chan *icmpp.ICMP_In, seqChan map[uint16](chan *icmpp.ICMP_In), terminate chan bool) {
	// TODO verify IPs
	for {
		select {
		case <-terminate:
			logs.Info.Println("Terminating seq dealer")
			return
		case packet := <-idInput:
			// logs.Info.Println("icmp in =", packet.Header.Opt)
			seqNum := uint16(packet.Header.Opt)
			if _, ok := seqChan[seqNum]; ok {
				seqChan[seqNum] <- packet
			} else {
				logs.Info.Println("Dropping bad seq num packet with existing identifier")
			}
		}
	}
}

func (pm *Ping_Manager) SendPing(ip string, interval, timeout time.Duration, numPings uint16) error {
	terminate := make(chan bool)
	id, seqChannel, err := pm.initIdentifier(terminate)
	if err != nil {
		logs.Error.Println(err)
		return err
	}
	defer func(){ terminate <- true }()

	// get ip writer
	writer, err := pm.getIP_Writer(ip)
	if err != nil {
		logs.Error.Println(err)
		return err
	}

	for i := uint16(1); i <= numPings; i++ {
		seqChannel[i] = make(chan *icmpp.ICMP_In)

		sendSinglePing(writer, id, i, timeout, seqChannel[i]) // function is non-blocking

		// not last
		if i != numPings {
			time.Sleep(interval)
		}
	}

	time.Sleep(timeout)
	return nil
}
<|endoftext|>