<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/tarm/goserial"

	log "github.com/cihub/seelog"
	"github.com/stampzilla/stampzilla-go/nodes/basenode"
	"github.com/stampzilla/stampzilla-go/protocol"
)

type SerialConnection struct {
	Name string
	Baud int
	Port io.ReadWriteCloser
}

var node *protocol.Node
var state *State = &State{}
var serverConnection *basenode.Connection
var ard *SerialConnection
var disp *SerialConnection

func main() {
	config := basenode.NewConfig()
	basenode.SetConfig(config)

	var arddev string
	var dispdev string
	flag.StringVar(&arddev, "arduino-dev", "/dev/arduino", "Arduino serial port")
	flag.StringVar(&dispdev, "display-dev", "/dev/poleDisplay", "Display serial port")

	flag.Parse()

	log.Info("Starting Aquarium node")

	// Create new node description
	node = protocol.NewNode("aquarium")
	node.SetState(state)

	node.AddElement(&protocol.Element{
		Type:     protocol.ElementTypeText,
		Name:     "Water temperature",
		Feedback: `WaterTemperature`,
	})

	serverConnection = basenode.Connect()
	go monitorState(serverConnection)

	// This worker recives all incomming commands
	go serverRecv(serverConnection)

	// Setup the serial connection
	ard = &SerialConnection{Name: arddev, Baud: 115200}
	ard.run(serverConnection, func(data string, connection *basenode.Connection) {
		processArduinoData(data, connection)
	})

	disp = &SerialConnection{Name: dispdev, Baud: 9600}
	disp.run(serverConnection, func(data string, connection *basenode.Connection) {
		log.Debug("Incomming from display", data)
	})

	go updateDisplay(serverConnection, disp)

	select {}
}

var updateInhibit bool = false
var changed bool = false

func processArduinoData(msg string, connection *basenode.Connection) { // {{{
	log.Debug(msg)

	var prevState State = *state

	values := strings.Split(msg, "|")
	if len(values) != 6 {
		return
	}

	// Temperature// {{{
	value, err := strconv.ParseFloat(values[0], 32)
	if err != nil {
		return
	}
	//if value != state.WaterTemperature {
	//changed = true
	//}
	state.WaterTemperature = value // }}}

	// Filling stat0e // {{{
	value, err = strconv.ParseFloat(values[1], 32)
	if err != nil {
		return
	}
	//if state.FillingTime != value {
	//changed = true
	//}
	state.FillingTime = value // }}}

	// Cooling %// {{{
	value, err = strconv.ParseFloat(values[2], 32)
	if err != nil {
		return
	}
	//if state.Cooling != value {
	//changed = true
	//}
	state.Cooling = value // }}}

	bits := values[3][0]

	state.CirculationPumps = bits&0x01 > 0
	state.Skimmer = bits&0x02 > 0
	state.Heating = bits&0x04 > 0
	state.Filling = bits&0x08 == 0
	state.WaterLevelOk = bits&0x0F == 0
	state.FilterOk = bits&0x10 == 0

	// Check if something have changed
	changed = reflect.DeepEqual(prevState, *state)

	if !updateInhibit && changed {
		changed = false
		connection.Send <- node.Node()
		updateInhibit = true
		//log.Warn("Invalid pacakge: ", msg)
		go func() {
			<-time.After(time.Millisecond * 100)
			updateInhibit = false
		}()
	}
}                                                                               // }}}
func updateDisplay(connection *basenode.Connection, serial *SerialConnection) { // {{{
	for {
		select {
		case <-time.After(time.Second / 15):
			if serial.Port == nil {
				continue
			}

			var msg, row1, row2 string
			var wLevel string = "OK"

			//if !state.WaterLevelOk {
			//wLevel = "WLVL"
			//}
			if state.FillingTime > 0 {
				wLevel = "FILL ERR"
				if state.FillingTime < 10000 {
					wLevel = "FILLING"
				}
			}

			//msg += "\x1B" // Reset
			//msg += "\x0E" // Clear
			//msg += "\x13" // Cursor on
			msg += "\x11" // Character over write mode
			//msg += "\x12" // Vertical scroll mode
			//msg += "\x13" // Cursor on
			msg += "\x14" // Cursor off

			msg += "\r"

			row1 = time.Now().Format("15:04:05") + strings.Repeat(" ", 12-len(wLevel)) + wLevel
			row2 = "Temp " + strconv.FormatFloat(state.WaterTemperature, 'f', 2, 64) + " Cool " + strconv.FormatFloat(state.Cooling, 'f', 2, 64)

			if 20-len(row1) > 0 {
				row1 += strings.Repeat(" ", 20-len(row1))
			}

			if 20-len(row2) > 0 {
				row2 += strings.Repeat(" ", 20-len(row2))
			}
			msg += row1[0:20]
			msg += row2[0:20]

			_, err := serial.Port.Write([]byte(msg))
			if err != nil {
				log.Debug("Failed write to display:", err)
			}
		}
	}
} // }}}

// WORKER that monitors the current connection state
func monitorState(connection *basenode.Connection) { // {{{
	for s := range connection.State {
		switch s {
		case basenode.ConnectionStateConnected:
			connection.Send <- node.Node()
		case basenode.ConnectionStateDisconnected:
		}
	}
} // }}}

// WORKER that recives all incomming commands
func serverRecv(connection *basenode.Connection) { // {{{
	for d := range connection.Receive {
		if err := processCommand(d); err != nil {
			log.Error(err)
		}
	}
} // }}}

func processCommand(cmd protocol.Command) error { // {{{
	if len(cmd.Args) < 1 {
		return fmt.Errorf("Missing arguments, ignoring command")
	}

	switch cmd.Cmd {
	case "CirculationPumps":
		if cmd.Args[0] != "" && cmd.Args[0] != "0" {
			ard.Port.Write([]byte{0x02, 0x01, 0x01, 0x03}) // Turn on
		} else {
			ard.Port.Write([]byte{0x02, 0x01, 0x00, 0x03}) // Turn off
		}
		break
	case "Skimmer":
		if cmd.Args[0] != "" && cmd.Args[0] != "0" {
			ard.Port.Write([]byte{0x02, 0x02, 0x01, 0x03}) // Turn on
		} else {
			ard.Port.Write([]byte{0x02, 0x02, 0x00, 0x03}) // Turn off
		}
		break
	case "Heating":
		if cmd.Args[0] != "" && cmd.Args[0] != "0" {
			ard.Port.Write([]byte{0x02, 0x03, 0x01, 0x03}) // Turn on
		} else {
			ard.Port.Write([]byte{0x02, 0x03, 0x00, 0x03}) // Turn off
		}
		break
	case "CoolingP":
		i, err := strconv.Atoi(cmd.Args[0])
		if err != nil {
			return fmt.Errorf("Failed to decode arg[0] to int %s %s", err, cmd.Args[0])
		}

		ard.Port.Write([]byte{0x02, 0x04, byte(i), 0x03}) // Turn on
		break
	case "CoolingI":
		i, err := strconv.Atoi(cmd.Args[0])
		if err != nil {
			return fmt.Errorf("Failed to decode arg[0] to int %s %s", err, cmd.Args[0])
		}

		ard.Port.Write([]byte{0x02, 0x05, byte(i), 0x03}) // Turn on
		break
	case "CoolingD":
		i, err := strconv.Atoi(cmd.Args[0])
		if err != nil {
			return fmt.Errorf("Failed to decode arg[0] to int %s %s", err, cmd.Args[0])
		}

		ard.Port.Write([]byte{0x02, 0x06, byte(i), 0x03}) // Turn on
		break
	}
	return nil
} // }}}

// Serial connection workers
func (config *SerialConnection) run(connection *basenode.Connection, callback func(data string, connection *basenode.Connection)) { // {{{
	go func() {
		for {
			config.connect(connection, callback)
			<-time.After(time.Second * 4)
		}
	}()
}                                                                                // }}}
func (config *SerialConnection) connect(connection *basenode.Connection, callback func(data string, connection *basenode.Connection)) { // {{{

	c := &serial.Config{Name: config.Name, Baud: config.Baud}
	var err error

	config.Port, err = serial.OpenPort(c)
	if err != nil {
		log.Error("Serial connect failed: ", err)
		return
	}

	var incomming string = ""

	for {
		buf := make([]byte, 128)

		n, err := config.Port.Read(buf)
		if err != nil {
			log.Error("Serial read failed: ", err)
			return
		}

		incomming += string(buf[:n])

		for {
			n := strings.Index(incomming, "\r")

			if n < 0 {
				break
			}

			msg := strings.TrimSpace(incomming[:n])
			incomming = incomming[n+1:]
			//log.Info(msg)

			go callback(msg, connection)
		}
	}
} // }}}
<commit_msg>New values and added a timeout when server is not reciving data<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/tarm/goserial"

	log "github.com/cihub/seelog"
	"github.com/stampzilla/stampzilla-go/nodes/basenode"
	"github.com/stampzilla/stampzilla-go/protocol"
)

type SerialConnection struct {
	Name string
	Baud int
	Port io.ReadWriteCloser
}

var node *protocol.Node
var state *State = &State{}
var serverConnection *basenode.Connection
var ard *SerialConnection
var disp *SerialConnection
var frameCount int

func main() {
	config := basenode.NewConfig()
	basenode.SetConfig(config)

	var arddev string
	var dispdev string
	flag.StringVar(&arddev, "arduino-dev", "/dev/arduino", "Arduino serial port")
	flag.StringVar(&dispdev, "display-dev", "/dev/poleDisplay", "Display serial port")

	flag.Parse()

	log.Info("Starting Aquarium node")

	// Create new node description
	node = protocol.NewNode("aquarium")
	node.SetState(state)

	node.AddElement(&protocol.Element{
		Type:     protocol.ElementTypeText,
		Name:     "Water temperature",
		Feedback: `WaterTemperature`,
	})
	node.AddElement(&protocol.Element{
		Type:     protocol.ElementTypeText,
		Name:     "Water level - OK",
		Feedback: `WaterLevelOk`,
	})
	node.AddElement(&protocol.Element{
		Type:     protocol.ElementTypeText,
		Name:     "Filter - OK",
		Feedback: `FilterOk`,
	})

	node.AddElement(&protocol.Element{
		Type:     protocol.ElementTypeText,
		Name:     "Cooling",
		Feedback: `Cooling`,
	})
	node.AddElement(&protocol.Element{
		Type: protocol.ElementTypeToggle,
		Name: "Heating",
		Command: &protocol.Command{
			Cmd:  "Heating",
			Args: []string{},
		},
		Feedback: `Heating`,
	})

	node.AddElement(&protocol.Element{
		Type: protocol.ElementTypeToggle,
		Name: "Skimmer",
		Command: &protocol.Command{
			Cmd:  "Skimmer",
			Args: []string{},
		},
		Feedback: `Skimmer`,
	})
	node.AddElement(&protocol.Element{
		Type: protocol.ElementTypeToggle,
		Name: "Circulation pumps",
		Command: &protocol.Command{
			Cmd:  "CirculationPumps",
			Args: []string{},
		},
		Feedback: `CirculationPumps`,
	})

	serverConnection = basenode.Connect()
	go monitorState(serverConnection)

	// This worker recives all incomming commands
	go serverRecv(serverConnection)

	// Setup the serial connection
	ard = &SerialConnection{Name: arddev, Baud: 115200}
	ard.run(serverConnection, func(data string, connection *basenode.Connection) {
		processArduinoData(data, connection)
	})

	disp = &SerialConnection{Name: dispdev, Baud: 9600}
	disp.run(serverConnection, func(data string, connection *basenode.Connection) {
		log.Debug("Incomming from display", data)
	})

	go updateDisplay(serverConnection, disp)

	go func() {
		for {
			fmt.Print("\r", frameCount)
			<-time.After(time.Millisecond * 10)
		}
	}()

	select {}
}

var updateInhibit bool = false
var changed bool = false

func processArduinoData(msg string, connection *basenode.Connection) { // {{{
	//log.Debug(msg)
	frameCount++

	var prevState State = *state

	values := strings.Split(msg, "|")
	if len(values) != 6 {
		return
	}

	// Temperature// {{{
	value, err := strconv.ParseFloat(values[0], 32)
	if err != nil {
		return
	}
	if value != state.WaterTemperature {
		changed = true
	}
	state.WaterTemperature = value // }}}

	// Filling stat0e // {{{
	value, err = strconv.ParseFloat(values[1], 32)
	if err != nil {
		return
	}
	if state.FillingTime != value {
		changed = true
	}
	state.FillingTime = value // }}}

	// Cooling %// {{{
	value, err = strconv.ParseFloat(values[2], 32)
	if err != nil {
		return
	}
	if state.Cooling != value {
		changed = true
	}
	state.Cooling = value // }}}

	bits := values[3][0]

	state.CirculationPumps = bits&0x01 != 0
	state.Skimmer = bits&0x02 != 0
	state.Heating = bits&0x04 != 0
	state.Filling = bits&0x08 == 0
	state.WaterLevelOk = bits&0x0F == 0
	state.FilterOk = bits&0x10 == 0

	// Check if something have changed
	if !reflect.DeepEqual(prevState, *state) {
		changed = true
	}

	if !updateInhibit && changed {
		changed = false
		//fmt.Print("\n")

		go func() {
			select {
			case connection.Send <- node.Node():
			case <-time.After(time.Second * 1):
				log.Warn("TIMEOUT: Failed to send update to server")
			}
		}()

		updateInhibit = true
		//log.Warn("Invalid pacakge: ", msg)
		go func() {
			<-time.After(time.Millisecond * 200)
			updateInhibit = false
		}()
	}
}                                                                               // }}}
func updateDisplay(connection *basenode.Connection, serial *SerialConnection) { // {{{
	for {
		select {
		case <-time.After(time.Second / 15):
			if serial.Port == nil {
				continue
			}

			var msg, row1, row2 string
			var wLevel string = "OK"

			//if !state.WaterLevelOk {
			//wLevel = "WLVL"
			//}
			if state.FillingTime > 0 {
				wLevel = "FILL ERR"
				if state.FillingTime < 20000 {
					wLevel = "FILLING"
				}
			}

			//msg += "\x1B" // Reset
			//msg += "\x0E" // Clear
			//msg += "\x13" // Cursor on
			msg += "\x11" // Character over write mode
			//msg += "\x12" // Vertical scroll mode
			//msg += "\x13" // Cursor on
			msg += "\x14" // Cursor off

			msg += "\r"

			row1 = time.Now().Format("15:04:05") + strings.Repeat(" ", 12-len(wLevel)) + wLevel
			row2 = "Temp " + strconv.FormatFloat(state.WaterTemperature, 'f', 2, 64) + " Cool " + strconv.FormatFloat(state.Cooling, 'f', 2, 64)

			if 20-len(row1) > 0 {
				row1 += strings.Repeat(" ", 20-len(row1))
			}

			if 20-len(row2) > 0 {
				row2 += strings.Repeat(" ", 20-len(row2))
			}
			msg += row1[0:20]
			msg += row2[0:20]

			_, err := serial.Port.Write([]byte(msg))
			if err != nil {
				log.Debug("Failed write to display:", err)
			}
		}
	}
} // }}}

// WORKER that monitors the current connection state
func monitorState(connection *basenode.Connection) { // {{{
	for s := range connection.State {
		switch s {
		case basenode.ConnectionStateConnected:
			connection.Send <- node.Node()
		case basenode.ConnectionStateDisconnected:
		}
	}
} // }}}

// WORKER that recives all incomming commands
func serverRecv(connection *basenode.Connection) { // {{{
	for d := range connection.Receive {
		if err := processCommand(d); err != nil {
			log.Error(err)
		}
	}
} // }}}

func processCommand(cmd protocol.Command) error { // {{{
	var target bool

	if len(cmd.Args) < 1 {
		if len(cmd.Params) < 1 {
			return fmt.Errorf("Missing arguments, ignoring command")
		} else {
			target = cmd.Params[0] != "" && cmd.Params[0] != "false"
		}
	} else {
		target = cmd.Args[0] != "" && cmd.Args[0] != "0"
	}

	switch cmd.Cmd {
	case "CirculationPumps":
		if target {
			ard.Port.Write([]byte{0x02, 0x01, 0x01, 0x03}) // Turn on
		} else {
			ard.Port.Write([]byte{0x02, 0x01, 0x00, 0x03}) // Turn off
		}
		break
	case "Skimmer":
		if target {
			ard.Port.Write([]byte{0x02, 0x02, 0x01, 0x03}) // Turn on
		} else {
			ard.Port.Write([]byte{0x02, 0x02, 0x00, 0x03}) // Turn off
		}
		break
	case "Heating":
		if target {
			ard.Port.Write([]byte{0x02, 0x03, 0x01, 0x03}) // Turn on
		} else {
			ard.Port.Write([]byte{0x02, 0x03, 0x00, 0x03}) // Turn off
		}
		break
	case "CoolingP":
		i, err := strconv.Atoi(cmd.Args[0])
		if err != nil {
			return fmt.Errorf("Failed to decode arg[0] to int %s %s", err, cmd.Args[0])
		}

		ard.Port.Write([]byte{0x02, 0x04, byte(i), 0x03}) // Turn on
		break
	case "CoolingI":
		i, err := strconv.Atoi(cmd.Args[0])
		if err != nil {
			return fmt.Errorf("Failed to decode arg[0] to int %s %s", err, cmd.Args[0])
		}

		ard.Port.Write([]byte{0x02, 0x05, byte(i), 0x03}) // Turn on
		break
	case "CoolingD":
		i, err := strconv.Atoi(cmd.Args[0])
		if err != nil {
			return fmt.Errorf("Failed to decode arg[0] to int %s %s", err, cmd.Args[0])
		}

		ard.Port.Write([]byte{0x02, 0x06, byte(i), 0x03}) // Turn on
		break
	}
	return nil
} // }}}

// Serial connection workers
func (config *SerialConnection) run(connection *basenode.Connection, callback func(data string, connection *basenode.Connection)) { // {{{
	go func() {
		for {
			config.connect(connection, callback)
			<-time.After(time.Second * 4)
		}
	}()
}                                                                                // }}}
func (config *SerialConnection) connect(connection *basenode.Connection, callback func(data string, connection *basenode.Connection)) { // {{{

	c := &serial.Config{Name: config.Name, Baud: config.Baud}
	var err error

	config.Port, err = serial.OpenPort(c)
	if err != nil {
		log.Error("Serial connect failed: ", err)
		return
	}

	var incomming string = ""

	for {
		buf := make([]byte, 128)

		n, err := config.Port.Read(buf)
		if err != nil {
			log.Error("Serial read failed: ", err)
			return
		}

		incomming += string(buf[:n])

		for {
			n := strings.Index(incomming, "\r")

			if n < 0 {
				break
			}

			msg := strings.TrimSpace(incomming[:n])
			incomming = incomming[n+1:]
			fmt.Print(msg, "\r")

			go callback(msg, connection)
		}
	}
} // }}}
<|endoftext|>


<commit_before>package structs

import (
	"bytes"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"time"

	// note: this is aliased so that it's more noticeable if someone
	// accidentally swaps it out for math/rand via running goimports
	cryptorand "crypto/rand"

	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/helper/uuid"
)

const (
	// SecureVariablesUpsertRPCMethod is the RPC method for upserting
	// secure variables into Nomad state.
	//
	// Args: SecureVariablesUpsertRequest
	// Reply: SecureVariablesUpsertResponse
	SecureVariablesUpsertRPCMethod = "SecureVariables.Upsert"

	// SecureVariablesDeleteRPCMethod is the RPC method for deleting
	// a secure variable by its namespace and path.
	//
	// Args: SecureVariablesDeleteRequest
	// Reply: SecureVariablesDeleteResponse
	SecureVariablesDeleteRPCMethod = "SecureVariables.Delete"

	// SecureVariablesListRPCMethod is the RPC method for listing secure
	// variables within Nomad.
	//
	// Args: SecureVariablesListRequest
	// Reply: SecureVariablesListResponse
	SecureVariablesListRPCMethod = "SecureVariables.List"

	// SecureVariablesGetServiceRPCMethod is the RPC method for fetching a
	// secure variable according to its namepace and path.
	//
	// Args: SecureVariablesByNameRequest
	// Reply: SecureVariablesByNameResponse
	SecureVariablesReadRPCMethod = "SecureVariables.Read"
)

// SecureVariableMetadata is the metadata envelope for a Secure Variable, it
// is the list object and is shared data between an SecureVariableEncrypted and
// a SecureVariableDecrypted object.
type SecureVariableMetadata struct {
	Namespace   string
	Path        string
	CreateIndex uint64
	CreateTime  int64
	ModifyIndex uint64
	ModifyTime  int64
}

// SecureVariableEncrypted structs are returned from the Encrypter's encrypt
// method. They are the only form that should ever be persisted to storage.
type SecureVariableEncrypted struct {
	SecureVariableMetadata
	SecureVariableData
}

// SecureVariableData is the secret data for a Secure Variable
type SecureVariableData struct {
	Data  []byte // includes nonce
	KeyID string // ID of root key used to encrypt this entry
}

// SecureVariableDecrypted structs are returned from the Encrypter's decrypt
// method. Since they contains sensitive material, they should never be
// persisted to disk.
type SecureVariableDecrypted struct {
	SecureVariableMetadata
	Items SecureVariableItems
}

// SecureVariableItems are the actual secrets stored in a secure variable. They
// are always encrypted and decrypted as a single unit.
type SecureVariableItems map[string]string

// Equals checks both the metadata and items in a SecureVariableDecrypted
// struct
func (v1 SecureVariableDecrypted) Equals(v2 SecureVariableDecrypted) bool {
	return v1.SecureVariableMetadata.Equals(v2.SecureVariableMetadata) &&
		v1.Items.Equals(v2.Items)
}

// Equals is a convenience method to provide similar equality checking
// syntax for metadata and the SecureVariablesData or SecureVariableItems
// struct
func (sv SecureVariableMetadata) Equals(sv2 SecureVariableMetadata) bool {
	return sv == sv2
}

// Equals performs deep equality checking on the cleartext items
// of a SecureVariableDecrypted. Uses reflect.DeepEqual
func (i1 SecureVariableItems) Equals(i2 SecureVariableItems) bool {
	return reflect.DeepEqual(i1, i2)
}

// Equals checks both the metadata and encrypted data for a
// SecureVariableEncrypted struct
func (v1 SecureVariableEncrypted) Equals(v2 SecureVariableEncrypted) bool {
	return v1.SecureVariableMetadata.Equals(v2.SecureVariableMetadata) &&
		v1.SecureVariableData.Equals(v2.SecureVariableData)
}

// Equals performs deep equality checking on the encrypted data part
// of a SecureVariableEncrypted
func (d1 SecureVariableData) Equals(d2 SecureVariableData) bool {
	return d1.KeyID == d2.KeyID &&
		bytes.Equal(d1.Data, d2.Data)
}

func (sv SecureVariableDecrypted) Copy() SecureVariableDecrypted {
	return SecureVariableDecrypted{
		SecureVariableMetadata: sv.SecureVariableMetadata,
		Items:                  sv.Items.Copy(),
	}
}

func (sv SecureVariableItems) Copy() SecureVariableItems {
	out := make(SecureVariableItems, len(sv))
	for k, v := range sv {
		out[k] = v
	}
	return out
}

func (sv SecureVariableEncrypted) Copy() SecureVariableEncrypted {
	return SecureVariableEncrypted{
		SecureVariableMetadata: sv.SecureVariableMetadata,
		SecureVariableData:     sv.SecureVariableData.Copy(),
	}
}

func (sv SecureVariableData) Copy() SecureVariableData {
	out := make([]byte, len(sv.Data))
	copy(out, sv.Data)
	return SecureVariableData{
		Data:  out,
		KeyID: sv.KeyID,
	}
}

func (sv SecureVariableDecrypted) Validate() error {

	if len(sv.Path) == 0 {
		return fmt.Errorf("variable requires path")
	}
	parts := strings.Split(sv.Path, "/")
	switch {
	case len(parts) == 1 && parts[0] == "nomad":
		return fmt.Errorf("\"nomad\" is a reserved top-level directory path, but you may write variables to \"nomad/jobs\" or below")
	case len(parts) >= 2 && parts[0] == "nomad" && parts[1] != "jobs":
		return fmt.Errorf("only paths at \"nomad/jobs\" or below are valid paths under the top-level \"nomad\" directory")
	}

	if len(sv.Items) == 0 {
		return errors.New("empty variables are invalid")
	}
	if sv.Namespace == AllNamespacesSentinel {
		return errors.New("can not target wildcard (\"*\")namespace")
	}
	return nil
}

func (sv *SecureVariableDecrypted) Canonicalize() {
	if sv.Namespace == "" {
		sv.Namespace = DefaultNamespace
	}
}

// GetNamespace returns the secure variable's namespace. Used for pagination.
func (sv *SecureVariableMetadata) Copy() *SecureVariableMetadata {
	var out SecureVariableMetadata = *sv
	return &out
}

// GetNamespace returns the secure variable's namespace. Used for pagination.
func (sv SecureVariableMetadata) GetNamespace() string {
	return sv.Namespace
}

// GetID returns the secure variable's path. Used for pagination.
func (sv SecureVariableMetadata) GetID() string {
	return sv.Path
}

// GetCreateIndex returns the secure variable's create index. Used for pagination.
func (sv SecureVariableMetadata) GetCreateIndex() uint64 {
	return sv.CreateIndex
}

// SecureVariablesQuota is used to track the total size of secure
// variables entries per namespace. The total length of
// SecureVariable.EncryptedData will be added to the SecureVariablesQuota
// table in the same transaction as a write, update, or delete.
type SecureVariablesQuota struct {
	Namespace   string
	Size        uint64
	CreateIndex uint64
	ModifyIndex uint64
}

func (svq *SecureVariablesQuota) Copy() *SecureVariablesQuota {
	if svq == nil {
		return nil
	}
	nq := new(SecureVariablesQuota)
	*nq = *svq
	return nq
}

type SecureVariablesUpsertRequest struct {
	Data       []*SecureVariableDecrypted
	CheckIndex *uint64
	WriteRequest
}

func (svur *SecureVariablesUpsertRequest) SetCheckIndex(ci uint64) {
	svur.CheckIndex = &ci
}

type SecureVariablesEncryptedUpsertRequest struct {
	Data []*SecureVariableEncrypted
	WriteRequest
}

type SecureVariablesUpsertResponse struct {
	Conflicts []*SecureVariableDecrypted
	WriteMeta
}

type SecureVariablesListRequest struct {
	QueryOptions
}

type SecureVariablesListResponse struct {
	Data []*SecureVariableMetadata
	QueryMeta
}

type SecureVariablesReadRequest struct {
	Path string
	QueryOptions
}

type SecureVariablesReadResponse struct {
	Data *SecureVariableDecrypted
	QueryMeta
}

type SecureVariablesDeleteRequest struct {
	Path       string
	CheckIndex *uint64
	WriteRequest
}

func (svdr *SecureVariablesDeleteRequest) SetCheckIndex(ci uint64) {
	svdr.CheckIndex = &ci
}

type SecureVariablesDeleteResponse struct {
	Conflict *SecureVariableDecrypted
	WriteMeta
}

// RootKey is used to encrypt and decrypt secure variables. It is
// never stored in raft.
type RootKey struct {
	Meta *RootKeyMeta
	Key  []byte // serialized to keystore as base64 blob
}

// NewRootKey returns a new root key and its metadata.
func NewRootKey(algorithm EncryptionAlgorithm) (*RootKey, error) {
	meta := NewRootKeyMeta()
	meta.Algorithm = algorithm

	rootKey := &RootKey{
		Meta: meta,
	}

	switch algorithm {
	case EncryptionAlgorithmAES256GCM:
		const keyBytes = 32
		key := make([]byte, keyBytes)
		n, err := cryptorand.Read(key)
		if err != nil {
			return nil, err
		}
		if n < keyBytes {
			return nil, fmt.Errorf("failed to generate key: entropy exhausted")
		}
		rootKey.Key = key
	}

	return rootKey, nil
}

// RootKeyMeta is the metadata used to refer to a RootKey. It is
// stored in raft.
type RootKeyMeta struct {
	KeyID       string // UUID
	Algorithm   EncryptionAlgorithm
	CreateTime  time.Time
	CreateIndex uint64
	ModifyIndex uint64
	State       RootKeyState
}

// RootKeyState enum describes the lifecycle of a root key.
type RootKeyState string

const (
	RootKeyStateInactive   RootKeyState = "inactive"
	RootKeyStateActive                  = "active"
	RootKeyStateRekeying                = "rekeying"
	RootKeyStateDeprecated              = "deprecated"
)

// NewRootKeyMeta returns a new RootKeyMeta with default values
func NewRootKeyMeta() *RootKeyMeta {
	return &RootKeyMeta{
		KeyID:      uuid.Generate(),
		Algorithm:  EncryptionAlgorithmAES256GCM,
		State:      RootKeyStateInactive,
		CreateTime: time.Now(),
	}
}

// RootKeyMetaStub is for serializing root key metadata to the
// keystore, not for the List API. It excludes frequently-changing
// fields such as ModifyIndex so we don't have to sync them to the
// on-disk keystore when the fields are already in raft.
type RootKeyMetaStub struct {
	KeyID      string
	Algorithm  EncryptionAlgorithm
	CreateTime time.Time
	State      RootKeyState
}

// Active indicates his key is the one currently being used for
// crypto operations (at most one key can be Active)
func (rkm *RootKeyMeta) Active() bool {
	return rkm.State == RootKeyStateActive
}

func (rkm *RootKeyMeta) SetActive() {
	rkm.State = RootKeyStateActive
}

// Rekeying indicates that variables encrypted with this key should be
// rekeyed
func (rkm *RootKeyMeta) Rekeying() bool {
	return rkm.State == RootKeyStateRekeying
}

func (rkm *RootKeyMeta) SetRekeying() {
	rkm.State = RootKeyStateRekeying
}

func (rkm *RootKeyMeta) SetInactive() {
	rkm.State = RootKeyStateInactive
}

// Deprecated indicates that variables encrypted with this key
// have been rekeyed
func (rkm *RootKeyMeta) Deprecated() bool {
	return rkm.State == RootKeyStateDeprecated
}

func (rkm *RootKeyMeta) SetDeprecated() {
	rkm.State = RootKeyStateDeprecated
}

func (rkm *RootKeyMeta) Stub() *RootKeyMetaStub {
	if rkm == nil {
		return nil
	}
	return &RootKeyMetaStub{
		KeyID:      rkm.KeyID,
		Algorithm:  rkm.Algorithm,
		CreateTime: rkm.CreateTime,
		State:      rkm.State,
	}

}
func (rkm *RootKeyMeta) Copy() *RootKeyMeta {
	if rkm == nil {
		return nil
	}
	out := *rkm
	return &out
}

func (rkm *RootKeyMeta) Validate() error {
	if rkm == nil {
		return fmt.Errorf("root key metadata is required")
	}
	if rkm.KeyID == "" || !helper.IsUUID(rkm.KeyID) {
		return fmt.Errorf("root key UUID is required")
	}
	if rkm.Algorithm == "" {
		return fmt.Errorf("root key algorithm is required")
	}
	switch rkm.State {
	case RootKeyStateInactive, RootKeyStateActive,
		RootKeyStateRekeying, RootKeyStateDeprecated:
	default:
		return fmt.Errorf("root key state %q is invalid", rkm.State)
	}
	return nil
}

// EncryptionAlgorithm chooses which algorithm is used for
// encrypting / decrypting entries with this key
type EncryptionAlgorithm string

const (
	EncryptionAlgorithmAES256GCM EncryptionAlgorithm = "aes256-gcm"
)

type KeyringRotateRootKeyRequest struct {
	Algorithm EncryptionAlgorithm
	Full      bool
	WriteRequest
}

// KeyringRotateRootKeyResponse returns the full key metadata
type KeyringRotateRootKeyResponse struct {
	Key *RootKeyMeta
	WriteMeta
}

type KeyringListRootKeyMetaRequest struct {
	// TODO: do we need any fields here?
	QueryOptions
}

type KeyringListRootKeyMetaResponse struct {
	Keys []*RootKeyMeta
	QueryMeta
}

// KeyringUpdateRootKeyRequest is used internally for key replication
// only and for keyring restores. The RootKeyMeta will be extracted
// for applying to the FSM with the KeyringUpdateRootKeyMetaRequest
// (see below)
type KeyringUpdateRootKeyRequest struct {
	RootKey *RootKey
	Rekey   bool
	WriteRequest
}

type KeyringUpdateRootKeyResponse struct {
	WriteMeta
}

// KeyringGetRootKeyRequest is used internally for key replication
// only and for keyring restores.
type KeyringGetRootKeyRequest struct {
	KeyID string
	QueryOptions
}

type KeyringGetRootKeyResponse struct {
	Key *RootKey
	QueryMeta
}

// KeyringUpdateRootKeyMetaRequest is used internally for key
// replication so that we have a request wrapper for writing the
// metadata to the FSM without including the key material
type KeyringUpdateRootKeyMetaRequest struct {
	RootKeyMeta *RootKeyMeta
	Rekey       bool
	WriteRequest
}

type KeyringUpdateRootKeyMetaResponse struct {
	WriteMeta
}

type KeyringDeleteRootKeyRequest struct {
	KeyID string
	WriteRequest
}

type KeyringDeleteRootKeyResponse struct {
	WriteMeta
}
<commit_msg>secure vars: limit maximum size of variable data (#13743)<commit_after>package structs

import (
	"bytes"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"time"

	// note: this is aliased so that it's more noticeable if someone
	// accidentally swaps it out for math/rand via running goimports
	cryptorand "crypto/rand"

	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/helper/uuid"
)

const (
	// SecureVariablesUpsertRPCMethod is the RPC method for upserting
	// secure variables into Nomad state.
	//
	// Args: SecureVariablesUpsertRequest
	// Reply: SecureVariablesUpsertResponse
	SecureVariablesUpsertRPCMethod = "SecureVariables.Upsert"

	// SecureVariablesDeleteRPCMethod is the RPC method for deleting
	// a secure variable by its namespace and path.
	//
	// Args: SecureVariablesDeleteRequest
	// Reply: SecureVariablesDeleteResponse
	SecureVariablesDeleteRPCMethod = "SecureVariables.Delete"

	// SecureVariablesListRPCMethod is the RPC method for listing secure
	// variables within Nomad.
	//
	// Args: SecureVariablesListRequest
	// Reply: SecureVariablesListResponse
	SecureVariablesListRPCMethod = "SecureVariables.List"

	// SecureVariablesGetServiceRPCMethod is the RPC method for fetching a
	// secure variable according to its namepace and path.
	//
	// Args: SecureVariablesByNameRequest
	// Reply: SecureVariablesByNameResponse
	SecureVariablesReadRPCMethod = "SecureVariables.Read"

	// maxVariableSize is the maximum size of the unencrypted contents of
	// a variable. This size is deliberately set low and is not
	// configurable, to discourage DoS'ing the cluster
	maxVariableSize = 16384
)

// SecureVariableMetadata is the metadata envelope for a Secure Variable, it
// is the list object and is shared data between an SecureVariableEncrypted and
// a SecureVariableDecrypted object.
type SecureVariableMetadata struct {
	Namespace   string
	Path        string
	CreateIndex uint64
	CreateTime  int64
	ModifyIndex uint64
	ModifyTime  int64
}

// SecureVariableEncrypted structs are returned from the Encrypter's encrypt
// method. They are the only form that should ever be persisted to storage.
type SecureVariableEncrypted struct {
	SecureVariableMetadata
	SecureVariableData
}

// SecureVariableData is the secret data for a Secure Variable
type SecureVariableData struct {
	Data  []byte // includes nonce
	KeyID string // ID of root key used to encrypt this entry
}

// SecureVariableDecrypted structs are returned from the Encrypter's decrypt
// method. Since they contains sensitive material, they should never be
// persisted to disk.
type SecureVariableDecrypted struct {
	SecureVariableMetadata
	Items SecureVariableItems
}

// SecureVariableItems are the actual secrets stored in a secure variable. They
// are always encrypted and decrypted as a single unit.
type SecureVariableItems map[string]string

func (svi SecureVariableItems) Size() uint64 {
	var out uint64
	for k, v := range svi {
		out += uint64(len(k))
		out += uint64(len(v))
	}
	return out
}

// Equals checks both the metadata and items in a SecureVariableDecrypted
// struct
func (v1 SecureVariableDecrypted) Equals(v2 SecureVariableDecrypted) bool {
	return v1.SecureVariableMetadata.Equals(v2.SecureVariableMetadata) &&
		v1.Items.Equals(v2.Items)
}

// Equals is a convenience method to provide similar equality checking
// syntax for metadata and the SecureVariablesData or SecureVariableItems
// struct
func (sv SecureVariableMetadata) Equals(sv2 SecureVariableMetadata) bool {
	return sv == sv2
}

// Equals performs deep equality checking on the cleartext items
// of a SecureVariableDecrypted. Uses reflect.DeepEqual
func (i1 SecureVariableItems) Equals(i2 SecureVariableItems) bool {
	return reflect.DeepEqual(i1, i2)
}

// Equals checks both the metadata and encrypted data for a
// SecureVariableEncrypted struct
func (v1 SecureVariableEncrypted) Equals(v2 SecureVariableEncrypted) bool {
	return v1.SecureVariableMetadata.Equals(v2.SecureVariableMetadata) &&
		v1.SecureVariableData.Equals(v2.SecureVariableData)
}

// Equals performs deep equality checking on the encrypted data part
// of a SecureVariableEncrypted
func (d1 SecureVariableData) Equals(d2 SecureVariableData) bool {
	return d1.KeyID == d2.KeyID &&
		bytes.Equal(d1.Data, d2.Data)
}

func (sv SecureVariableDecrypted) Copy() SecureVariableDecrypted {
	return SecureVariableDecrypted{
		SecureVariableMetadata: sv.SecureVariableMetadata,
		Items:                  sv.Items.Copy(),
	}
}

func (sv SecureVariableItems) Copy() SecureVariableItems {
	out := make(SecureVariableItems, len(sv))
	for k, v := range sv {
		out[k] = v
	}
	return out
}

func (sv SecureVariableEncrypted) Copy() SecureVariableEncrypted {
	return SecureVariableEncrypted{
		SecureVariableMetadata: sv.SecureVariableMetadata,
		SecureVariableData:     sv.SecureVariableData.Copy(),
	}
}

func (sv SecureVariableData) Copy() SecureVariableData {
	out := make([]byte, len(sv.Data))
	copy(out, sv.Data)
	return SecureVariableData{
		Data:  out,
		KeyID: sv.KeyID,
	}
}

func (sv SecureVariableDecrypted) Validate() error {

	if len(sv.Path) == 0 {
		return fmt.Errorf("variable requires path")
	}
	parts := strings.Split(sv.Path, "/")
	switch {
	case len(parts) == 1 && parts[0] == "nomad":
		return fmt.Errorf("\"nomad\" is a reserved top-level directory path, but you may write variables to \"nomad/jobs\" or below")
	case len(parts) >= 2 && parts[0] == "nomad" && parts[1] != "jobs":
		return fmt.Errorf("only paths at \"nomad/jobs\" or below are valid paths under the top-level \"nomad\" directory")
	}

	if len(sv.Items) == 0 {
		return errors.New("empty variables are invalid")
	}
	if sv.Items.Size() > maxVariableSize {
		return errors.New("variables are limited to 16KiB in total size")
	}
	if sv.Namespace == AllNamespacesSentinel {
		return errors.New("can not target wildcard (\"*\")namespace")
	}
	return nil
}

func (sv *SecureVariableDecrypted) Canonicalize() {
	if sv.Namespace == "" {
		sv.Namespace = DefaultNamespace
	}
}

// GetNamespace returns the secure variable's namespace. Used for pagination.
func (sv *SecureVariableMetadata) Copy() *SecureVariableMetadata {
	var out SecureVariableMetadata = *sv
	return &out
}

// GetNamespace returns the secure variable's namespace. Used for pagination.
func (sv SecureVariableMetadata) GetNamespace() string {
	return sv.Namespace
}

// GetID returns the secure variable's path. Used for pagination.
func (sv SecureVariableMetadata) GetID() string {
	return sv.Path
}

// GetCreateIndex returns the secure variable's create index. Used for pagination.
func (sv SecureVariableMetadata) GetCreateIndex() uint64 {
	return sv.CreateIndex
}

// SecureVariablesQuota is used to track the total size of secure
// variables entries per namespace. The total length of
// SecureVariable.EncryptedData will be added to the SecureVariablesQuota
// table in the same transaction as a write, update, or delete.
type SecureVariablesQuota struct {
	Namespace   string
	Size        uint64
	CreateIndex uint64
	ModifyIndex uint64
}

func (svq *SecureVariablesQuota) Copy() *SecureVariablesQuota {
	if svq == nil {
		return nil
	}
	nq := new(SecureVariablesQuota)
	*nq = *svq
	return nq
}

type SecureVariablesUpsertRequest struct {
	Data       []*SecureVariableDecrypted
	CheckIndex *uint64
	WriteRequest
}

func (svur *SecureVariablesUpsertRequest) SetCheckIndex(ci uint64) {
	svur.CheckIndex = &ci
}

type SecureVariablesEncryptedUpsertRequest struct {
	Data []*SecureVariableEncrypted
	WriteRequest
}

type SecureVariablesUpsertResponse struct {
	Conflicts []*SecureVariableDecrypted
	WriteMeta
}

type SecureVariablesListRequest struct {
	QueryOptions
}

type SecureVariablesListResponse struct {
	Data []*SecureVariableMetadata
	QueryMeta
}

type SecureVariablesReadRequest struct {
	Path string
	QueryOptions
}

type SecureVariablesReadResponse struct {
	Data *SecureVariableDecrypted
	QueryMeta
}

type SecureVariablesDeleteRequest struct {
	Path       string
	CheckIndex *uint64
	WriteRequest
}

func (svdr *SecureVariablesDeleteRequest) SetCheckIndex(ci uint64) {
	svdr.CheckIndex = &ci
}

type SecureVariablesDeleteResponse struct {
	Conflict *SecureVariableDecrypted
	WriteMeta
}

// RootKey is used to encrypt and decrypt secure variables. It is
// never stored in raft.
type RootKey struct {
	Meta *RootKeyMeta
	Key  []byte // serialized to keystore as base64 blob
}

// NewRootKey returns a new root key and its metadata.
func NewRootKey(algorithm EncryptionAlgorithm) (*RootKey, error) {
	meta := NewRootKeyMeta()
	meta.Algorithm = algorithm

	rootKey := &RootKey{
		Meta: meta,
	}

	switch algorithm {
	case EncryptionAlgorithmAES256GCM:
		const keyBytes = 32
		key := make([]byte, keyBytes)
		n, err := cryptorand.Read(key)
		if err != nil {
			return nil, err
		}
		if n < keyBytes {
			return nil, fmt.Errorf("failed to generate key: entropy exhausted")
		}
		rootKey.Key = key
	}

	return rootKey, nil
}

// RootKeyMeta is the metadata used to refer to a RootKey. It is
// stored in raft.
type RootKeyMeta struct {
	KeyID       string // UUID
	Algorithm   EncryptionAlgorithm
	CreateTime  time.Time
	CreateIndex uint64
	ModifyIndex uint64
	State       RootKeyState
}

// RootKeyState enum describes the lifecycle of a root key.
type RootKeyState string

const (
	RootKeyStateInactive   RootKeyState = "inactive"
	RootKeyStateActive                  = "active"
	RootKeyStateRekeying                = "rekeying"
	RootKeyStateDeprecated              = "deprecated"
)

// NewRootKeyMeta returns a new RootKeyMeta with default values
func NewRootKeyMeta() *RootKeyMeta {
	return &RootKeyMeta{
		KeyID:      uuid.Generate(),
		Algorithm:  EncryptionAlgorithmAES256GCM,
		State:      RootKeyStateInactive,
		CreateTime: time.Now(),
	}
}

// RootKeyMetaStub is for serializing root key metadata to the
// keystore, not for the List API. It excludes frequently-changing
// fields such as ModifyIndex so we don't have to sync them to the
// on-disk keystore when the fields are already in raft.
type RootKeyMetaStub struct {
	KeyID      string
	Algorithm  EncryptionAlgorithm
	CreateTime time.Time
	State      RootKeyState
}

// Active indicates his key is the one currently being used for
// crypto operations (at most one key can be Active)
func (rkm *RootKeyMeta) Active() bool {
	return rkm.State == RootKeyStateActive
}

func (rkm *RootKeyMeta) SetActive() {
	rkm.State = RootKeyStateActive
}

// Rekeying indicates that variables encrypted with this key should be
// rekeyed
func (rkm *RootKeyMeta) Rekeying() bool {
	return rkm.State == RootKeyStateRekeying
}

func (rkm *RootKeyMeta) SetRekeying() {
	rkm.State = RootKeyStateRekeying
}

func (rkm *RootKeyMeta) SetInactive() {
	rkm.State = RootKeyStateInactive
}

// Deprecated indicates that variables encrypted with this key
// have been rekeyed
func (rkm *RootKeyMeta) Deprecated() bool {
	return rkm.State == RootKeyStateDeprecated
}

func (rkm *RootKeyMeta) SetDeprecated() {
	rkm.State = RootKeyStateDeprecated
}

func (rkm *RootKeyMeta) Stub() *RootKeyMetaStub {
	if rkm == nil {
		return nil
	}
	return &RootKeyMetaStub{
		KeyID:      rkm.KeyID,
		Algorithm:  rkm.Algorithm,
		CreateTime: rkm.CreateTime,
		State:      rkm.State,
	}

}
func (rkm *RootKeyMeta) Copy() *RootKeyMeta {
	if rkm == nil {
		return nil
	}
	out := *rkm
	return &out
}

func (rkm *RootKeyMeta) Validate() error {
	if rkm == nil {
		return fmt.Errorf("root key metadata is required")
	}
	if rkm.KeyID == "" || !helper.IsUUID(rkm.KeyID) {
		return fmt.Errorf("root key UUID is required")
	}
	if rkm.Algorithm == "" {
		return fmt.Errorf("root key algorithm is required")
	}
	switch rkm.State {
	case RootKeyStateInactive, RootKeyStateActive,
		RootKeyStateRekeying, RootKeyStateDeprecated:
	default:
		return fmt.Errorf("root key state %q is invalid", rkm.State)
	}
	return nil
}

// EncryptionAlgorithm chooses which algorithm is used for
// encrypting / decrypting entries with this key
type EncryptionAlgorithm string

const (
	EncryptionAlgorithmAES256GCM EncryptionAlgorithm = "aes256-gcm"
)

type KeyringRotateRootKeyRequest struct {
	Algorithm EncryptionAlgorithm
	Full      bool
	WriteRequest
}

// KeyringRotateRootKeyResponse returns the full key metadata
type KeyringRotateRootKeyResponse struct {
	Key *RootKeyMeta
	WriteMeta
}

type KeyringListRootKeyMetaRequest struct {
	// TODO: do we need any fields here?
	QueryOptions
}

type KeyringListRootKeyMetaResponse struct {
	Keys []*RootKeyMeta
	QueryMeta
}

// KeyringUpdateRootKeyRequest is used internally for key replication
// only and for keyring restores. The RootKeyMeta will be extracted
// for applying to the FSM with the KeyringUpdateRootKeyMetaRequest
// (see below)
type KeyringUpdateRootKeyRequest struct {
	RootKey *RootKey
	Rekey   bool
	WriteRequest
}

type KeyringUpdateRootKeyResponse struct {
	WriteMeta
}

// KeyringGetRootKeyRequest is used internally for key replication
// only and for keyring restores.
type KeyringGetRootKeyRequest struct {
	KeyID string
	QueryOptions
}

type KeyringGetRootKeyResponse struct {
	Key *RootKey
	QueryMeta
}

// KeyringUpdateRootKeyMetaRequest is used internally for key
// replication so that we have a request wrapper for writing the
// metadata to the FSM without including the key material
type KeyringUpdateRootKeyMetaRequest struct {
	RootKeyMeta *RootKeyMeta
	Rekey       bool
	WriteRequest
}

type KeyringUpdateRootKeyMetaResponse struct {
	WriteMeta
}

type KeyringDeleteRootKeyRequest struct {
	KeyID string
	WriteRequest
}

type KeyringDeleteRootKeyResponse struct {
	WriteMeta
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package cli

import (
	"fmt"

	"github.com/cockroachdb/cockroach/security"

	"github.com/spf13/cobra"
)

const defaultKeySize = 2048

var keySize int

// A createCACert command generates a CA certificate and stores it
// in the cert directory.
var createCACertCmd = &cobra.Command{
	Use:   "create-ca [options]",
	Short: "create CA cert and key",
	Long: `
Generates a new key pair and CA certificate, writing them to
individual files in the directory specified by --certs (required).
`,
	Run: runCreateCACert,
}

// runCreateCACert generates key pair and CA certificate and writes them
// to their corresponding files.
func runCreateCACert(cmd *cobra.Command, args []string) {
	err := security.RunCreateCACert(Context.Certs, keySize)
	if err != nil {
		fmt.Fprintf(osStderr, "failed to generate CA certificate: %s\n", err)
		osExit(1)
		return
	}
}

// A createNodeCert command generates a node certificate and stores it
// in the cert directory.
var createNodeCertCmd = &cobra.Command{
	Use:   "create-node [options] <host 1> <host 2> ... <host N>",
	Short: "create node cert and key",
	Long: `
Generates server and client certificates and keys for a given node, writing them to
individual files in the directory specified by --certs (required).
The certs directory should contain a CA cert and key.
At least one host should be passed in (either IP address of dns name).
`,
	Run: runCreateNodeCert,
}

// runCreateNodeCert generates key pair and CA certificate and writes them
// to their corresponding files.
func runCreateNodeCert(cmd *cobra.Command, args []string) {
	err := security.RunCreateNodeCert(Context.Certs, keySize, args)
	if err != nil {
		fmt.Fprintf(osStderr, "failed to generate node certificate: %s\n", err)
		osExit(1)
		return
	}
}

// A createClientCert command generates a client certificate and stores it
// in the cert directory under <username>.crt and key under <username>.key.
var createClientCertCmd = &cobra.Command{
	Use:   "create-client [options] username",
	Short: "create client cert and key",
	Long: `
Generates a new key pair and client certificate, writing them to
individual files in the directory specified by --certs (required).
The certs directory should contain a CA cert and key.
`,
	Run: runCreateClientCert,
}

// runCreateClientCert generates key pair and CA certificate and writes them
// to their corresponding files.
func runCreateClientCert(cmd *cobra.Command, args []string) {
	err := security.RunCreateClientCert(Context.Certs, keySize, args[0])
	if err != nil {
		fmt.Fprintf(osStderr, "failed to generate clent certificate: %s\n", err)
		osExit(1)
		return
	}
}

var certCmds = []*cobra.Command{
	createCACertCmd,
	createNodeCertCmd,
	createClientCertCmd,
}

var certCmd = &cobra.Command{
	Use:   "cert",
	Short: "create ca, node, and client certs",
	Run: func(cmd *cobra.Command, args []string) {
		cmd.Usage()
	},
}

func init() {
	certCmd.AddCommand(certCmds...)
}
<commit_msg>Fix NPE in cert generation command<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package cli

import (
	"fmt"

	"github.com/cockroachdb/cockroach/security"

	"github.com/spf13/cobra"
)

const defaultKeySize = 2048

var keySize int

// A createCACert command generates a CA certificate and stores it
// in the cert directory.
var createCACertCmd = &cobra.Command{
	Use:   "create-ca [options]",
	Short: "create CA cert and key",
	Long: `
Generates a new key pair and CA certificate, writing them to
individual files in the directory specified by --certs (required).
`,
	Run: runCreateCACert,
}

// runCreateCACert generates key pair and CA certificate and writes them
// to their corresponding files.
func runCreateCACert(cmd *cobra.Command, args []string) {
	err := security.RunCreateCACert(Context.Certs, keySize)
	if err != nil {
		fmt.Fprintf(osStderr, "failed to generate CA certificate: %s\n", err)
		osExit(1)
		return
	}
}

// A createNodeCert command generates a node certificate and stores it
// in the cert directory.
var createNodeCertCmd = &cobra.Command{
	Use:   "create-node [options] <host 1> <host 2> ... <host N>",
	Short: "create node cert and key",
	Long: `
Generates server and client certificates and keys for a given node, writing them to
individual files in the directory specified by --certs (required).
The certs directory should contain a CA cert and key.
At least one host should be passed in (either IP address of dns name).
`,
	Run: runCreateNodeCert,
}

// runCreateNodeCert generates key pair and CA certificate and writes them
// to their corresponding files.
func runCreateNodeCert(cmd *cobra.Command, args []string) {
	err := security.RunCreateNodeCert(Context.Certs, keySize, args)
	if err != nil {
		fmt.Fprintf(osStderr, "failed to generate node certificate: %s\n", err)
		osExit(1)
		return
	}
}

// A createClientCert command generates a client certificate and stores it
// in the cert directory under <username>.crt and key under <username>.key.
var createClientCertCmd = &cobra.Command{
	Use:   "create-client [options] username",
	Short: "create client cert and key",
	Long: `
Generates a new key pair and client certificate, writing them to
individual files in the directory specified by --certs (required).
The certs directory should contain a CA cert and key.
`,
	Run: runCreateClientCert,
}

// runCreateClientCert generates key pair and CA certificate and writes them
// to their corresponding files.
func runCreateClientCert(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		cmd.Usage()
		return
	}
	err := security.RunCreateClientCert(Context.Certs, keySize, args[0])
	if err != nil {
		fmt.Fprintf(osStderr, "failed to generate clent certificate: %s\n", err)
		osExit(1)
		return
	}
}

var certCmds = []*cobra.Command{
	createCACertCmd,
	createNodeCertCmd,
	createClientCertCmd,
}

var certCmd = &cobra.Command{
	Use:   "cert",
	Short: "create ca, node, and client certs",
	Run: func(cmd *cobra.Command, args []string) {
		cmd.Usage()
	},
}

func init() {
	certCmd.AddCommand(certCmds...)
}
<|endoftext|>


<commit_before>// Package observer provides a facility for publishing progress updates and
// state changes from parts of the daemon, an a SSE http handler for consumers
// of these events.
package observer

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
)

type ctxkey string

// CtxRequestID is the context WithValue key for a request id.
var CtxRequestID ctxkey = "id"

// EventType represents all possible types of events that can be observed.
type EventType string

// All values for EventType
const (
	Progress EventType = "progress"
	Started  EventType = "started"
	Finished EventType = "finished"
	Errored  EventType = "errored"
	Aborted  EventType = "aborted"
)

type event struct {
	ID        string    `json:"id"`
	Type      EventType `json:"-"` // type is included in the SSE
	Message   string    `json:"message"`
	Completed uint      `json:"completed"`
	Total     uint      `json:"total"`
}

type notification struct {
	Type      EventType
	Message   string
	Increment bool
}

// Observer receives events via Notify, and publishes them as SSEs via its
// ServeHTTP function.
type Observer struct {
	notify chan *event
	closed chan int

	observers map[chan []byte]bool

	newObservers    chan chan []byte
	closedObservers chan chan []byte
}

type transaction struct {
	requestID      string
	total          uint
	current        uint
	events         chan<- *event
	totalUpdates   <-chan uint
	notifications  <-chan *notification
	observerClosed <-chan int
	ctxDone        <-chan struct{}
}

// Notifier belongs to a transactions and represents one segment in a series of
// actions. A Notifier can send many messages.
type Notifier struct {
	total          uint
	current        uint
	transaction    *transaction
	totalUpdates   chan<- uint
	notifications  chan<- *notification
	observerClosed <-chan int
	ctxDone        <-chan struct{}
}

func (t *transaction) start() {
	go func() {
		for {
			select {
			case notification := <-t.notifications:
				if notification.Increment {
					t.current++
				}

				evt := &event{
					ID:        t.requestID,
					Type:      notification.Type,
					Message:   notification.Message,
					Completed: t.current,
					Total:     t.total,
				}

				t.events <- evt
			case size := <-t.totalUpdates:
				t.total += size
			case <-t.ctxDone: // transaction has completed
				return
			case <-t.observerClosed: // observer has shutdown
				return
			}
		}
	}()
}

// Notifier creates a child notifier to this Notifier
func (n *Notifier) Notifier(total uint) *Notifier {
	notifier := &Notifier{
		total:          total,
		current:        0,
		transaction:    nil,
		totalUpdates:   n.totalUpdates,
		notifications:  n.notifications,
		observerClosed: n.observerClosed,
		ctxDone:        n.ctxDone,
	}

	n.totalUpdates <- total
	return notifier
}

// Notify publishes an event to all SSE observers.
func (n *Notifier) Notify(eventType EventType, message string, increment bool) {
	notif := &notification{
		Type:      eventType,
		Message:   message,
		Increment: increment,
	}

	if increment {
		n.current++
	}

	if n.current > n.total {
		panic(fmt.Sprintf(
			"notifications exceed maximum %d/%d", n.current, n.total))
	}

	select {
	case n.notifications <- notif:
		return
	case <-n.observerClosed:
		return
	case <-n.ctxDone:
		return
	}
}

// New returns a new initialized Observer
func New() *Observer {
	return &Observer{
		notify: make(chan *event),
		closed: make(chan int),

		observers: make(map[chan []byte]bool),

		newObservers:    make(chan chan []byte),
		closedObservers: make(chan chan []byte),
	}
}

// Notifier creates a new transaction for sending notifications
func (o *Observer) Notifier(ctx context.Context, total uint) (*Notifier, error) {

	if ctx == nil {
		return nil, errors.New("Context must be provided")
	}

	id, ok := ctx.Value(CtxRequestID).(string)
	if !ok {
		return nil, errors.New("Missing 'id' property in Context")
	}

	totalUpdates := make(chan uint)
	notifications := make(chan *notification)

	t := &transaction{
		requestID:      id,
		total:          total,
		current:        0,
		totalUpdates:   totalUpdates,
		events:         o.notify,
		notifications:  notifications,
		observerClosed: o.closed,
		ctxDone:        ctx.Done(),
	}

	t.start()

	n := &Notifier{
		total:          total,
		current:        0,
		transaction:    t,
		totalUpdates:   totalUpdates,
		notifications:  notifications,
		observerClosed: t.observerClosed,
		ctxDone:        t.ctxDone,
	}

	return n, nil
}

// ServeHTTP implements the http.Handler interface for providing server-sent
// events of observed notifications.
func (o *Observer) ServeHTTP(rw http.ResponseWriter, r *http.Request) {

	rwf := rw.(http.Flusher)
	closed := rw.(http.CloseNotifier).CloseNotify()

	notify := make(chan []byte)
	o.newObservers <- notify

	defer func() {
		o.closedObservers <- notify
	}()

	rw.Header().Set("Content-Type", "text/event-stream")
	rw.Header().Set("Connection", "keep-alive")
	rwf.Flush()

	for {
		select {
		case evt := <-notify:
			// Write the event to the client. we ignore errors here, and
			// let the close channel tell us when to stop writing.
			rw.Write(evt)
			rwf.Flush()

		case <-closed: // client has disconnected. Exit.
			return
		case <-o.closed: // The Observer is shutting down. Exit.
			return
		}
	}
}

// Start begins listening for notifications of observable events. It returns
// after stop has been called.
func (o *Observer) Start() {
	for {
		select {
		case evt := <-o.notify: // We have an event to observe
			if len(o.observers) == 0 {
				continue
			}

			evtb, err := json.Marshal(evt)
			if err != nil {
				log.Printf("Error marshaling event: %s", err)
				continue
			}

			sse := []byte("event: " + evt.Type + "\ndata: ")
			sse = append(sse, append(evtb, []byte("\n\n")...)...)

			for n := range o.observers {
				n <- sse
			}

		case n := <-o.newObservers:
			o.observers[n] = true
		case n := <-o.closedObservers:
			delete(o.observers, n)

		case <-o.closed: // The Observer has been closed.
			return
		}
	}
}

// Stop terminates propagation of events through the observer
func (o *Observer) Stop() {
	close(o.closed)
}
<commit_msg>Observer: put a Read/Write mutex in place.<commit_after>// Package observer provides a facility for publishing progress updates and
// state changes from parts of the daemon, an a SSE http handler for consumers
// of these events.
package observer

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"sync"
)

type ctxkey string

// CtxRequestID is the context WithValue key for a request id.
var CtxRequestID ctxkey = "id"

// EventType represents all possible types of events that can be observed.
type EventType string

// All values for EventType
const (
	Progress EventType = "progress"
	Started  EventType = "started"
	Finished EventType = "finished"
	Errored  EventType = "errored"
	Aborted  EventType = "aborted"
)

type event struct {
	ID        string    `json:"id"`
	Type      EventType `json:"-"` // type is included in the SSE
	Message   string    `json:"message"`
	Completed uint      `json:"completed"`
	Total     uint      `json:"total"`
}

type notification struct {
	Type      EventType
	Message   string
	Increment bool
}

// Observer receives events via Notify, and publishes them as SSEs via its
// ServeHTTP function.
type Observer struct {
	notify chan *event
	closed chan int

	observers map[chan []byte]bool

	newObservers    chan chan []byte
	closedObservers chan chan []byte
}

type transaction struct {
	requestID      string
	total          uint
	current        uint
	events         chan<- *event
	totalUpdates   <-chan uint
	notifications  <-chan *notification
	observerClosed <-chan int
	ctxDone        <-chan struct{}
}

// Notifier belongs to a transactions and represents one segment in a series of
// actions. A Notifier can send many messages.
type Notifier struct {
	total          uint
	current        uint
	transaction    *transaction
	totalUpdates   chan<- uint
	notifications  chan<- *notification
	observerClosed <-chan int
	ctxDone        <-chan struct{}
	sync.RWMutex
}

func (t *transaction) start() {
	go func() {
		for {
			select {
			case notification := <-t.notifications:
				if notification.Increment {
					t.current++
				}

				evt := &event{
					ID:        t.requestID,
					Type:      notification.Type,
					Message:   notification.Message,
					Completed: t.current,
					Total:     t.total,
				}

				t.events <- evt
			case size := <-t.totalUpdates:
				t.total += size
			case <-t.ctxDone: // transaction has completed
				return
			case <-t.observerClosed: // observer has shutdown
				return
			}
		}
	}()
}

// Notifier creates a child notifier to this Notifier
func (n *Notifier) Notifier(total uint) *Notifier {
	notifier := &Notifier{
		total:          total,
		current:        0,
		transaction:    nil,
		totalUpdates:   n.totalUpdates,
		notifications:  n.notifications,
		observerClosed: n.observerClosed,
		ctxDone:        n.ctxDone,
	}

	n.totalUpdates <- total
	return notifier
}

// Notify publishes an event to all SSE observers.
func (n *Notifier) Notify(eventType EventType, message string, increment bool) {
	notif := &notification{
		Type:      eventType,
		Message:   message,
		Increment: increment,
	}

	if increment {
		n.Lock()
		n.current++
		n.Unlock()
	}

	n.RLock()
	if n.current > n.total {
		panic(fmt.Sprintf(
			"notifications exceed maximum %d/%d", n.current, n.total))
	}
	n.RUnlock()

	select {
	case n.notifications <- notif:
		return
	case <-n.observerClosed:
		return
	case <-n.ctxDone:
		return
	}
}

// New returns a new initialized Observer
func New() *Observer {
	return &Observer{
		notify: make(chan *event),
		closed: make(chan int),

		observers: make(map[chan []byte]bool),

		newObservers:    make(chan chan []byte),
		closedObservers: make(chan chan []byte),
	}
}

// Notifier creates a new transaction for sending notifications
func (o *Observer) Notifier(ctx context.Context, total uint) (*Notifier, error) {

	if ctx == nil {
		return nil, errors.New("Context must be provided")
	}

	id, ok := ctx.Value(CtxRequestID).(string)
	if !ok {
		return nil, errors.New("Missing 'id' property in Context")
	}

	totalUpdates := make(chan uint)
	notifications := make(chan *notification)

	t := &transaction{
		requestID:      id,
		total:          total,
		current:        0,
		totalUpdates:   totalUpdates,
		events:         o.notify,
		notifications:  notifications,
		observerClosed: o.closed,
		ctxDone:        ctx.Done(),
	}

	t.start()

	n := &Notifier{
		total:          total,
		current:        0,
		transaction:    t,
		totalUpdates:   totalUpdates,
		notifications:  notifications,
		observerClosed: t.observerClosed,
		ctxDone:        t.ctxDone,
	}

	return n, nil
}

// ServeHTTP implements the http.Handler interface for providing server-sent
// events of observed notifications.
func (o *Observer) ServeHTTP(rw http.ResponseWriter, r *http.Request) {

	rwf := rw.(http.Flusher)
	closed := rw.(http.CloseNotifier).CloseNotify()

	notify := make(chan []byte)
	o.newObservers <- notify

	defer func() {
		o.closedObservers <- notify
	}()

	rw.Header().Set("Content-Type", "text/event-stream")
	rw.Header().Set("Connection", "keep-alive")
	rwf.Flush()

	for {
		select {
		case evt := <-notify:
			// Write the event to the client. we ignore errors here, and
			// let the close channel tell us when to stop writing.
			rw.Write(evt)
			rwf.Flush()

		case <-closed: // client has disconnected. Exit.
			return
		case <-o.closed: // The Observer is shutting down. Exit.
			return
		}
	}
}

// Start begins listening for notifications of observable events. It returns
// after stop has been called.
func (o *Observer) Start() {
	for {
		select {
		case evt := <-o.notify: // We have an event to observe
			if len(o.observers) == 0 {
				continue
			}

			evtb, err := json.Marshal(evt)
			if err != nil {
				log.Printf("Error marshaling event: %s", err)
				continue
			}

			sse := []byte("event: " + evt.Type + "\ndata: ")
			sse = append(sse, append(evtb, []byte("\n\n")...)...)

			for n := range o.observers {
				n <- sse
			}

		case n := <-o.newObservers:
			o.observers[n] = true
		case n := <-o.closedObservers:
			delete(o.observers, n)

		case <-o.closed: // The Observer has been closed.
			return
		}
	}
}

// Stop terminates propagation of events through the observer
func (o *Observer) Stop() {
	close(o.closed)
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"os"
	"strings"
	"time"

	log "github.com/cihub/seelog"
	"github.com/golang/protobuf/proto"
	"github.com/streadway/amqp"

	"github.com/b2aio/typhon/errors"
	"github.com/b2aio/typhon/rabbit"
)

type AMQPServer struct {
	// this is the routing key prefix for all endpoints
	ServiceName        string
	ServiceDescription string
	endpointRegistry   *EndpointRegistry
	connection         *rabbit.RabbitConnection
	notifyConnected    []chan bool

	closeChan chan struct{}
}

func NewAMQPServer() Server {
	return &AMQPServer{
		endpointRegistry: NewEndpointRegistry(),
		connection:       rabbit.NewRabbitConnection(),
		closeChan:        make(chan struct{}),
	}
}

func (s *AMQPServer) Name() string {
	if s == nil {
		return ""
	}
	return s.ServiceName
}

func (s *AMQPServer) Description() string {
	if s == nil {
		return ""
	}
	return s.ServiceDescription
}

func (s *AMQPServer) Init(c *Config) {
	s.ServiceName = c.Name
	s.ServiceDescription = c.Description
}

func (s *AMQPServer) NotifyConnected() chan bool {
	ch := make(chan bool)
	s.notifyConnected = append(s.notifyConnected, ch)
	return ch
}

func (s *AMQPServer) RegisterEndpoint(endpoint Endpoint) {
	s.endpointRegistry.Register(endpoint)
}

func (s *AMQPServer) DeregisterEndpoint(endpointName string) {
	s.endpointRegistry.Deregister(endpointName)
}

// Run the server, connecting to our transport and serving requests
func (s *AMQPServer) Run() {

	// Connect to AMQP
	select {
	case <-s.connection.Init():
		log.Info("[Server] Connected to RabbitMQ")
		for _, notify := range s.notifyConnected {
			notify <- true
		}
	case <-time.After(10 * time.Second):
		log.Critical("[Server] Failed to connect to RabbitMQ")
		os.Exit(1)
	}

	// Get a delivery channel from the connection
	log.Infof("[Server] Listening for deliveries on %s.#", s.ServiceName)
	deliveries, err := s.connection.Consume(s.ServiceName)
	if err != nil {
		log.Criticalf("[Server] [%s] Failed to consume from Rabbit", s.ServiceName)
	}

	// Handle deliveries
	for {
		select {
		case req := <-deliveries:
			log.Infof("[Server] [%s] Received new delivery", s.ServiceName)
			go s.handleRequest(req)
		case <-s.closeChan:
			// shut down server
			s.connection.Close()
			break
		}
	}

	log.Infof("Exiting")
	log.Flush()
}

func (s *AMQPServer) Close() {
	close(s.closeChan)
}

// handleRequest takes a delivery from AMQP, attempts to process it and return a response
func (s *AMQPServer) handleRequest(delivery amqp.Delivery) {

	log.Infof("Received delivery for %+v", delivery)

	// See if we have a matching endpoint for this request
	endpointName := strings.Replace(delivery.RoutingKey, fmt.Sprintf("%s.", s.ServiceName), "", -1)
	endpoint := s.endpointRegistry.Get(endpointName)
	if endpoint == nil {
		log.Errorf("[Server] Endpoint '%s' not found, cannot handle request", endpointName)
		s.respondWithError(delivery, errors.BadRequest("endpoint.notfound", "Endpoint not found"))
		return
	}

	// Handle the delivery
	req := NewAMQPRequest(&delivery)
	rsp, err := endpoint.HandleRequest(req)
	if err != nil {
		s.respondWithError(delivery, err)
		return
	}

	// TODO deal with rsp == nil (programmer error, but still)

	// Marshal the response
	body, err := rsp.Encode()
	if err != nil {
		log.Errorf("[Server] Failed to marshal response")
		s.respondWithError(delivery, errors.BadResponse("response.marshal", err.Error()))
		return
	}

	// Build return delivery, and publish
	msg := amqp.Publishing{
		CorrelationId: delivery.CorrelationId,
		Timestamp:     time.Now().UTC(),
		Body:          body,
		Headers: map[string]interface{}{
			"Content-Encoding": "RESPONSE",
		},
	}
	s.connection.Publish("", delivery.ReplyTo, msg)
}

// respondWithError to a delivery, with the provided error
func (s *AMQPServer) respondWithError(delivery amqp.Delivery, err error) {

	// Construct a return message with an error
	msg := amqp.Publishing{
		CorrelationId: delivery.CorrelationId,
		Timestamp:     time.Now().UTC(),
		Body:          []byte(err.Error()),
	}

	// Publish the error back to the client
	s.connection.Publish("", delivery.ReplyTo, msg)
}
<commit_msg>Marshal service errors to proto, and set content encoding<commit_after>package server

import (
	"fmt"
	"os"
	"strings"
	"time"

	log "github.com/cihub/seelog"
	"github.com/golang/protobuf/proto"
	"github.com/streadway/amqp"

	"github.com/b2aio/typhon/errors"
	"github.com/b2aio/typhon/rabbit"
)

type AMQPServer struct {
	// this is the routing key prefix for all endpoints
	ServiceName        string
	ServiceDescription string
	endpointRegistry   *EndpointRegistry
	connection         *rabbit.RabbitConnection
	notifyConnected    []chan bool

	closeChan chan struct{}
}

func NewAMQPServer() Server {
	return &AMQPServer{
		endpointRegistry: NewEndpointRegistry(),
		connection:       rabbit.NewRabbitConnection(),
		closeChan:        make(chan struct{}),
	}
}

func (s *AMQPServer) Name() string {
	if s == nil {
		return ""
	}
	return s.ServiceName
}

func (s *AMQPServer) Description() string {
	if s == nil {
		return ""
	}
	return s.ServiceDescription
}

func (s *AMQPServer) Init(c *Config) {
	s.ServiceName = c.Name
	s.ServiceDescription = c.Description
}

func (s *AMQPServer) NotifyConnected() chan bool {
	ch := make(chan bool)
	s.notifyConnected = append(s.notifyConnected, ch)
	return ch
}

func (s *AMQPServer) RegisterEndpoint(endpoint Endpoint) {
	s.endpointRegistry.Register(endpoint)
}

func (s *AMQPServer) DeregisterEndpoint(endpointName string) {
	s.endpointRegistry.Deregister(endpointName)
}

// Run the server, connecting to our transport and serving requests
func (s *AMQPServer) Run() {

	// Connect to AMQP
	select {
	case <-s.connection.Init():
		log.Info("[Server] Connected to RabbitMQ")
		for _, notify := range s.notifyConnected {
			notify <- true
		}
	case <-time.After(10 * time.Second):
		log.Critical("[Server] Failed to connect to RabbitMQ")
		os.Exit(1)
	}

	// Get a delivery channel from the connection
	log.Infof("[Server] Listening for deliveries on %s.#", s.ServiceName)
	deliveries, err := s.connection.Consume(s.ServiceName)
	if err != nil {
		log.Criticalf("[Server] [%s] Failed to consume from Rabbit", s.ServiceName)
	}

	// Handle deliveries
	for {
		select {
		case req := <-deliveries:
			log.Infof("[Server] [%s] Received new delivery", s.ServiceName)
			go s.handleRequest(req)
		case <-s.closeChan:
			// shut down server
			s.connection.Close()
			break
		}
	}

	log.Infof("Exiting")
	log.Flush()
}

func (s *AMQPServer) Close() {
	close(s.closeChan)
}

// handleRequest takes a delivery from AMQP, attempts to process it and return a response
func (s *AMQPServer) handleRequest(delivery amqp.Delivery) {

	log.Infof("Received delivery for %+v", delivery)

	// See if we have a matching endpoint for this request
	endpointName := strings.Replace(delivery.RoutingKey, fmt.Sprintf("%s.", s.ServiceName), "", -1)
	endpoint := s.endpointRegistry.Get(endpointName)
	if endpoint == nil {
		log.Errorf("[Server] Endpoint '%s' not found, cannot handle request", endpointName)
		s.respondWithError(delivery, errors.BadRequest("endpoint.notfound", "Endpoint not found"))
		return
	}

	// Handle the delivery
	req := NewAMQPRequest(&delivery)
	rsp, err := endpoint.HandleRequest(req)
	if err != nil {
		s.respondWithError(delivery, err)
		return
	}

	// TODO deal with rsp == nil (programmer error, but still)

	// Marshal the response
	body, err := rsp.Encode()
	if err != nil {
		log.Errorf("[Server] Failed to marshal response")
		s.respondWithError(delivery, errors.BadResponse("response.marshal", err.Error()))
		return
	}

	// Build return delivery, and publish
	msg := amqp.Publishing{
		CorrelationId: delivery.CorrelationId,
		Timestamp:     time.Now().UTC(),
		Body:          body,
		Headers: map[string]interface{}{
			"Content-Encoding": "RESPONSE",
		},
	}
	s.connection.Publish("", delivery.ReplyTo, msg)
}

// respondWithError to a delivery, with the provided error
func (s *AMQPServer) respondWithError(delivery amqp.Delivery, err error) {

	// Ensure we have a service error in proto form
	// and marshal this for transmission
	svcErr := errorToServiceError(err)
	b, err := proto.Marshal(errors.Marshal(svcErr))
	if err != nil {
		// shit
	}

	// Construct a return message with an error
	msg := amqp.Publishing{
		CorrelationId: delivery.CorrelationId,
		Timestamp:     time.Now().UTC(),
		Body:          b,
		Headers: map[string]interface{}{
			"Content-Encoding": "ERROR",
		},
	}

	// Publish the error back to the client
	s.connection.Publish("", delivery.ReplyTo, msg)
}

// errorToServiceError converts an error interface to the concrete
// errors.ServiceError type which we pass between services (as proto)
func errorToServiceError(err error) *errors.ServiceError {

	// If we already have a service error, return this
	if svcErr, ok := err.(*errors.ServiceError); ok {
		return svcErr
	}

	// Otherwise make a generic internal service error
	// Encapsulating the error we were given
	return errors.InternalService("internalservice", err.Error()).(*errors.ServiceError)
}
<|endoftext|>


<commit_before>// Package service provides a micro service
package service

import (
	"os"
	"strings"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/proxy"
	"github.com/micro/go-micro/proxy/grpc"
	"github.com/micro/go-micro/proxy/http"
	"github.com/micro/go-micro/proxy/mucp"
	"github.com/micro/go-micro/runtime"
	"github.com/micro/go-micro/server"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/go-micro/util/mux"
)

func run(ctx *cli.Context, opts ...micro.Option) {
	log.Name("service")

	name := ctx.String("name")
	address := ctx.String("address")
	endpoint := ctx.String("endpoint")

	if len(name) > 0 {
		opts = append(opts, micro.Name(name))
	}

	if len(address) > 0 {
		opts = append(opts, micro.Address(address))
	}

	if len(endpoint) == 0 {
		endpoint = proxy.DefaultEndpoint
	}

	var p proxy.Proxy

	switch {
	case strings.HasPrefix(endpoint, "grpc"):
		p = grpc.NewProxy(proxy.WithEndpoint(endpoint))
	case strings.HasPrefix(endpoint, "http"):
		p = http.NewProxy(proxy.WithEndpoint(endpoint))
	default:
		p = mucp.NewProxy(proxy.WithEndpoint(endpoint))
	}

	// run the service if asked to
	if len(ctx.Args()) > 0 {
		args := []runtime.CreateOption{
			runtime.WithCommand(ctx.Args()[0], ctx.Args()[1:]...),
			runtime.WithOutput(os.Stdout),
		}

		// register the service
		runtime.Create(&runtime.Service{
			Name: name,
		}, args...)

		// start the runtime
		runtime.Start()

		// stop the runtime
		defer func() {
			runtime.Delete(&runtime.Service{
				Name: name,
			})
			runtime.Stop()
		}()
	}

	log.Logf("Service [%s] Serving %s at endpoint %s\n", p.String(), name, endpoint)

	// new service
	service := micro.NewService(opts...)

	// create new muxer
	muxer := mux.New(name, p)

	// set the router
	service.Server().Init(
		server.WithRouter(muxer),
	)

	// run service
	service.Run()
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "service",
		Usage: "Run a micro service",
		Action: func(ctx *cli.Context) {
			run(ctx, options...)
		},
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "name",
				Usage:  "Name of the service",
				EnvVar: "MICRO_SERVICE_NAME",
			},
			cli.StringFlag{
				Name:   "address",
				Usage:  "Address of the service",
				EnvVar: "MICRO_SERVICE_ADDRESS",
			},
			cli.StringFlag{
				Name:   "endpoint",
				Usage:  "The local service endpoint. Defaults to localhost:9090",
				EnvVar: "MICRO_SERVICE_ENDPOINT",
			},
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<commit_msg>fix the service runtime break<commit_after>// Package service provides a micro service
package service

import (
	"os"
	"strings"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/proxy"
	"github.com/micro/go-micro/proxy/grpc"
	"github.com/micro/go-micro/proxy/http"
	"github.com/micro/go-micro/proxy/mucp"
	"github.com/micro/go-micro/runtime"
	"github.com/micro/go-micro/server"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/go-micro/util/mux"
)

func run(ctx *cli.Context, opts ...micro.Option) {
	log.Name("service")

	name := ctx.String("name")
	address := ctx.String("address")
	endpoint := ctx.String("endpoint")

	if len(name) > 0 {
		opts = append(opts, micro.Name(name))
	}

	if len(address) > 0 {
		opts = append(opts, micro.Address(address))
	}

	if len(endpoint) == 0 {
		endpoint = proxy.DefaultEndpoint
	}

	var p proxy.Proxy

	switch {
	case strings.HasPrefix(endpoint, "grpc"):
		p = grpc.NewProxy(proxy.WithEndpoint(endpoint))
	case strings.HasPrefix(endpoint, "http"):
		p = http.NewProxy(proxy.WithEndpoint(endpoint))
	default:
		p = mucp.NewProxy(proxy.WithEndpoint(endpoint))
	}

	// run the service if asked to
	if len(ctx.Args()) > 0 {
		args := []runtime.CreateOption{
			runtime.WithCommand(ctx.Args()[0], ctx.Args()[1:]...),
			runtime.WithOutput(os.Stdout),
		}

		r := runtime.NewRuntime()

		// register the service
		r.Create(&runtime.Service{
			Name: name,
		}, args...)

		// start the runtime
		r.Start()

		// stop the runtime
		defer func() {
			r.Delete(&runtime.Service{
				Name: name,
			})
			r.Stop()
		}()
	}

	log.Logf("Service [%s] Serving %s at endpoint %s\n", p.String(), name, endpoint)

	// new service
	service := micro.NewService(opts...)

	// create new muxer
	muxer := mux.New(name, p)

	// set the router
	service.Server().Init(
		server.WithRouter(muxer),
	)

	// run service
	service.Run()
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "service",
		Usage: "Run a micro service",
		Action: func(ctx *cli.Context) {
			run(ctx, options...)
		},
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "name",
				Usage:  "Name of the service",
				EnvVar: "MICRO_SERVICE_NAME",
			},
			cli.StringFlag{
				Name:   "address",
				Usage:  "Address of the service",
				EnvVar: "MICRO_SERVICE_ADDRESS",
			},
			cli.StringFlag{
				Name:   "endpoint",
				Usage:  "The local service endpoint. Defaults to localhost:9090",
				EnvVar: "MICRO_SERVICE_ENDPOINT",
			},
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"testing"
	"time"

	"encoding/json"

	"github.com/cosminrentea/gobbler/client"
	"github.com/cosminrentea/gobbler/server/connector"
	"github.com/cosminrentea/gobbler/server/fcm"
	"github.com/cosminrentea/gobbler/server/service"
	"github.com/cosminrentea/gobbler/testutil"
	"github.com/stretchr/testify/assert"
	"github.com/cosminrentea/gobbler/server/configstring"
)

var (
	testHttpPort         = 11000
	timeoutForOneMessage = 150 * time.Millisecond
)

type fcmMetricsMap struct {
	CurrentErrorsCount            int `json:"current_errors_count"`
	CurrentMessagesCount          int `json:"current_messages_count"`
	CurrentMessagesTotalLatencies int `json:"current_messages_total_latencies_nanos"`
	CurrentErrorsTotalLatencies   int `json:"current_errors_total_latencies_nanos"`
}

type fcmMetrics struct {
	TotalSentMessages      int           `json:"fcm.total_sent_messages"`
	TotalSentMessageErrors int           `json:"fcm.total_sent_message_errors"`
	Minute                 fcmMetricsMap `json:"fcm.minute"`
	Hour                   fcmMetricsMap `json:"fcm.hour"`
	Day                    fcmMetricsMap `json:"fcm.day"`
}

type routerMetrics struct {
	CurrentRoutes        int `json:"router.current_routes"`
	CurrentSubscriptions int `json:"router.current_subscriptions"`
}

type expectedValues struct {
	ZeroLatencies        bool
	MessageCount         int
	CurrentRoutes        int
	CurrentSubscriptions int
}

// Test that restarting the service continues to fetch messages from store for a subscription from lastID
func TestFCMRestart(t *testing.T) {
	//defer testutil.EnableDebugForMethod()()
	defer testutil.ResetDefaultRegistryHealthCheck()
	a := assert.New(t)

	receiveC := make(chan bool)
	s, cleanup := serviceSetUp(t)
	defer cleanup()

	assertMetrics(a, s, expectedValues{true, 0, 0, 0})

	var fcmConn connector.ResponsiveConnector
	var ok bool
	for _, iface := range s.ModulesSortedByStartOrder() {
		fcmConn, ok = iface.(connector.ResponsiveConnector)
		if ok {
			break
		}
	}
	a.True(ok, "There should be a module of type FCMConnector")

	// add a high timeout so the messages are processed slow
	sender, err := fcm.CreateFcmSender(fcm.SuccessFCMResponse, receiveC, 10*time.Millisecond)
	a.NoError(err)
	fcmConn.SetSender(sender)

	// create subscription on topic
	subscriptionSetUp(t, s)

	assertMetrics(a, s, expectedValues{true, 0, 1, 1})

	c := clientSetUp(t, s)
	defer c.Close()

	// send 3 messages in the router but read only one and close the service
	for i := 0; i < 3; i++ {
		c.Send(testTopic, "dummy body", `{"Correlation-Id": "id"}`)
	}

	// receive one message only from FCM
	select {
	case <-receiveC:
	case <-time.After(timeoutForOneMessage):
		a.Fail("Initial FCM message not received")
	}

	assertMetrics(a, s, expectedValues{false, 1, 1, 1})
	close(receiveC)
	// restart the service
	a.NoError(s.Stop())

	// remake the sender
	receiveC = make(chan bool)
	sender, err = fcm.CreateFcmSender(fcm.SuccessFCMResponse, receiveC, 10*time.Millisecond)
	a.NoError(err)
	fcmConn.SetSender(sender)

	time.Sleep(50 * time.Millisecond)
	testutil.ResetDefaultRegistryHealthCheck()
	a.NoError(s.Start())

	//TODO Cosmin Bogdan add 2 calls to assertMetrics before and after the next block

	// read the other 2 messages
	for i := 0; i < 1; i++ {
		select {
		case <-receiveC:
		case <-time.After(2 * timeoutForOneMessage):
			a.Fail("FCM message not received")
		}
	}
}

func serviceSetUp(t *testing.T) (*service.Service, func()) {
	dir, errTempDir := ioutil.TempDir("", "guble_fcm_test")
	assert.NoError(t, errTempDir)

	*Config.KVS = "memory"
	*Config.MS = "file"
	*Config.Cluster.NodeID = 0
	*Config.StoragePath = dir
	*Config.MetricsEndpoint = "/admin/metrics-old"
	*Config.WS.Enabled = true
	*Config.WS.Prefix = "/stream/"
	*Config.FCM.Enabled = true
	*Config.FCM.APIKey = "WILL BE OVERWRITTEN"
	*Config.FCM.Prefix = "/fcm/"
	*Config.FCM.Workers = 1 // use only one worker so we can control the number of messages that go to FCM
	*Config.APNS.Enabled = false
	*Config.KafkaProducer.Brokers = configstring.List{}

	var s *service.Service
	for s == nil {
		testHttpPort++
		logger.WithField("port", testHttpPort).Debug("trying to use HTTP Port")
		*Config.HttpListen = fmt.Sprintf("127.0.0.1:%d", testHttpPort)
		s = StartService()
	}
	return s, func() {
		errRemove := os.RemoveAll(dir)
		if errRemove != nil {
			logger.WithError(errRemove).WithField("module", "testing").Error("Could not remove directory")
		}
	}
}

func clientSetUp(t *testing.T, service *service.Service) client.Client {
	wsURL := "ws://" + service.WebServer().GetAddr() + "/stream/user/user01"
	c, err := client.Open(wsURL, "http://localhost/", 1000, false)
	assert.NoError(t, err)
	return c
}

func subscriptionSetUp(t *testing.T, service *service.Service) {
	a := assert.New(t)

	urlFormat := fmt.Sprintf("http://%s/fcm/%%d/gcmId%%d/%%s", service.WebServer().GetAddr())
	// create GCM subscription
	response, errPost := http.Post(
		fmt.Sprintf(urlFormat, 1, 1, strings.TrimPrefix(testTopic, "/")),
		"text/plain",
		bytes.NewBufferString(""),
	)
	a.NoError(errPost)
	a.Equal(response.StatusCode, 200)

	body, errReadAll := ioutil.ReadAll(response.Body)
	a.NoError(errReadAll)
	a.Equal(fmt.Sprintf(`{"subscribed":"%s"}`, testTopic), string(body))
}

func assertMetrics(a *assert.Assertions, s *service.Service, expected expectedValues) {
	time.Sleep(50 * time.Millisecond)

	httpClient := &http.Client{}
	u := fmt.Sprintf("http://%s%s", s.WebServer().GetAddr(), defaultMetricsEndpoint)
	request, err := http.NewRequest(http.MethodGet, u, nil)
	a.NoError(err)

	response, err := httpClient.Do(request)
	a.NoError(err)
	defer response.Body.Close()

	a.Equal(http.StatusOK, response.StatusCode)
	bodyBytes, err := ioutil.ReadAll(response.Body)
	a.NoError(err)
	logger.WithField("body", string(bodyBytes)).Debug("metrics response")

	mFCM := &fcmMetrics{}
	err = json.Unmarshal(bodyBytes, mFCM)
	a.NoError(err)

	a.Equal(0, mFCM.TotalSentMessageErrors)
	a.Equal(expected.MessageCount, mFCM.TotalSentMessages)

	a.Equal(0, mFCM.Minute.CurrentErrorsCount)
	a.Equal(expected.MessageCount, mFCM.Minute.CurrentMessagesCount)
	a.Equal(0, mFCM.Minute.CurrentErrorsTotalLatencies)
	a.Equal(expected.ZeroLatencies, mFCM.Minute.CurrentMessagesTotalLatencies == 0)

	a.Equal(0, mFCM.Hour.CurrentErrorsCount)
	a.Equal(expected.MessageCount, mFCM.Hour.CurrentMessagesCount)
	a.Equal(0, mFCM.Hour.CurrentErrorsTotalLatencies)
	a.Equal(expected.ZeroLatencies, mFCM.Hour.CurrentMessagesTotalLatencies == 0)

	a.Equal(0, mFCM.Day.CurrentErrorsCount)
	a.Equal(expected.MessageCount, mFCM.Day.CurrentMessagesCount)
	a.Equal(0, mFCM.Day.CurrentErrorsTotalLatencies)
	a.Equal(expected.ZeroLatencies, mFCM.Day.CurrentMessagesTotalLatencies == 0)

	mRouter := &routerMetrics{}
	err = json.Unmarshal(bodyBytes, mRouter)
	a.NoError(err)

	a.Equal(expected.CurrentRoutes, mRouter.CurrentRoutes)
	a.Equal(expected.CurrentSubscriptions, mRouter.CurrentSubscriptions)
}
<commit_msg>some test.<commit_after>package server

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"testing"
	"time"

	"encoding/json"

	"github.com/cosminrentea/gobbler/client"
	"github.com/cosminrentea/gobbler/server/configstring"
	"github.com/cosminrentea/gobbler/server/connector"
	"github.com/cosminrentea/gobbler/server/fcm"
	"github.com/cosminrentea/gobbler/server/service"
	"github.com/cosminrentea/gobbler/testutil"
	"github.com/stretchr/testify/assert"
)

var (
	testHttpPort         = 11000
	timeoutForOneMessage = 150 * time.Millisecond
)

type fcmMetricsMap struct {
	CurrentErrorsCount            int `json:"current_errors_count"`
	CurrentMessagesCount          int `json:"current_messages_count"`
	CurrentMessagesTotalLatencies int `json:"current_messages_total_latencies_nanos"`
	CurrentErrorsTotalLatencies   int `json:"current_errors_total_latencies_nanos"`
}

type fcmMetrics struct {
	TotalSentMessages      int           `json:"fcm.total_sent_messages"`
	TotalSentMessageErrors int           `json:"fcm.total_sent_message_errors"`
	Minute                 fcmMetricsMap `json:"fcm.minute"`
	Hour                   fcmMetricsMap `json:"fcm.hour"`
	Day                    fcmMetricsMap `json:"fcm.day"`
}

type routerMetrics struct {
	CurrentRoutes        int `json:"router.current_routes"`
	CurrentSubscriptions int `json:"router.current_subscriptions"`
}

type expectedValues struct {
	ZeroLatencies        bool
	MessageCount         int
	CurrentRoutes        int
	CurrentSubscriptions int
}

// Test that restarting the service continues to fetch messages from store for a subscription from lastID
func TestFCMRestart(t *testing.T) {
	//defer testutil.EnableDebugForMethod()()
	defer testutil.ResetDefaultRegistryHealthCheck()
	a := assert.New(t)

	receiveC := make(chan bool)
	s, cleanup := serviceSetUp(t)
	defer cleanup()

	assertMetrics(a, s, expectedValues{true, 0, 0, 0})

	var fcmConn connector.ResponsiveConnector
	var ok bool
	for _, iface := range s.ModulesSortedByStartOrder() {
		fcmConn, ok = iface.(connector.ResponsiveConnector)
		if ok {
			break
		}
	}
	a.True(ok, "There should be a module of type FCMConnector")

	// add a high timeout so the messages are processed slow
	sender, err := fcm.CreateFcmSender(fcm.SuccessFCMResponse, receiveC, 10*time.Millisecond)
	a.NoError(err)
	fcmConn.SetSender(sender)

	// create subscription on topic
	subscriptionSetUp(t, s)
	time.Sleep(time.Second)
	assertMetrics(a, s, expectedValues{true, 0, 1, 1})

	c := clientSetUp(t, s)
	defer c.Close()

	// send 3 messages in the router but read only one and close the service
	for i := 0; i < 3; i++ {
		c.Send(testTopic, "dummy body", `{"Correlation-Id": "id"}`)
	}

	// receive one message only from FCM
	select {
	case <-receiveC:
	case <-time.After(timeoutForOneMessage):
		a.Fail("Initial FCM message not received")
	}

	time.Sleep(time.Second)

	assertMetrics(a, s, expectedValues{false, 1, 1, 1})
	close(receiveC)
	// restart the service
	a.NoError(s.Stop())

	// remake the sender
	receiveC = make(chan bool)
	sender, err = fcm.CreateFcmSender(fcm.SuccessFCMResponse, receiveC, 10*time.Millisecond)
	a.NoError(err)
	fcmConn.SetSender(sender)

	time.Sleep(50 * time.Millisecond)
	testutil.ResetDefaultRegistryHealthCheck()
	a.NoError(s.Start())

	//TODO Cosmin Bogdan add 2 calls to assertMetrics before and after the next block

	// read the other 2 messages
	for i := 0; i < 1; i++ {
		select {
		case <-receiveC:
		case <-time.After(2 * timeoutForOneMessage):
			a.Fail("FCM message not received")
		}
	}
}

func serviceSetUp(t *testing.T) (*service.Service, func()) {
	dir, errTempDir := ioutil.TempDir("", "guble_fcm_test")
	assert.NoError(t, errTempDir)

	*Config.KVS = "memory"
	*Config.MS = "file"
	*Config.Cluster.NodeID = 0
	*Config.StoragePath = dir
	*Config.MetricsEndpoint = "/admin/metrics-old"
	*Config.WS.Enabled = true
	*Config.WS.Prefix = "/stream/"
	*Config.FCM.Enabled = true
	*Config.FCM.APIKey = "WILL BE OVERWRITTEN"
	*Config.FCM.Prefix = "/fcm/"
	*Config.FCM.Workers = 1 // use only one worker so we can control the number of messages that go to FCM
	*Config.APNS.Enabled = false
	*Config.KafkaProducer.Brokers = configstring.List{}

	var s *service.Service
	for s == nil {
		testHttpPort++
		logger.WithField("port", testHttpPort).Debug("trying to use HTTP Port")
		*Config.HttpListen = fmt.Sprintf("127.0.0.1:%d", testHttpPort)
		s = StartService()
	}
	return s, func() {
		errRemove := os.RemoveAll(dir)
		if errRemove != nil {
			logger.WithError(errRemove).WithField("module", "testing").Error("Could not remove directory")
		}
	}
}

func clientSetUp(t *testing.T, service *service.Service) client.Client {
	wsURL := "ws://" + service.WebServer().GetAddr() + "/stream/user/user01"
	c, err := client.Open(wsURL, "http://localhost/", 1000, false)
	assert.NoError(t, err)
	return c
}

func subscriptionSetUp(t *testing.T, service *service.Service) {
	a := assert.New(t)

	urlFormat := fmt.Sprintf("http://%s/fcm/%%d/gcmId%%d/%%s", service.WebServer().GetAddr())
	// create GCM subscription
	response, errPost := http.Post(
		fmt.Sprintf(urlFormat, 1, 1, strings.TrimPrefix(testTopic, "/")),
		"text/plain",
		bytes.NewBufferString(""),
	)
	a.NoError(errPost)
	a.Equal(response.StatusCode, 200)

	body, errReadAll := ioutil.ReadAll(response.Body)
	a.NoError(errReadAll)
	a.Equal(fmt.Sprintf(`{"subscribed":"%s"}`, testTopic), string(body))
}

func assertMetrics(a *assert.Assertions, s *service.Service, expected expectedValues) {
	time.Sleep(50 * time.Millisecond)

	httpClient := &http.Client{}
	u := fmt.Sprintf("http://%s%s", s.WebServer().GetAddr(), defaultMetricsEndpoint)
	request, err := http.NewRequest(http.MethodGet, u, nil)
	a.NoError(err)

	response, err := httpClient.Do(request)
	a.NoError(err)
	defer response.Body.Close()

	a.Equal(http.StatusOK, response.StatusCode)
	bodyBytes, err := ioutil.ReadAll(response.Body)
	a.NoError(err)
	logger.WithField("body", string(bodyBytes)).Debug("metrics response")

	mFCM := &fcmMetrics{}
	err = json.Unmarshal(bodyBytes, mFCM)
	a.NoError(err)

	a.Equal(0, mFCM.TotalSentMessageErrors)
	a.Equal(expected.MessageCount, mFCM.TotalSentMessages)

	a.Equal(0, mFCM.Minute.CurrentErrorsCount)
	a.Equal(expected.MessageCount, mFCM.Minute.CurrentMessagesCount)
	a.Equal(0, mFCM.Minute.CurrentErrorsTotalLatencies)
	a.Equal(expected.ZeroLatencies, mFCM.Minute.CurrentMessagesTotalLatencies == 0)

	a.Equal(0, mFCM.Hour.CurrentErrorsCount)
	a.Equal(expected.MessageCount, mFCM.Hour.CurrentMessagesCount)
	a.Equal(0, mFCM.Hour.CurrentErrorsTotalLatencies)
	a.Equal(expected.ZeroLatencies, mFCM.Hour.CurrentMessagesTotalLatencies == 0)

	a.Equal(0, mFCM.Day.CurrentErrorsCount)
	a.Equal(expected.MessageCount, mFCM.Day.CurrentMessagesCount)
	a.Equal(0, mFCM.Day.CurrentErrorsTotalLatencies)
	a.Equal(expected.ZeroLatencies, mFCM.Day.CurrentMessagesTotalLatencies == 0)

	mRouter := &routerMetrics{}
	err = json.Unmarshal(bodyBytes, mRouter)
	a.NoError(err)

	a.Equal(expected.CurrentRoutes, mRouter.CurrentRoutes)
	a.Equal(expected.CurrentSubscriptions, mRouter.CurrentSubscriptions)
}
<|endoftext|>