<commit_before>package ledger

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
)

type Entry struct {
	Date        string // "Y-m-d"
	Description string
	Change      int // in cents
}

func isValidCurrency(currency string) bool {
	return currency == "USD" || currency == "EUR"
}

func isValidLocale(locale string) bool {
	return locale == "en-US" || locale == "nl-NL"
}

func isValidDate(entries []Entry) bool {
	for _, entry := range entries {
		if len(entry.Date) != 10 {
			return false
		}
		_, d2, _, d4, _ := entry.Date[0:4], entry.Date[4], entry.Date[5:7], entry.Date[7], entry.Date[8:10]
		if d2 != '-' || d4 != '-' {
			return false
		}
	}
	return true
}

func formatDescription(description string) string {
	if len(description) > 25 {
		return description[:22] + "..."
	}
	return description + strings.Repeat(" ", 25-len(description))
}
func formatDate(locale string, date string) string {
	year, month, day := date[0:4], date[5:7], date[8:10]

	if locale == "nl-NL" {
		return fmt.Sprintf("%v-%v-%v", day, month, year)
	} else if locale == "en-US" {
		return fmt.Sprintf("%v/%v/%v", month, day, year)
	}
	panic("invalid locale")
}

func header(locale string) (output string) {
	if locale == "nl-NL" {
		return "Datum" +
			strings.Repeat(" ", 10-len("Datum")) +
			" | " +
			"Omschrijving" +
			strings.Repeat(" ", 25-len("Omschrijving")) +
			" | " + "Verandering" + "\n"
	} else if locale == "en-US" {
		return "Date" +
			strings.Repeat(" ", 10-len("Date")) +
			" | " +
			"Description" +
			strings.Repeat(" ", 25-len("Description")) +
			" | " + "Change" + "\n"
	}
	panic("invalid locale")
}

func FormatLedger(currency string, locale string, entries []Entry) (output string, e error) {
	if !isValidCurrency(currency) {
		return "", errors.New("invalid currency")
	}
	if !isValidLocale(locale) {
		return "", errors.New("invalid locale")
	}
	if !isValidDate(entries) {
		return "", errors.New("invalid date")
	}

	entriesCopy := append([]Entry{}, entries...)

	m1 := map[bool]int{true: 0, false: 1}
	m2 := map[bool]int{true: -1, false: 1}
	es := entriesCopy
	for len(es) > 1 {
		first, rest := es[0], es[1:]
		success := false
		for !success {
			success = true
			for i, e := range rest {
				if (m1[e.Date == first.Date]*m2[e.Date < first.Date]*4 +
					m1[e.Description == first.Description]*m2[e.Description < first.Description]*2 +
					m1[e.Change == first.Change]*m2[e.Change < first.Change]*1) < 0 {
					es[0], es[i+1] = es[i+1], es[0]
					success = false
				}
			}
		}
		es = es[1:]
	}

	output += header(locale)
	for _, entry := range entriesCopy {
		description := formatDescription(entry.Description)
		date := formatDate(locale, entry.Date)
		cents := entry.Change
		if cents < 0 {
			cents = cents * -1
		}
		var a string
		if locale == "nl-NL" {
			if currency == "EUR" {
				a += "€"
			} else if currency == "USD" {
				a += "$"
			}
			a += " "
			centsStr := strconv.Itoa(cents)
			switch len(centsStr) {
			case 1:
				centsStr = "00" + centsStr
			case 2:
				centsStr = "0" + centsStr
			}
			rest := centsStr[:len(centsStr)-2]
			var parts []string
			for len(rest) > 3 {
				parts = append(parts, rest[len(rest)-3:])
				rest = rest[:len(rest)-3]
			}
			if len(rest) > 0 {
				parts = append(parts, rest)
			}
			for i := len(parts) - 1; i >= 0; i-- {
				a += parts[i] + "."
			}
			a = a[:len(a)-1]
			a += ","
			a += centsStr[len(centsStr)-2:]
			if cents < 0 {
				a += "-"
			} else {
				a += " "
			}
		} else if locale == "en-US" {
			if cents < 0 {
				a += "("
			}
			if currency == "EUR" {
				a += "€"
			} else if currency == "USD" {
				a += "$"
			}
			centsStr := strconv.Itoa(cents)
			switch len(centsStr) {
			case 1:
				centsStr = "00" + centsStr
			case 2:
				centsStr = "0" + centsStr
			}
			rest := centsStr[:len(centsStr)-2]
			var parts []string
			for len(rest) > 3 {
				parts = append(parts, rest[len(rest)-3:])
				rest = rest[:len(rest)-3]
			}
			if len(rest) > 0 {
				parts = append(parts, rest)
			}
			for i := len(parts) - 1; i >= 0; i-- {
				a += parts[i] + ","
			}
			a = a[:len(a)-1]
			a += "."
			a += centsStr[len(centsStr)-2:]
			if cents < 0 {
				a += ")"
			} else {
				a += " "
			}
		}
		output += date + strings.Repeat(" ", 10-len(date)) + " | " + description + " | " +
			strings.Repeat(" ", 13-len(a)) + a + "\n"
	}
	return output, nil
}
<commit_msg>Extract formatChange<commit_after>package ledger

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
)

type Entry struct {
	Date        string // "Y-m-d"
	Description string
	Change      int // in cents
}

func isValidCurrency(currency string) bool {
	return currency == "USD" || currency == "EUR"
}

func isValidLocale(locale string) bool {
	return locale == "en-US" || locale == "nl-NL"
}

func isValidDate(entries []Entry) bool {
	for _, entry := range entries {
		if len(entry.Date) != 10 {
			return false
		}
		_, d2, _, d4, _ := entry.Date[0:4], entry.Date[4], entry.Date[5:7], entry.Date[7], entry.Date[8:10]
		if d2 != '-' || d4 != '-' {
			return false
		}
	}
	return true
}

func formatDescription(description string) string {
	if len(description) > 25 {
		return description[:22] + "..."
	}
	return description + strings.Repeat(" ", 25-len(description))
}
func formatDate(locale string, date string) string {
	year, month, day := date[0:4], date[5:7], date[8:10]

	if locale == "nl-NL" {
		return fmt.Sprintf("%v-%v-%v", day, month, year)
	} else if locale == "en-US" {
		return fmt.Sprintf("%v/%v/%v", month, day, year)
	}
	panic("invalid locale")
}

func header(locale string) (output string) {
	if locale == "nl-NL" {
		return "Datum" +
			strings.Repeat(" ", 10-len("Datum")) +
			" | " +
			"Omschrijving" +
			strings.Repeat(" ", 25-len("Omschrijving")) +
			" | " + "Verandering" + "\n"
	} else if locale == "en-US" {
		return "Date" +
			strings.Repeat(" ", 10-len("Date")) +
			" | " +
			"Description" +
			strings.Repeat(" ", 25-len("Description")) +
			" | " + "Change" + "\n"
	}
	panic("invalid locale")
}

func FormatLedger(currency string, locale string, entries []Entry) (output string, e error) {
	if !isValidCurrency(currency) {
		return "", errors.New("invalid currency")
	}
	if !isValidLocale(locale) {
		return "", errors.New("invalid locale")
	}
	if !isValidDate(entries) {
		return "", errors.New("invalid date")
	}

	entriesCopy := append([]Entry{}, entries...)

	m1 := map[bool]int{true: 0, false: 1}
	m2 := map[bool]int{true: -1, false: 1}
	es := entriesCopy
	for len(es) > 1 {
		first, rest := es[0], es[1:]
		success := false
		for !success {
			success = true
			for i, e := range rest {
				if (m1[e.Date == first.Date]*m2[e.Date < first.Date]*4 +
					m1[e.Description == first.Description]*m2[e.Description < first.Description]*2 +
					m1[e.Change == first.Change]*m2[e.Change < first.Change]*1) < 0 {
					es[0], es[i+1] = es[i+1], es[0]
					success = false
				}
			}
		}
		es = es[1:]
	}

	output += header(locale)
	for _, entry := range entriesCopy {
		description := formatDescription(entry.Description)
		date := formatDate(locale, entry.Date)
		change := formatChange(entry, locale, currency)
		output += date + strings.Repeat(" ", 10-len(date)) + " | " + description + " | " +
			strings.Repeat(" ", 13-len(change)) + change + "\n"
	}
	return output, nil
}

func formatChange(entry Entry, locale string, currency string) string {
	cents := entry.Change
	if cents < 0 {
		cents = cents * -1
	}
	var a string
	if locale == "nl-NL" {
		if currency == "EUR" {
			a += "€"
		} else if currency == "USD" {
			a += "$"
		}
		a += " "
		centsStr := strconv.Itoa(cents)
		switch len(centsStr) {
		case 1:
			centsStr = "00" + centsStr
		case 2:
			centsStr = "0" + centsStr
		}
		rest := centsStr[:len(centsStr)-2]
		var parts []string
		for len(rest) > 3 {
			parts = append(parts, rest[len(rest)-3:])
			rest = rest[:len(rest)-3]
		}
		if len(rest) > 0 {
			parts = append(parts, rest)
		}
		for i := len(parts) - 1; i >= 0; i-- {
			a += parts[i] + "."
		}
		a = a[:len(a)-1]
		a += ","
		a += centsStr[len(centsStr)-2:]
		if cents < 0 {
			a += "-"
		} else {
			a += " "
		}
	} else if locale == "en-US" {
		if cents < 0 {
			a += "("
		}
		if currency == "EUR" {
			a += "€"
		} else if currency == "USD" {
			a += "$"
		}
		centsStr := strconv.Itoa(cents)
		switch len(centsStr) {
		case 1:
			centsStr = "00" + centsStr
		case 2:
			centsStr = "0" + centsStr
		}
		rest := centsStr[:len(centsStr)-2]
		var parts []string
		for len(rest) > 3 {
			parts = append(parts, rest[len(rest)-3:])
			rest = rest[:len(rest)-3]
		}
		if len(rest) > 0 {
			parts = append(parts, rest)
		}
		for i := len(parts) - 1; i >= 0; i-- {
			a += parts[i] + ","
		}
		a = a[:len(a)-1]
		a += "."
		a += centsStr[len(centsStr)-2:]
		if cents < 0 {
			a += ")"
		} else {
			a += " "
		}
	}
	return a
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package libkb

import (
	"bytes"
	"fmt"
	"io"
	"os/exec"
	"sync"
)

type GpgCLI struct {
	path    string
	options []string

	mutex *sync.Mutex

	logUI LogUI
}

type GpgCLIArg struct {
	LogUI LogUI // If nil, use the global
}

func NewGpgCLI(arg GpgCLIArg) *GpgCLI {
	logUI := arg.LogUI
	if logUI == nil {
		logUI = G.Log
	}
	return &GpgCLI{
		mutex: new(sync.Mutex),
		logUI: logUI,
	}
}

func (g *GpgCLI) Configure() (err error) {

	g.mutex.Lock()
	defer g.mutex.Unlock()

	prog := G.Env.GetGpg()
	opts := G.Env.GetGpgOptions()

	if len(prog) > 0 {
		err = canExec(prog)
	} else {
		prog, err = exec.LookPath("gpg2")
		if err != nil {
			prog, err = exec.LookPath("gpg")
		}
	}
	if err != nil {
		return err
	}

	g.logUI.Debug("| configured GPG w/ path: %s", prog)

	g.path = prog
	g.options = opts

	return
}

// CanExec returns true if a gpg executable exists.
func (g *GpgCLI) CanExec() (bool, error) {
	if err := g.Configure(); err != nil {
		if oerr, ok := err.(*exec.Error); ok {
			if oerr.Err == exec.ErrNotFound {
				return false, nil
			}
		}
		return false, err
	}
	return true, nil
}

// Path returns the path of the gpg executable.
// Path is only available if CanExec() is true.
func (g *GpgCLI) Path() string {
	canExec, err := g.CanExec()
	if err == nil && canExec {
		return g.path
	}
	return ""
}

type RunGpgArg struct {
	Arguments []string
	Stdin     bool
	Stderr    io.WriteCloser
	Stdout    io.WriteCloser
}

type RunGpgRes struct {
	Stdin io.WriteCloser
	Err   error
	Wait  func() error
}

func (g *GpgCLI) ImportKey(secret bool, fp PGPFingerprint) (*PGPKeyBundle, error) {
	var cmd string
	var which string
	if secret {
		cmd = "--export-secret-key"
		which = "secret "
	} else {
		cmd = "--export"
	}

	arg := RunGpg2Arg{
		Arguments: []string{"--armor", cmd, fp.String()},
		Stdout:    true,
	}

	res := g.Run2(arg)
	if res.Err != nil {
		return nil, res.Err
	}

	buf := new(bytes.Buffer)
	buf.ReadFrom(res.Stdout)
	armored := buf.String()

	// Convert to posix style on windows
	armored = PosixLineEndings(armored)

	if err := res.Wait(); err != nil {
		return nil, err
	}

	if len(armored) == 0 {
		return nil, NoKeyError{fmt.Sprintf("No %skey found for fingerprint %s", which, fp)}
	}

	bundle, err := ReadOneKeyFromString(armored)
	if err != nil {
		return nil, err
	}

	// For secret keys, *also* import the key in public mode, and then grab the
	// ArmoredPublicKey from that. That's because the public import goes out of
	// its way to preserve the exact armored string from GPG.
	if secret {
		publicBundle, err := g.ImportKey(false, fp)
		if err != nil {
			return nil, err
		}
		bundle.ArmoredPublicKey = publicBundle.ArmoredPublicKey
	}

	return bundle, nil
}

func (g *GpgCLI) ExportKey(k PGPKeyBundle) (err error) {
	arg := RunGpg2Arg{
		Arguments: []string{"--import"},
		Stdin:     true,
	}

	res := g.Run2(arg)
	if res.Err != nil {
		return res.Err
	}

	e1 := k.EncodeToStream(res.Stdin)
	e2 := res.Stdin.Close()
	e3 := res.Wait()
	return PickFirstError(e1, e2, e3)
}

func (g *GpgCLI) Sign(fp PGPFingerprint, payload []byte) (string, error) {
	arg := RunGpg2Arg{
		Arguments: []string{"--armor", "--sign", "--default-key", fp.String()},
		Stdout:    true,
		Stdin:     true,
	}

	res := g.Run2(arg)
	if res.Err != nil {
		return "", res.Err
	}

	res.Stdin.Write(payload)
	res.Stdin.Close()

	buf := new(bytes.Buffer)
	buf.ReadFrom(res.Stdout)
	armored := buf.String()

	// Convert to posix style on windows
	armored = PosixLineEndings(armored)

	if err := res.Wait(); err != nil {
		return "", err
	}

	return armored, nil
}

type RunGpg2Arg struct {
	Arguments []string
	Stdin     bool
	Stderr    bool
	Stdout    bool
}

type RunGpg2Res struct {
	Stdin  io.WriteCloser
	Stdout io.ReadCloser
	Stderr io.ReadCloser
	Wait   func() error
	Err    error
}

func (g *GpgCLI) Run2(arg RunGpg2Arg) (res RunGpg2Res) {

	cmd := g.MakeCmd(arg.Arguments)

	if arg.Stdin {
		if res.Stdin, res.Err = cmd.StdinPipe(); res.Err != nil {
			return
		}
	}

	var stdout, stderr io.ReadCloser

	if stdout, res.Err = cmd.StdoutPipe(); res.Err != nil {
		return
	}
	if stderr, res.Err = cmd.StderrPipe(); res.Err != nil {
		return
	}

	if res.Err = cmd.Start(); res.Err != nil {
		return
	}

	waited := false
	out := 0
	ch := make(chan error)
	var fep FirstErrorPicker

	res.Wait = func() error {
		for out > 0 {
			fep.Push(<-ch)
			out--
		}
		if !waited {
			waited = true
			err := cmd.Wait()
			if err != nil {
				fep.Push(ErrorToGpgError(err))
			}
			return fep.Error()
		}
		return nil
	}

	if !arg.Stdout {
		out++
		go func() {
			ch <- DrainPipe(stdout, func(s string) { g.logUI.Info(s) })
		}()
	} else {
		res.Stdout = stdout
	}

	if !arg.Stderr {
		out++
		go func() {
			ch <- DrainPipe(stderr, func(s string) { g.logUI.Info(s) })
		}()
	} else {
		res.Stderr = stderr
	}

	return
}

func (g *GpgCLI) MakeCmd(args []string) *exec.Cmd {
	var nargs []string
	if g.options != nil {
		nargs = make([]string, len(g.options))
		copy(nargs, g.options)
		nargs = append(nargs, args...)
	} else {
		nargs = args
	}
	g.logUI.Debug("| running Gpg: %s %v", g.path, nargs)
	return exec.Command(g.path, nargs...)
}

func (g *GpgCLI) Run(arg RunGpgArg) (res RunGpgRes) {

	cmd := g.MakeCmd(arg.Arguments)

	waited := false

	var stdout, stderr io.ReadCloser

	if arg.Stdin {
		if res.Stdin, res.Err = cmd.StdinPipe(); res.Err != nil {
			return
		}
	}
	if stdout, res.Err = cmd.StdoutPipe(); res.Err != nil {
		return
	}
	if stderr, res.Err = cmd.StderrPipe(); res.Err != nil {
		return
	}

	if res.Err = cmd.Start(); res.Err != nil {
		return
	}

	waitfn := func() error {
		if !waited {
			waited = true
			return cmd.Wait()
		}
		return nil
	}

	if arg.Stdin {
		res.Wait = waitfn
	} else {
		defer waitfn()
	}

	var e1, e2, e3 error

	if arg.Stdout != nil {
		_, e1 = io.Copy(arg.Stdout, stdout)
	} else {
		e1 = DrainPipe(stdout, func(s string) { g.logUI.Info(s) })
	}

	if arg.Stderr != nil {
		_, e2 = io.Copy(arg.Stderr, stderr)
	} else {
		e2 = DrainPipe(stderr, func(s string) { g.logUI.Warning(s) })
	}

	if !arg.Stdin {
		e3 = waitfn()
	}

	res.Err = PickFirstError(e1, e2, e3)
	return
}
<commit_msg>Change GPG --default-key arg to -u<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package libkb

import (
	"bytes"
	"fmt"
	"io"
	"os/exec"
	"sync"
)

type GpgCLI struct {
	path    string
	options []string

	mutex *sync.Mutex

	logUI LogUI
}

type GpgCLIArg struct {
	LogUI LogUI // If nil, use the global
}

func NewGpgCLI(arg GpgCLIArg) *GpgCLI {
	logUI := arg.LogUI
	if logUI == nil {
		logUI = G.Log
	}
	return &GpgCLI{
		mutex: new(sync.Mutex),
		logUI: logUI,
	}
}

func (g *GpgCLI) Configure() (err error) {

	g.mutex.Lock()
	defer g.mutex.Unlock()

	prog := G.Env.GetGpg()
	opts := G.Env.GetGpgOptions()

	if len(prog) > 0 {
		err = canExec(prog)
	} else {
		prog, err = exec.LookPath("gpg2")
		if err != nil {
			prog, err = exec.LookPath("gpg")
		}
	}
	if err != nil {
		return err
	}

	g.logUI.Debug("| configured GPG w/ path: %s", prog)

	g.path = prog
	g.options = opts

	return
}

// CanExec returns true if a gpg executable exists.
func (g *GpgCLI) CanExec() (bool, error) {
	if err := g.Configure(); err != nil {
		if oerr, ok := err.(*exec.Error); ok {
			if oerr.Err == exec.ErrNotFound {
				return false, nil
			}
		}
		return false, err
	}
	return true, nil
}

// Path returns the path of the gpg executable.
// Path is only available if CanExec() is true.
func (g *GpgCLI) Path() string {
	canExec, err := g.CanExec()
	if err == nil && canExec {
		return g.path
	}
	return ""
}

type RunGpgArg struct {
	Arguments []string
	Stdin     bool
	Stderr    io.WriteCloser
	Stdout    io.WriteCloser
}

type RunGpgRes struct {
	Stdin io.WriteCloser
	Err   error
	Wait  func() error
}

func (g *GpgCLI) ImportKey(secret bool, fp PGPFingerprint) (*PGPKeyBundle, error) {
	var cmd string
	var which string
	if secret {
		cmd = "--export-secret-key"
		which = "secret "
	} else {
		cmd = "--export"
	}

	arg := RunGpg2Arg{
		Arguments: []string{"--armor", cmd, fp.String()},
		Stdout:    true,
	}

	res := g.Run2(arg)
	if res.Err != nil {
		return nil, res.Err
	}

	buf := new(bytes.Buffer)
	buf.ReadFrom(res.Stdout)
	armored := buf.String()

	// Convert to posix style on windows
	armored = PosixLineEndings(armored)

	if err := res.Wait(); err != nil {
		return nil, err
	}

	if len(armored) == 0 {
		return nil, NoKeyError{fmt.Sprintf("No %skey found for fingerprint %s", which, fp)}
	}

	bundle, err := ReadOneKeyFromString(armored)
	if err != nil {
		return nil, err
	}

	// For secret keys, *also* import the key in public mode, and then grab the
	// ArmoredPublicKey from that. That's because the public import goes out of
	// its way to preserve the exact armored string from GPG.
	if secret {
		publicBundle, err := g.ImportKey(false, fp)
		if err != nil {
			return nil, err
		}
		bundle.ArmoredPublicKey = publicBundle.ArmoredPublicKey
	}

	return bundle, nil
}

func (g *GpgCLI) ExportKey(k PGPKeyBundle) (err error) {
	arg := RunGpg2Arg{
		Arguments: []string{"--import"},
		Stdin:     true,
	}

	res := g.Run2(arg)
	if res.Err != nil {
		return res.Err
	}

	e1 := k.EncodeToStream(res.Stdin)
	e2 := res.Stdin.Close()
	e3 := res.Wait()
	return PickFirstError(e1, e2, e3)
}

func (g *GpgCLI) Sign(fp PGPFingerprint, payload []byte) (string, error) {
	arg := RunGpg2Arg{
		Arguments: []string{"--armor", "--sign", "-u", fp.String()},
		Stdout:    true,
		Stdin:     true,
	}

	res := g.Run2(arg)
	if res.Err != nil {
		return "", res.Err
	}

	res.Stdin.Write(payload)
	res.Stdin.Close()

	buf := new(bytes.Buffer)
	buf.ReadFrom(res.Stdout)
	armored := buf.String()

	// Convert to posix style on windows
	armored = PosixLineEndings(armored)

	if err := res.Wait(); err != nil {
		return "", err
	}

	return armored, nil
}

type RunGpg2Arg struct {
	Arguments []string
	Stdin     bool
	Stderr    bool
	Stdout    bool
}

type RunGpg2Res struct {
	Stdin  io.WriteCloser
	Stdout io.ReadCloser
	Stderr io.ReadCloser
	Wait   func() error
	Err    error
}

func (g *GpgCLI) Run2(arg RunGpg2Arg) (res RunGpg2Res) {

	cmd := g.MakeCmd(arg.Arguments)

	if arg.Stdin {
		if res.Stdin, res.Err = cmd.StdinPipe(); res.Err != nil {
			return
		}
	}

	var stdout, stderr io.ReadCloser

	if stdout, res.Err = cmd.StdoutPipe(); res.Err != nil {
		return
	}
	if stderr, res.Err = cmd.StderrPipe(); res.Err != nil {
		return
	}

	if res.Err = cmd.Start(); res.Err != nil {
		return
	}

	waited := false
	out := 0
	ch := make(chan error)
	var fep FirstErrorPicker

	res.Wait = func() error {
		for out > 0 {
			fep.Push(<-ch)
			out--
		}
		if !waited {
			waited = true
			err := cmd.Wait()
			if err != nil {
				fep.Push(ErrorToGpgError(err))
			}
			return fep.Error()
		}
		return nil
	}

	if !arg.Stdout {
		out++
		go func() {
			ch <- DrainPipe(stdout, func(s string) { g.logUI.Info(s) })
		}()
	} else {
		res.Stdout = stdout
	}

	if !arg.Stderr {
		out++
		go func() {
			ch <- DrainPipe(stderr, func(s string) { g.logUI.Info(s) })
		}()
	} else {
		res.Stderr = stderr
	}

	return
}

func (g *GpgCLI) MakeCmd(args []string) *exec.Cmd {
	var nargs []string
	if g.options != nil {
		nargs = make([]string, len(g.options))
		copy(nargs, g.options)
		nargs = append(nargs, args...)
	} else {
		nargs = args
	}
	g.logUI.Debug("| running Gpg: %s %v", g.path, nargs)
	return exec.Command(g.path, nargs...)
}

func (g *GpgCLI) Run(arg RunGpgArg) (res RunGpgRes) {

	cmd := g.MakeCmd(arg.Arguments)

	waited := false

	var stdout, stderr io.ReadCloser

	if arg.Stdin {
		if res.Stdin, res.Err = cmd.StdinPipe(); res.Err != nil {
			return
		}
	}
	if stdout, res.Err = cmd.StdoutPipe(); res.Err != nil {
		return
	}
	if stderr, res.Err = cmd.StderrPipe(); res.Err != nil {
		return
	}

	if res.Err = cmd.Start(); res.Err != nil {
		return
	}

	waitfn := func() error {
		if !waited {
			waited = true
			return cmd.Wait()
		}
		return nil
	}

	if arg.Stdin {
		res.Wait = waitfn
	} else {
		defer waitfn()
	}

	var e1, e2, e3 error

	if arg.Stdout != nil {
		_, e1 = io.Copy(arg.Stdout, stdout)
	} else {
		e1 = DrainPipe(stdout, func(s string) { g.logUI.Info(s) })
	}

	if arg.Stderr != nil {
		_, e2 = io.Copy(arg.Stderr, stderr)
	} else {
		e2 = DrainPipe(stderr, func(s string) { g.logUI.Warning(s) })
	}

	if !arg.Stdin {
		e3 = waitfn()
	}

	res.Err = PickFirstError(e1, e2, e3)
	return
}
<|endoftext|>


<commit_before>package render

import (
	"net"
	"regexp"

	"github.com/weaveworks/scope/probe/host"
	"github.com/weaveworks/scope/report"
)

var (
	// ServiceNodeIDPrefix is how the ID of all service pseudo nodes begin
	ServiceNodeIDPrefix = "service-"

	knownServiceMatchers = []*regexp.Regexp{
		// See http://docs.aws.amazon.com/general/latest/gr/rande.html for fainer grained
		// details
		regexp.MustCompile(`^.+\.amazonaws\.com$`),
		regexp.MustCompile(`^.+\.googleapis\.com$`),
		regexp.MustCompile(`^.+\.core\.windows\.net$`),	// Azure Storage - Blob, Tables, Files & Queues
		regexp.MustCompile(`^.+\.servicebus\.windows\.net$`), // Azure Service Bus
		regexp.MustCompile(`^.+\.azure-api\.net$`), // Azure API Management
		regexp.MustCompile(`^.+\.onmicrosoft\.com$`), // Azure Active Directory
		regexp.MustCompile(`^.+\.cloudapp\.azure\.com$`),	// Azure IaaS
		regexp.MustCompile(`^.+\.database\.windows\.net$`),	// Azure SQL DB
		regexp.MustCompile(`^.+\.documents\.azure\.com$`),	// Azure DocumentDB/CosmosDB
	}

	knownServiceExcluders = []*regexp.Regexp{
		// We exclude ec2 machines because they are too generic
		// and having separate nodes for them makes visualizations worse
		regexp.MustCompile(`^ec2.*\.amazonaws\.com$`),
	}
)

// TODO: Make it user-customizable https://github.com/weaveworks/scope/issues/1876
func isKnownService(hostname string) bool {
	foundMatch := false
	for _, matcher := range knownServiceMatchers {
		if matcher.MatchString(hostname) {
			foundMatch = true
			break
		}
	}
	if !foundMatch {
		return false
	}

	for _, excluder := range knownServiceExcluders {
		if excluder.MatchString(hostname) {
			return false
		}
	}
	return true
}

// LocalNetworks returns a superset of the networks (think: CIDRs) that are
// "local" from the perspective of each host represented in the report. It's
// used to determine which nodes in the report are "remote", i.e. outside of
// our infrastructure.
func LocalNetworks(r report.Report) report.Networks {
	var (
		result   = report.Networks{}
		networks = map[string]struct{}{}
	)

	for _, topology := range []report.Topology{r.Host, r.Overlay} {
		for _, md := range topology.Nodes {
			nets, _ := md.Sets.Lookup(host.LocalNetworks)
			for _, s := range nets {
				_, ipNet, err := net.ParseCIDR(s)
				if err != nil {
					continue
				}
				_, ok := networks[ipNet.String()]
				if !ok {
					result = append(result, ipNet)
					networks[ipNet.String()] = struct{}{}
				}
			}
		}
	}
	return result
}
<commit_msg>Standardized formatting.<commit_after>package render

import (
	"net"
	"regexp"

	"github.com/weaveworks/scope/probe/host"
	"github.com/weaveworks/scope/report"
)

var (
	// ServiceNodeIDPrefix is how the ID of all service pseudo nodes begin
	ServiceNodeIDPrefix = "service-"

	knownServiceMatchers = []*regexp.Regexp{
		// See http://docs.aws.amazon.com/general/latest/gr/rande.html for fainer grained
		// details
		regexp.MustCompile(`^.+\.amazonaws\.com$`),
		regexp.MustCompile(`^.+\.googleapis\.com$`),
		regexp.MustCompile(`^.+\.core\.windows\.net$`),       // Azure Storage - Blob, Tables, Files & Queues
		regexp.MustCompile(`^.+\.servicebus\.windows\.net$`), // Azure Service Bus
		regexp.MustCompile(`^.+\.azure-api\.net$`),           // Azure API Management
		regexp.MustCompile(`^.+\.onmicrosoft\.com$`),         // Azure Active Directory
		regexp.MustCompile(`^.+\.cloudapp\.azure\.com$`),     // Azure IaaS
		regexp.MustCompile(`^.+\.database\.windows\.net$`),   // Azure SQL DB
		regexp.MustCompile(`^.+\.documents\.azure\.com$`),    // Azure DocumentDB/CosmosDB
	}

	knownServiceExcluders = []*regexp.Regexp{
		// We exclude ec2 machines because they are too generic
		// and having separate nodes for them makes visualizations worse
		regexp.MustCompile(`^ec2.*\.amazonaws\.com$`),
	}
)

// TODO: Make it user-customizable https://github.com/weaveworks/scope/issues/1876
func isKnownService(hostname string) bool {
	foundMatch := false
	for _, matcher := range knownServiceMatchers {
		if matcher.MatchString(hostname) {
			foundMatch = true
			break
		}
	}
	if !foundMatch {
		return false
	}

	for _, excluder := range knownServiceExcluders {
		if excluder.MatchString(hostname) {
			return false
		}
	}
	return true
}

// LocalNetworks returns a superset of the networks (think: CIDRs) that are
// "local" from the perspective of each host represented in the report. It's
// used to determine which nodes in the report are "remote", i.e. outside of
// our infrastructure.
func LocalNetworks(r report.Report) report.Networks {
	var (
		result   = report.Networks{}
		networks = map[string]struct{}{}
	)

	for _, topology := range []report.Topology{r.Host, r.Overlay} {
		for _, md := range topology.Nodes {
			nets, _ := md.Sets.Lookup(host.LocalNetworks)
			for _, s := range nets {
				_, ipNet, err := net.ParseCIDR(s)
				if err != nil {
					continue
				}
				_, ok := networks[ipNet.String()]
				if !ok {
					result = append(result, ipNet)
					networks[ipNet.String()] = struct{}{}
				}
			}
		}
	}
	return result
}
<|endoftext|>


<commit_before>package report

import (
	"bytes"
	"fmt"
	"sort"

	"github.com/ugorji/go/codec"
	"github.com/weaveworks/ps"
)

// Helper functions for ps.Map, without considering what is inside

// Return a new map containing all the elements of the two input maps
// and where the same key is in both, pick 'b' where prefer(a,b) is true
func mergeMaps(m, n ps.Map, prefer func(a, b interface{}) bool) ps.Map {
	switch {
	case m == nil:
		return n
	case n == nil:
		return m
	case m.Size() < n.Size():
		m, n = n, m
	}

	n.ForEach(func(key string, val interface{}) {
		if existingVal, found := m.Lookup(key); found {
			if prefer(existingVal, val) {
				m = m.Set(key, val)
			}
		} else {
			m = m.Set(key, val)
		}
	})

	return m
}

func mapEqual(m, n ps.Map, equalf func(a, b interface{}) bool) bool {
	var mSize, nSize int
	if m != nil {
		mSize = m.Size()
	}
	if n != nil {
		nSize = n.Size()
	}
	if mSize != nSize {
		return false
	}
	if mSize == 0 {
		return true
	}
	equal := true
	m.ForEach(func(k string, val interface{}) {
		if otherValue, ok := n.Lookup(k); !ok {
			equal = false
		} else {
			equal = equal && equalf(val, otherValue)
		}
	})
	return equal
}

// very similar to ps.Map.String() but with keys sorted
func mapToString(m ps.Map) string {
	buf := bytes.NewBufferString("{")
	for _, key := range mapKeys(m) {
		val, _ := m.Lookup(key)
		fmt.Fprintf(buf, "%s: %s,\n", key, val)
	}
	fmt.Fprintf(buf, "}")
	return buf.String()
}

func mapKeys(m ps.Map) []string {
	if m == nil {
		return nil
	}
	keys := m.Keys()
	sort.Strings(keys)
	return keys
}

// constants from https://github.com/ugorji/go/blob/master/codec/helper.go#L207
const (
	containerMapKey   = 2
	containerMapValue = 3
	containerMapEnd   = 4
	// from https://github.com/ugorji/go/blob/master/codec/helper.go#L152
	cUTF8 = 2
)

// This implementation does not use the intermediate form as that was a
// performance issue; skipping it saved almost 10% CPU.  Note this means
// we are using undocumented, internal APIs, which could break in the future.
// See https://github.com/weaveworks/scope/pull/1709 for more information.
func mapRead(decoder *codec.Decoder, decodeValue func(isNil bool) interface{}) ps.Map {
	z, r := codec.GenHelperDecoder(decoder)
	if r.TryDecodeAsNil() {
		return ps.NewMap()
	}

	length := r.ReadMapStart()
	out := ps.NewMap()
	for i := 0; length < 0 || i < length; i++ {
		if length < 0 && r.CheckBreak() {
			break
		}

		var key string
		z.DecSendContainerState(containerMapKey)
		if !r.TryDecodeAsNil() {
			key = r.DecodeString()
		}

		z.DecSendContainerState(containerMapValue)
		value := decodeValue(r.TryDecodeAsNil())
		out = out.UnsafeMutableSet(key, value)
	}
	z.DecSendContainerState(containerMapEnd)
	return out
}

// Inverse of mapRead, done for performance. Same comments about
// undocumented internal APIs apply.
func mapWrite(m ps.Map, encoder *codec.Encoder, encodeValue func(*codec.Encoder, interface{})) {
	z, r := codec.GenHelperEncoder(encoder)
	if m == nil {
		r.EncodeNil()
		return
	}
	r.EncodeMapStart(m.Size())
	m.ForEach(func(key string, val interface{}) {
		z.EncSendContainerState(containerMapKey)
		r.EncodeString(cUTF8, key)
		z.EncSendContainerState(containerMapValue)
		encodeValue(encoder, val)
	})
	z.EncSendContainerState(containerMapEnd)
}
<commit_msg>encode empty ps.Maps as nil<commit_after>package report

import (
	"bytes"
	"fmt"
	"sort"

	"github.com/ugorji/go/codec"
	"github.com/weaveworks/ps"
)

// Helper functions for ps.Map, without considering what is inside

// Return a new map containing all the elements of the two input maps
// and where the same key is in both, pick 'b' where prefer(a,b) is true
func mergeMaps(m, n ps.Map, prefer func(a, b interface{}) bool) ps.Map {
	switch {
	case m == nil:
		return n
	case n == nil:
		return m
	case m.Size() < n.Size():
		m, n = n, m
	}

	n.ForEach(func(key string, val interface{}) {
		if existingVal, found := m.Lookup(key); found {
			if prefer(existingVal, val) {
				m = m.Set(key, val)
			}
		} else {
			m = m.Set(key, val)
		}
	})

	return m
}

func mapEqual(m, n ps.Map, equalf func(a, b interface{}) bool) bool {
	var mSize, nSize int
	if m != nil {
		mSize = m.Size()
	}
	if n != nil {
		nSize = n.Size()
	}
	if mSize != nSize {
		return false
	}
	if mSize == 0 {
		return true
	}
	equal := true
	m.ForEach(func(k string, val interface{}) {
		if otherValue, ok := n.Lookup(k); !ok {
			equal = false
		} else {
			equal = equal && equalf(val, otherValue)
		}
	})
	return equal
}

// very similar to ps.Map.String() but with keys sorted
func mapToString(m ps.Map) string {
	buf := bytes.NewBufferString("{")
	for _, key := range mapKeys(m) {
		val, _ := m.Lookup(key)
		fmt.Fprintf(buf, "%s: %s,\n", key, val)
	}
	fmt.Fprintf(buf, "}")
	return buf.String()
}

func mapKeys(m ps.Map) []string {
	if m == nil {
		return nil
	}
	keys := m.Keys()
	sort.Strings(keys)
	return keys
}

// constants from https://github.com/ugorji/go/blob/master/codec/helper.go#L207
const (
	containerMapKey   = 2
	containerMapValue = 3
	containerMapEnd   = 4
	// from https://github.com/ugorji/go/blob/master/codec/helper.go#L152
	cUTF8 = 2
)

// This implementation does not use the intermediate form as that was a
// performance issue; skipping it saved almost 10% CPU.  Note this means
// we are using undocumented, internal APIs, which could break in the future.
// See https://github.com/weaveworks/scope/pull/1709 for more information.
func mapRead(decoder *codec.Decoder, decodeValue func(isNil bool) interface{}) ps.Map {
	z, r := codec.GenHelperDecoder(decoder)
	if r.TryDecodeAsNil() {
		return ps.NewMap()
	}

	length := r.ReadMapStart()
	out := ps.NewMap()
	for i := 0; length < 0 || i < length; i++ {
		if length < 0 && r.CheckBreak() {
			break
		}

		var key string
		z.DecSendContainerState(containerMapKey)
		if !r.TryDecodeAsNil() {
			key = r.DecodeString()
		}

		z.DecSendContainerState(containerMapValue)
		value := decodeValue(r.TryDecodeAsNil())
		out = out.UnsafeMutableSet(key, value)
	}
	z.DecSendContainerState(containerMapEnd)
	return out
}

// Inverse of mapRead, done for performance. Same comments about
// undocumented internal APIs apply.
func mapWrite(m ps.Map, encoder *codec.Encoder, encodeValue func(*codec.Encoder, interface{})) {
	z, r := codec.GenHelperEncoder(encoder)
	if m == nil || m.IsNil() {
		r.EncodeNil()
		return
	}
	r.EncodeMapStart(m.Size())
	m.ForEach(func(key string, val interface{}) {
		z.EncSendContainerState(containerMapKey)
		r.EncodeString(cUTF8, key)
		z.EncSendContainerState(containerMapValue)
		encodeValue(encoder, val)
	})
	z.EncSendContainerState(containerMapEnd)
}
<|endoftext|>


<commit_before>package model

// A Atdn
type At struct {
	Events []AAt `json:"events"`
}

type AAt struct {
	Event AtdnEvent `json:"event"`
}

// C Connpass
type Cp struct {
	Events []ConnpassEvent `json:"events"`
}

// D Doorkeeper
type Dk struct {
	Event DoorkeeperEvent `json:"event"`
}

type Event struct {
	EventId    int `json:"event_id"`
	ApiId      int
	Title      string `json:"title"`
	Desc       string `json:"desc"`
	Url        string `json:"url"`
	Limit      int    `json:"limit"`
	Accepted   int    `json:"accepted"`
	Waitlisted int    `json:"waitlisted"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"strat_at"`
	EndAt      string `json:"end_at"`
	ID         int    `json:"id"`
}

type EventJson struct {
	EventId    string `json:"event_id"`
	ApiId      int
	Title      string `json:"title"`
	//Desc       string `json:"desc"`
	Url        string `json:"url"`
	//Limit      int    `json:"limit"`
	//Accepted   int    `json:"accepted"`
	//Waitlisted int    `json:"waitlisted"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"strat_at"`
	EndAt      string `json:"end_at"`
	ID         int    `json:"id"`
}

type AtdnEvent struct {
	EventId    int `json:"event_id"`
	ApiId      int
	Title      string `json:"title"`
	Desc       string `json:"description"`
	Url        string `json:"event_url"`
	Limit      int    `json:"limit"`
	Accepted   int    `json:"accepted"`
	Waitlisted int    `json:"waiting"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"started_at"`
	EndAt      string `json:"ended_at"`
	ID         int    `json:"id"`
}

type ConnpassEvent struct {
	EventId    int `json:"event_id"`
	ApiId      int
	Title      string `json:"title"`
	Desc       string `json:"description"`
	Url        string `json:"event_url"`
	Limit      int    `json:"limit"`
	Accepted   int    `json:"accepted"`
	Waitlisted int    `json:"waiting"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"started_at"`
	EndAt      string `json:"ended_at"`
	ID         int    `json:"id"`
}

type DoorkeeperEvent struct {
	EventId    int `json:"id"`
	ApiId      int
	Title      string `json:"title"`
	Desc       string `json:"description"`
	Url        string `json:"public_url"`
	Limit      int    `json:"ticket_limit"`
	Accepted   int    `json:"participants"`
	Waitlisted int    `json:"waitlisted"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"starts_at"`
	EndAt      string `json:"ends_at"`
	ID         int    `json:"id"`
}
<commit_msg>[go][delete]jsonで無駄な情報が帰ってきていたので削除<commit_after>package model

// A Atdn
type At struct {
	Events []AAt `json:"events"`
}

type AAt struct {
	Event AtdnEvent `json:"event"`
}

// C Connpass
type Cp struct {
	Events []ConnpassEvent `json:"events"`
}

// D Doorkeeper
type Dk struct {
	Event DoorkeeperEvent `json:"event"`
}

type Event struct {
	EventId    int `json:"event_id"`
	ApiId      int
	Title      string `json:"title"`
	Desc       string `json:"desc"`
	Url        string `json:"url"`
	Limit      int    `json:"limit"`
	Accepted   int    `json:"accepted"`
	Waitlisted int    `json:"waitlisted"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"strat_at"`
	EndAt      string `json:"end_at"`
	ID         int    `json:"id"`
}

type EventJson struct {
	EventId    string `json:"event_id"`
	//ApiId      int
	Title      string `json:"title"`
	//Desc       string `json:"desc"`
	Url        string `json:"url"`
	//Limit      int    `json:"limit"`
	//Accepted   int    `json:"accepted"`
	//Waitlisted int    `json:"waitlisted"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"strat_at"`
	EndAt      string `json:"end_at"`
	ID         int    `json:"id"`
}

type AtdnEvent struct {
	EventId    int `json:"event_id"`
	ApiId      int
	Title      string `json:"title"`
	Desc       string `json:"description"`
	Url        string `json:"event_url"`
	Limit      int    `json:"limit"`
	Accepted   int    `json:"accepted"`
	Waitlisted int    `json:"waiting"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"started_at"`
	EndAt      string `json:"ended_at"`
	ID         int    `json:"id"`
}

type ConnpassEvent struct {
	EventId    int `json:"event_id"`
	ApiId      int
	Title      string `json:"title"`
	Desc       string `json:"description"`
	Url        string `json:"event_url"`
	Limit      int    `json:"limit"`
	Accepted   int    `json:"accepted"`
	Waitlisted int    `json:"waiting"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"started_at"`
	EndAt      string `json:"ended_at"`
	ID         int    `json:"id"`
}

type DoorkeeperEvent struct {
	EventId    int `json:"id"`
	ApiId      int
	Title      string `json:"title"`
	Desc       string `json:"description"`
	Url        string `json:"public_url"`
	Limit      int    `json:"ticket_limit"`
	Accepted   int    `json:"participants"`
	Waitlisted int    `json:"waitlisted"`
	Address    string `json:"address"`
	Place      string `json:"place"`
	StratAt    string `json:"starts_at"`
	EndAt      string `json:"ends_at"`
	ID         int    `json:"id"`
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"runtime/debug"
	"strconv"
	"strings"
	"sync"

	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/codec"
	c "github.com/micro/go-micro/context"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/transport"

	log "github.com/golang/glog"

	"golang.org/x/net/context"
)

type rpcServer struct {
	rpc  *server
	exit chan chan error

	sync.RWMutex
	opts        Options
	handlers    map[string]Handler
	subscribers map[*subscriber][]broker.Subscriber
}

func newRpcServer(opts ...Option) Server {
	options := newOptions(opts...)
	return &rpcServer{
		opts: options,
		rpc: &server{
			name:         options.Name,
			serviceMap:   make(map[string]*service),
			hdlrWrappers: options.HdlrWrappers,
		},
		handlers:    make(map[string]Handler),
		subscribers: make(map[*subscriber][]broker.Subscriber),
		exit:        make(chan chan error),
	}
}

func (s *rpcServer) accept(sock transport.Socket) {
	defer func() {
		if r := recover(); r != nil {
			log.Error(r, string(debug.Stack()))
			sock.Close()
		}
	}()

	var msg transport.Message
	if err := sock.Recv(&msg); err != nil {
		return
	}

	ct := msg.Header["Content-Type"]
	cf, err := s.newCodec(ct)
	// TODO: needs better error handling
	if err != nil {
		sock.Send(&transport.Message{
			Header: map[string]string{
				"Content-Type": "text/plain",
			},
			Body: []byte(err.Error()),
		})
		sock.Close()
		return
	}

	codec := newRpcPlusCodec(&msg, sock, cf)

	// strip our headers
	hdr := make(map[string]string)
	for k, v := range msg.Header {
		hdr[k] = v
	}
	delete(hdr, "Content-Type")

	ctx := c.WithMetadata(context.Background(), hdr)

	// TODO: needs better error handling
	if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
		log.Errorf("Unexpected error serving request, closing socket: %v", err)
		sock.Close()
	}
}

func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
	if cf, ok := s.opts.Codecs[contentType]; ok {
		return cf, nil
	}
	if cf, ok := defaultCodecs[contentType]; ok {
		return cf, nil
	}
	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}

func (s *rpcServer) Options() Options {
	s.RLock()
	opts := s.opts
	s.RUnlock()
	return opts
}

func (s *rpcServer) Init(opts ...Option) error {
	s.Lock()
	for _, opt := range opts {
		opt(&s.opts)
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) NewHandler(h interface{}) Handler {
	return newRpcHandler(h)
}

func (s *rpcServer) Handle(h Handler) error {
	if err := s.rpc.register(h.Handler()); err != nil {
		return err
	}
	s.Lock()
	s.handlers[h.Name()] = h
	s.Unlock()
	return nil
}

func (s *rpcServer) NewSubscriber(topic string, sb interface{}) Subscriber {
	return newSubscriber(topic, sb)
}

func (s *rpcServer) Subscribe(sb Subscriber) error {
	sub, ok := sb.(*subscriber)
	if !ok {
		return fmt.Errorf("invalid subscriber: expected *subscriber")
	}
	if len(sub.handlers) == 0 {
		return fmt.Errorf("invalid subscriber: no handler functions")
	}

	if err := validateSubscriber(sb); err != nil {
		return err
	}

	s.Lock()
	_, ok = s.subscribers[sub]
	if ok {
		return fmt.Errorf("subscriber %v already exists", s)
	}
	s.subscribers[sub] = nil
	s.Unlock()
	return nil
}

func (s *rpcServer) Register() error {
	// parse address for host, port
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := extractAddress(host)
	if err != nil {
		return err
	}

	// register service
	node := &registry.Node{
		Id:       config.Name + "-" + config.Id,
		Address:  addr,
		Port:     port,
		Metadata: config.Metadata,
	}

	node.Metadata["transport"] = config.Transport.String()

	s.RLock()
	var endpoints []*registry.Endpoint
	for _, e := range s.handlers {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	for e, _ := range s.subscribers {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	s.RUnlock()

	service := &registry.Service{
		Name:      config.Name,
		Version:   config.Version,
		Nodes:     []*registry.Node{node},
		Endpoints: endpoints,
	}

	log.Infof("Registering node: %s", node.Id)
	if err := config.Registry.Register(service); err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for sb, _ := range s.subscribers {
		handler := s.createSubHandler(sb, s.opts)
		sub, err := config.Broker.Subscribe(sb.Topic(), handler)
		if err != nil {
			return err
		}
		s.subscribers[sb] = []broker.Subscriber{sub}
	}

	return nil
}

func (s *rpcServer) Deregister() error {
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := extractAddress(host)
	if err != nil {
		return err
	}

	node := &registry.Node{
		Id:      config.Name + "-" + config.Id,
		Address: addr,
		Port:    port,
	}

	service := &registry.Service{
		Name:    config.Name,
		Version: config.Version,
		Nodes:   []*registry.Node{node},
	}

	log.Infof("Deregistering node: %s", node.Id)
	if err := config.Registry.Deregister(service); err != nil {
		return err
	}

	s.Lock()
	for sb, subs := range s.subscribers {
		for _, sub := range subs {
			log.Infof("Unsubscribing from topic: %s", sub.Topic())
			sub.Unsubscribe()
		}
		s.subscribers[sb] = nil
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) Start() error {
	registerHealthChecker(s)
	config := s.Options()

	ts, err := config.Transport.Listen(config.Address)
	if err != nil {
		return err
	}

	log.Infof("Listening on %s", ts.Addr())
	s.Lock()
	s.opts.Address = ts.Addr()
	s.Unlock()

	go ts.Accept(s.accept)

	go func() {
		ch := <-s.exit
		ch <- ts.Close()
		config.Broker.Disconnect()
	}()

	// TODO: subscribe to cruft
	return config.Broker.Connect()
}

func (s *rpcServer) Stop() error {
	ch := make(chan error)
	s.exit <- ch
	return <-ch
}

func (s *rpcServer) String() string {
	return "rpc"
}
<commit_msg>Add broker and server to metadata<commit_after>package server

import (
	"fmt"
	"runtime/debug"
	"strconv"
	"strings"
	"sync"

	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/codec"
	c "github.com/micro/go-micro/context"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/transport"

	log "github.com/golang/glog"

	"golang.org/x/net/context"
)

type rpcServer struct {
	rpc  *server
	exit chan chan error

	sync.RWMutex
	opts        Options
	handlers    map[string]Handler
	subscribers map[*subscriber][]broker.Subscriber
}

func newRpcServer(opts ...Option) Server {
	options := newOptions(opts...)
	return &rpcServer{
		opts: options,
		rpc: &server{
			name:         options.Name,
			serviceMap:   make(map[string]*service),
			hdlrWrappers: options.HdlrWrappers,
		},
		handlers:    make(map[string]Handler),
		subscribers: make(map[*subscriber][]broker.Subscriber),
		exit:        make(chan chan error),
	}
}

func (s *rpcServer) accept(sock transport.Socket) {
	defer func() {
		if r := recover(); r != nil {
			log.Error(r, string(debug.Stack()))
			sock.Close()
		}
	}()

	var msg transport.Message
	if err := sock.Recv(&msg); err != nil {
		return
	}

	ct := msg.Header["Content-Type"]
	cf, err := s.newCodec(ct)
	// TODO: needs better error handling
	if err != nil {
		sock.Send(&transport.Message{
			Header: map[string]string{
				"Content-Type": "text/plain",
			},
			Body: []byte(err.Error()),
		})
		sock.Close()
		return
	}

	codec := newRpcPlusCodec(&msg, sock, cf)

	// strip our headers
	hdr := make(map[string]string)
	for k, v := range msg.Header {
		hdr[k] = v
	}
	delete(hdr, "Content-Type")

	ctx := c.WithMetadata(context.Background(), hdr)

	// TODO: needs better error handling
	if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
		log.Errorf("Unexpected error serving request, closing socket: %v", err)
		sock.Close()
	}
}

func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
	if cf, ok := s.opts.Codecs[contentType]; ok {
		return cf, nil
	}
	if cf, ok := defaultCodecs[contentType]; ok {
		return cf, nil
	}
	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}

func (s *rpcServer) Options() Options {
	s.RLock()
	opts := s.opts
	s.RUnlock()
	return opts
}

func (s *rpcServer) Init(opts ...Option) error {
	s.Lock()
	for _, opt := range opts {
		opt(&s.opts)
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) NewHandler(h interface{}) Handler {
	return newRpcHandler(h)
}

func (s *rpcServer) Handle(h Handler) error {
	if err := s.rpc.register(h.Handler()); err != nil {
		return err
	}
	s.Lock()
	s.handlers[h.Name()] = h
	s.Unlock()
	return nil
}

func (s *rpcServer) NewSubscriber(topic string, sb interface{}) Subscriber {
	return newSubscriber(topic, sb)
}

func (s *rpcServer) Subscribe(sb Subscriber) error {
	sub, ok := sb.(*subscriber)
	if !ok {
		return fmt.Errorf("invalid subscriber: expected *subscriber")
	}
	if len(sub.handlers) == 0 {
		return fmt.Errorf("invalid subscriber: no handler functions")
	}

	if err := validateSubscriber(sb); err != nil {
		return err
	}

	s.Lock()
	_, ok = s.subscribers[sub]
	if ok {
		return fmt.Errorf("subscriber %v already exists", s)
	}
	s.subscribers[sub] = nil
	s.Unlock()
	return nil
}

func (s *rpcServer) Register() error {
	// parse address for host, port
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := extractAddress(host)
	if err != nil {
		return err
	}

	// register service
	node := &registry.Node{
		Id:       config.Name + "-" + config.Id,
		Address:  addr,
		Port:     port,
		Metadata: config.Metadata,
	}

	node.Metadata["transport"] = config.Transport.String()
	node.Metadata["broker"] = config.Broker.String()
	node.Metadata["server"] = s.String()

	s.RLock()
	var endpoints []*registry.Endpoint
	for _, e := range s.handlers {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	for e, _ := range s.subscribers {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	s.RUnlock()

	service := &registry.Service{
		Name:      config.Name,
		Version:   config.Version,
		Nodes:     []*registry.Node{node},
		Endpoints: endpoints,
	}

	log.Infof("Registering node: %s", node.Id)
	if err := config.Registry.Register(service); err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for sb, _ := range s.subscribers {
		handler := s.createSubHandler(sb, s.opts)
		sub, err := config.Broker.Subscribe(sb.Topic(), handler)
		if err != nil {
			return err
		}
		s.subscribers[sb] = []broker.Subscriber{sub}
	}

	return nil
}

func (s *rpcServer) Deregister() error {
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := extractAddress(host)
	if err != nil {
		return err
	}

	node := &registry.Node{
		Id:      config.Name + "-" + config.Id,
		Address: addr,
		Port:    port,
	}

	service := &registry.Service{
		Name:    config.Name,
		Version: config.Version,
		Nodes:   []*registry.Node{node},
	}

	log.Infof("Deregistering node: %s", node.Id)
	if err := config.Registry.Deregister(service); err != nil {
		return err
	}

	s.Lock()
	for sb, subs := range s.subscribers {
		for _, sub := range subs {
			log.Infof("Unsubscribing from topic: %s", sub.Topic())
			sub.Unsubscribe()
		}
		s.subscribers[sb] = nil
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) Start() error {
	registerHealthChecker(s)
	config := s.Options()

	ts, err := config.Transport.Listen(config.Address)
	if err != nil {
		return err
	}

	log.Infof("Listening on %s", ts.Addr())
	s.Lock()
	s.opts.Address = ts.Addr()
	s.Unlock()

	go ts.Accept(s.accept)

	go func() {
		ch := <-s.exit
		ch <- ts.Close()
		config.Broker.Disconnect()
	}()

	// TODO: subscribe to cruft
	return config.Broker.Connect()
}

func (s *rpcServer) Stop() error {
	ch := make(chan error)
	s.exit <- ch
	return <-ch
}

func (s *rpcServer) String() string {
	return "rpc"
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"fmt"
	"reflect"

	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/metadata"
	"github.com/micro/go-micro/registry"
	"golang.org/x/net/context"
)

const (
	subSig = "func(context.Context, interface{}) error"
)

type handler struct {
	method  reflect.Value
	reqType reflect.Type
	ctxType reflect.Type
}

type subscriber struct {
	topic      string
	rcvr       reflect.Value
	typ        reflect.Type
	subscriber interface{}
	handlers   []*handler
	endpoints  []*registry.Endpoint
	opts       SubscriberOptions
}

func newSubscriber(topic string, sub interface{}, opts ...SubscriberOption) Subscriber {
	var options SubscriberOptions
	for _, o := range opts {
		o(&options)
	}

	var endpoints []*registry.Endpoint
	var handlers []*handler

	if typ := reflect.TypeOf(sub); typ.Kind() == reflect.Func {
		h := &handler{
			method: reflect.ValueOf(sub),
		}

		switch typ.NumIn() {
		case 1:
			h.reqType = typ.In(0)
		case 2:
			h.ctxType = typ.In(0)
			h.reqType = typ.In(1)
		}

		handlers = append(handlers, h)

		endpoints = append(endpoints, &registry.Endpoint{
			Name:    "Func",
			Request: extractSubValue(typ),
			Metadata: map[string]string{
				"topic":      topic,
				"subscriber": "true",
			},
		})
	} else {
		hdlr := reflect.ValueOf(sub)
		name := reflect.Indirect(hdlr).Type().Name()

		for m := 0; m < typ.NumMethod(); m++ {
			method := typ.Method(m)
			h := &handler{
				method: method.Func,
			}

			switch method.Type.NumIn() {
			case 2:
				h.reqType = method.Type.In(1)
			case 3:
				h.ctxType = method.Type.In(1)
				h.reqType = method.Type.In(2)
			}

			handlers = append(handlers, h)

			endpoints = append(endpoints, &registry.Endpoint{
				Name:    name + "." + method.Name,
				Request: extractSubValue(method.Type),
				Metadata: map[string]string{
					"topic":      topic,
					"subscriber": "true",
				},
			})
		}
	}

	return &subscriber{
		rcvr:       reflect.ValueOf(sub),
		typ:        reflect.TypeOf(sub),
		topic:      topic,
		subscriber: sub,
		handlers:   handlers,
		endpoints:  endpoints,
		opts:       options,
	}
}

func validateSubscriber(sub Subscriber) error {
	typ := reflect.TypeOf(sub.Subscriber())
	var argType reflect.Type

	if typ.Kind() == reflect.Func {
		name := "Func"
		switch typ.NumIn() {
		case 2:
			argType = typ.In(1)
		default:
			return fmt.Errorf("subscriber %v takes wrong number of args: %v required signature %s", name, typ.NumIn(), subSig)
		}
		if !isExportedOrBuiltinType(argType) {
			return fmt.Errorf("subscriber %v argument type not exported: %v", name, argType)
		}
		if typ.NumOut() != 1 {
			return fmt.Errorf("subscriber %v has wrong number of outs: %v require signature %s",
				name, typ.NumOut(), subSig)
		}
		if returnType := typ.Out(0); returnType != typeOfError {
			return fmt.Errorf("subscriber %v returns %v not error", name, returnType.String())
		}
	} else {
		hdlr := reflect.ValueOf(sub.Subscriber())
		name := reflect.Indirect(hdlr).Type().Name()

		for m := 0; m < typ.NumMethod(); m++ {
			method := typ.Method(m)

			switch method.Type.NumIn() {
			case 3:
				argType = method.Type.In(2)
			default:
				return fmt.Errorf("subscriber %v.%v takes wrong number of args: %v required signature %s",
					name, method.Name, method.Type.NumIn(), subSig)
			}

			if !isExportedOrBuiltinType(argType) {
				return fmt.Errorf("%v argument type not exported: %v", name, argType)
			}
			if method.Type.NumOut() != 1 {
				return fmt.Errorf(
					"subscriber %v.%v has wrong number of outs: %v require signature %s",
					name, method.Name, method.Type.NumOut(), subSig)
			}
			if returnType := method.Type.Out(0); returnType != typeOfError {
				return fmt.Errorf("subscriber %v.%v returns %v not error", name, method.Name, returnType.String())
			}
		}
	}

	return nil
}

func (s *rpcServer) createSubHandler(sb *subscriber, opts Options) broker.Handler {
	return func(p broker.Publication) error {
		msg := p.Message()
		ct := msg.Header["Content-Type"]
		cf, err := s.newCodec(ct)
		if err != nil {
			return err
		}

		hdr := make(map[string]string)
		for k, v := range msg.Header {
			hdr[k] = v
		}
		delete(hdr, "Content-Type")
		ctx := metadata.NewContext(context.Background(), hdr)

		for i := 0; i < len(sb.handlers); i++ {
			handler := sb.handlers[i]

			var isVal bool
			var req reflect.Value

			if handler.reqType.Kind() == reflect.Ptr {
				req = reflect.New(handler.reqType.Elem())
			} else {
				req = reflect.New(handler.reqType)
				isVal = true
			}
			if isVal {
				req = req.Elem()
			}

			b := &buffer{bytes.NewBuffer(msg.Body)}
			co := cf(b)
			defer co.Close()

			if err := co.ReadHeader(&codec.Message{}, codec.Publication); err != nil {
				return err
			}

			if err := co.ReadBody(req.Interface()); err != nil {
				return err
			}

			fn := func(ctx context.Context, msg Publication) error {
				var vals []reflect.Value
				if sb.typ.Kind() != reflect.Func {
					vals = append(vals, sb.rcvr)
				}
				if handler.ctxType != nil {
					vals = append(vals, reflect.ValueOf(ctx))
				}

				vals = append(vals, reflect.ValueOf(msg.Message()))

				returnValues := handler.method.Call(vals)
				if err := returnValues[0].Interface(); err != nil {
					return err.(error)
				}
				return nil
			}

			for i := len(opts.SubWrappers); i > 0; i-- {
				fn = opts.SubWrappers[i-1](fn)
			}

			go fn(ctx, &rpcPublication{
				topic:       sb.topic,
				contentType: ct,
				message:     req.Interface(),
			})
		}
		return nil
	}
}

func (s *subscriber) Topic() string {
	return s.topic
}

func (s *subscriber) Subscriber() interface{} {
	return s.subscriber
}

func (s *subscriber) Endpoints() []*registry.Endpoint {
	return s.endpoints
}

func (s *subscriber) Options() SubscriberOptions {
	return s.opts
}
<commit_msg>wait for subscribers to complete<commit_after>package server

import (
	"bytes"
	"fmt"
	"reflect"

	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/metadata"
	"github.com/micro/go-micro/registry"
	"golang.org/x/net/context"
)

const (
	subSig = "func(context.Context, interface{}) error"
)

type handler struct {
	method  reflect.Value
	reqType reflect.Type
	ctxType reflect.Type
}

type subscriber struct {
	topic      string
	rcvr       reflect.Value
	typ        reflect.Type
	subscriber interface{}
	handlers   []*handler
	endpoints  []*registry.Endpoint
	opts       SubscriberOptions
}

func newSubscriber(topic string, sub interface{}, opts ...SubscriberOption) Subscriber {
	var options SubscriberOptions
	for _, o := range opts {
		o(&options)
	}

	var endpoints []*registry.Endpoint
	var handlers []*handler

	if typ := reflect.TypeOf(sub); typ.Kind() == reflect.Func {
		h := &handler{
			method: reflect.ValueOf(sub),
		}

		switch typ.NumIn() {
		case 1:
			h.reqType = typ.In(0)
		case 2:
			h.ctxType = typ.In(0)
			h.reqType = typ.In(1)
		}

		handlers = append(handlers, h)

		endpoints = append(endpoints, &registry.Endpoint{
			Name:    "Func",
			Request: extractSubValue(typ),
			Metadata: map[string]string{
				"topic":      topic,
				"subscriber": "true",
			},
		})
	} else {
		hdlr := reflect.ValueOf(sub)
		name := reflect.Indirect(hdlr).Type().Name()

		for m := 0; m < typ.NumMethod(); m++ {
			method := typ.Method(m)
			h := &handler{
				method: method.Func,
			}

			switch method.Type.NumIn() {
			case 2:
				h.reqType = method.Type.In(1)
			case 3:
				h.ctxType = method.Type.In(1)
				h.reqType = method.Type.In(2)
			}

			handlers = append(handlers, h)

			endpoints = append(endpoints, &registry.Endpoint{
				Name:    name + "." + method.Name,
				Request: extractSubValue(method.Type),
				Metadata: map[string]string{
					"topic":      topic,
					"subscriber": "true",
				},
			})
		}
	}

	return &subscriber{
		rcvr:       reflect.ValueOf(sub),
		typ:        reflect.TypeOf(sub),
		topic:      topic,
		subscriber: sub,
		handlers:   handlers,
		endpoints:  endpoints,
		opts:       options,
	}
}

func validateSubscriber(sub Subscriber) error {
	typ := reflect.TypeOf(sub.Subscriber())
	var argType reflect.Type

	if typ.Kind() == reflect.Func {
		name := "Func"
		switch typ.NumIn() {
		case 2:
			argType = typ.In(1)
		default:
			return fmt.Errorf("subscriber %v takes wrong number of args: %v required signature %s", name, typ.NumIn(), subSig)
		}
		if !isExportedOrBuiltinType(argType) {
			return fmt.Errorf("subscriber %v argument type not exported: %v", name, argType)
		}
		if typ.NumOut() != 1 {
			return fmt.Errorf("subscriber %v has wrong number of outs: %v require signature %s",
				name, typ.NumOut(), subSig)
		}
		if returnType := typ.Out(0); returnType != typeOfError {
			return fmt.Errorf("subscriber %v returns %v not error", name, returnType.String())
		}
	} else {
		hdlr := reflect.ValueOf(sub.Subscriber())
		name := reflect.Indirect(hdlr).Type().Name()

		for m := 0; m < typ.NumMethod(); m++ {
			method := typ.Method(m)

			switch method.Type.NumIn() {
			case 3:
				argType = method.Type.In(2)
			default:
				return fmt.Errorf("subscriber %v.%v takes wrong number of args: %v required signature %s",
					name, method.Name, method.Type.NumIn(), subSig)
			}

			if !isExportedOrBuiltinType(argType) {
				return fmt.Errorf("%v argument type not exported: %v", name, argType)
			}
			if method.Type.NumOut() != 1 {
				return fmt.Errorf(
					"subscriber %v.%v has wrong number of outs: %v require signature %s",
					name, method.Name, method.Type.NumOut(), subSig)
			}
			if returnType := method.Type.Out(0); returnType != typeOfError {
				return fmt.Errorf("subscriber %v.%v returns %v not error", name, method.Name, returnType.String())
			}
		}
	}

	return nil
}

func (s *rpcServer) createSubHandler(sb *subscriber, opts Options) broker.Handler {
	return func(p broker.Publication) error {
		msg := p.Message()
		ct := msg.Header["Content-Type"]
		cf, err := s.newCodec(ct)
		if err != nil {
			return err
		}

		hdr := make(map[string]string)
		for k, v := range msg.Header {
			hdr[k] = v
		}
		delete(hdr, "Content-Type")
		ctx := metadata.NewContext(context.Background(), hdr)

		for i := 0; i < len(sb.handlers); i++ {
			handler := sb.handlers[i]

			var isVal bool
			var req reflect.Value

			if handler.reqType.Kind() == reflect.Ptr {
				req = reflect.New(handler.reqType.Elem())
			} else {
				req = reflect.New(handler.reqType)
				isVal = true
			}
			if isVal {
				req = req.Elem()
			}

			b := &buffer{bytes.NewBuffer(msg.Body)}
			co := cf(b)
			defer co.Close()

			if err := co.ReadHeader(&codec.Message{}, codec.Publication); err != nil {
				return err
			}

			if err := co.ReadBody(req.Interface()); err != nil {
				return err
			}

			fn := func(ctx context.Context, msg Publication) error {
				var vals []reflect.Value
				if sb.typ.Kind() != reflect.Func {
					vals = append(vals, sb.rcvr)
				}
				if handler.ctxType != nil {
					vals = append(vals, reflect.ValueOf(ctx))
				}

				vals = append(vals, reflect.ValueOf(msg.Message()))

				returnValues := handler.method.Call(vals)
				if err := returnValues[0].Interface(); err != nil {
					return err.(error)
				}
				return nil
			}

			for i := len(opts.SubWrappers); i > 0; i-- {
				fn = opts.SubWrappers[i-1](fn)
			}

			s.wg.Add(1)
			go func() {
				fn(ctx, &rpcPublication{
					topic:       sb.topic,
					contentType: ct,
					message:     req.Interface(),
				})
				s.wg.Done()
			}()
		}
		return nil
	}
}

func (s *subscriber) Topic() string {
	return s.topic
}

func (s *subscriber) Subscriber() interface{} {
	return s.subscriber
}

func (s *subscriber) Endpoints() []*registry.Endpoint {
	return s.endpoints
}

func (s *subscriber) Options() SubscriberOptions {
	return s.opts
}
<|endoftext|>


<commit_before>package registryserver

import (
	"log"
	"time"

	"github.com/go-errors/errors"
	"github.com/golang/protobuf/ptypes/empty"
	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/ligato/networkservicemesh/controlplane/pkg/apis/registry"
	"github.com/ligato/networkservicemesh/k8s/pkg/apis/networkservice/v1"
	nsmClientset "github.com/ligato/networkservicemesh/k8s/pkg/networkservice/clientset/versioned"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	NodeNameLabelKey = "nodeName"
)

type registryService struct {
	clientset *nsmClientset.Clientset
	nsmName   string
}

func (rs registryService) RegisterNSE(ctx context.Context, request *registry.NSERegistration) (*registry.NSERegistration, error) {
	logrus.Infof("Received RegisterNSE(%v)", request)
	// get network service
	networkService, err := rs.clientset.Networkservicemesh().NetworkServices("default").Create(&v1.NetworkService{
		ObjectMeta: metav1.ObjectMeta{
			Name: request.NetworkService.GetName(),
		},
		Spec: v1.NetworkServiceSpec{
			Payload: request.NetworkService.GetPayload(),
		},
		Status: v1.NetworkServiceStatus{},
	})
	if err != nil && !apierrors.IsAlreadyExists(err) {
		return nil, err
	}

	if request.GetNetworkServiceManager().GetUrl() == "" {
		return nil, errors.New("NSERegistration.NetworkServiceManager.Url must be defined")
	}

	nsm := &v1.NetworkServiceManager{
		ObjectMeta: metav1.ObjectMeta{
			Name: rs.nsmName,
		},
		Spec: v1.NetworkServiceManagerSpec{},
		Status: v1.NetworkServiceManagerStatus{
			LastSeen: metav1.Time{Time: time.Now()},
			URL:      request.GetNetworkServiceManager().GetUrl(),
			State:    v1.RUNNING,
		},
	}

	_, err = rs.clientset.Networkservicemesh().NetworkServiceManagers("default").Create(nsm)
	if err != nil && !apierrors.IsAlreadyExists(err) {
		return nil, err
	}

	labels := request.GetNetworkserviceEndpoint().GetLabels()
	if labels == nil {
		labels = make(map[string]string)
	}
	labels["networkservicename"] = request.GetNetworkService().GetName()
	if request.GetNetworkserviceEndpoint() != nil {
		nseResponse, err := rs.clientset.Networkservicemesh().NetworkServiceEndpoints("default").Create(&v1.NetworkServiceEndpoint{
			ObjectMeta: metav1.ObjectMeta{
				GenerateName: request.GetNetworkService().GetName(),
				Labels:       labels,
			},
			Spec: v1.NetworkServiceEndpointSpec{
				NetworkServiceName: request.GetNetworkService().GetName(),
				NsmName:            rs.nsmName,
			},
			Status: v1.NetworkServiceEndpointStatus{
				State: v1.RUNNING,
			},
		})
		if err != nil {
			return nil, err
		}

		request.NetworkserviceEndpoint = &registry.NetworkServiceEndpoint{
			NetworkServiceName:        nseResponse.Spec.NetworkServiceName,
			Payload:                   networkService.Spec.Payload,
			NetworkServiceManagerName: nsm.GetObjectMeta().GetName(),
			EndpointName:              nseResponse.GetObjectMeta().GetName(),
			Labels:                    nseResponse.GetObjectMeta().GetLabels(),
			State:                     string(nseResponse.Status.State),
		}
	}
	logrus.Infof("Returned from RegisterNSE: %v", request)
	return request, nil

}

func (rs registryService) RemoveNSE(ctx context.Context, request *registry.RemoveNSERequest) (*empty.Empty, error) {
	if err := rs.clientset.Networkservicemesh().NetworkServiceEndpoints("default").Delete(request.EndpointName, &metav1.DeleteOptions{}); err != nil {
		return nil, err
	}
	return &empty.Empty{}, nil
}

func (rs registryService) FindNetworkService(ctx context.Context, request *registry.FindNetworkServiceRequest) (*registry.FindNetworkServiceResponse, error) {
	service, e := rs.clientset.Networkservicemesh().NetworkServices("default").Get(request.NetworkServiceName, metav1.GetOptions{})
	if e != nil {
		return nil, e
	}
	payload := service.Spec.Payload

	lo := metav1.ListOptions{}
	lo.LabelSelector = "networkservicename=" + request.NetworkServiceName
	endpointList, e := rs.clientset.Networkservicemesh().NetworkServiceEndpoints("default").List(lo)
	if e != nil {
		return nil, e
	}

	logrus.Println(len(endpointList.Items))
	NSEs := make([]*registry.NetworkServiceEndpoint, len(endpointList.Items))
	NSMs := make(map[string]*registry.NetworkServiceManager)
	for i, endpoint := range endpointList.Items {
		log.Println(endpoint.Name)
		NSEs[i] = &registry.NetworkServiceEndpoint{
			EndpointName:              endpoint.Name,
			NetworkServiceName:        endpoint.Spec.NetworkServiceName,
			NetworkServiceManagerName: endpoint.Spec.NsmName,
			Payload:                   payload,
			Labels:                    endpoint.ObjectMeta.Labels,
		}
		// TODO check in the NSMs map first before potentially looking up again
		manager, e := rs.clientset.Networkservicemesh().NetworkServiceManagers("default").Get(endpoint.Spec.NsmName, metav1.GetOptions{})
		if e != nil {
			return nil, e
		}
		NSMs[endpoint.Spec.NsmName] = &registry.NetworkServiceManager{
			Name: manager.ObjectMeta.Name,
			Url:  manager.Status.URL,
			LastSeen: &timestamp.Timestamp{
				Seconds: manager.Status.LastSeen.ProtoTime().Seconds,
				Nanos:   manager.Status.LastSeen.ProtoTime().Nanos,
			},
		}
	}

	response := &registry.FindNetworkServiceResponse{
		Payload: payload,
		NetworkService: &registry.NetworkService{
			Name:    service.ObjectMeta.Name,
			Payload: service.Spec.Payload,
		},
		NetworkServiceManagers:  NSMs,
		NetworkServiceEndpoints: NSEs,
	}
	return response, nil
}
<commit_msg>Fix bug with retrieving NSM during startup. (#526)<commit_after>package registryserver

import (
	"log"
	"time"

	"github.com/go-errors/errors"
	"github.com/golang/protobuf/ptypes/empty"
	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/ligato/networkservicemesh/controlplane/pkg/apis/registry"
	"github.com/ligato/networkservicemesh/k8s/pkg/apis/networkservice/v1"
	nsmClientset "github.com/ligato/networkservicemesh/k8s/pkg/networkservice/clientset/versioned"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	NodeNameLabelKey = "nodeName"
)

type registryService struct {
	clientset *nsmClientset.Clientset
	nsmName   string
}

func (rs registryService) RegisterNSE(ctx context.Context, request *registry.NSERegistration) (*registry.NSERegistration, error) {
	logrus.Infof("Received RegisterNSE(%v)", request)
	// get network service
	if request.GetNetworkServiceManager().GetUrl() == "" {
		return nil, errors.New("NSERegistration.NetworkServiceManager.Url must be defined")
	}

	nsm := &v1.NetworkServiceManager{
		ObjectMeta: metav1.ObjectMeta{
			Name: rs.nsmName,
		},
		Spec: v1.NetworkServiceManagerSpec{},
		Status: v1.NetworkServiceManagerStatus{
			LastSeen: metav1.Time{Time: time.Now()},
			URL:      request.GetNetworkServiceManager().GetUrl(),
			State:    v1.RUNNING,
		},
	}

	_, err := rs.clientset.Networkservicemesh().NetworkServiceManagers("default").Create(nsm)
	if err != nil && !apierrors.IsAlreadyExists(err) {
		logrus.Errorf("Failed to register nsm: %s", err)
		return nil, err
	}

	labels := request.GetNetworkserviceEndpoint().GetLabels()
	if labels == nil {
		labels = make(map[string]string)
	}
	labels["networkservicename"] = request.GetNetworkService().GetName()
	if request.GetNetworkserviceEndpoint() != nil && request.GetNetworkService() != nil {
		networkService, err := rs.clientset.Networkservicemesh().NetworkServices("default").Create(&v1.NetworkService{
			ObjectMeta: metav1.ObjectMeta{
				Name: request.NetworkService.GetName(),
			},
			Spec: v1.NetworkServiceSpec{
				Payload: request.NetworkService.GetPayload(),
			},
			Status: v1.NetworkServiceStatus{},
		})
		if err != nil && !apierrors.IsAlreadyExists(err) {
			logrus.Errorf("Failed to register nsm: %s", err)
			return nil, err
		}
		nseResponse, err := rs.clientset.Networkservicemesh().NetworkServiceEndpoints("default").Create(&v1.NetworkServiceEndpoint{
			ObjectMeta: metav1.ObjectMeta{
				GenerateName: request.GetNetworkService().GetName(),
				Labels:       labels,
			},
			Spec: v1.NetworkServiceEndpointSpec{
				NetworkServiceName: request.GetNetworkService().GetName(),
				NsmName:            rs.nsmName,
			},
			Status: v1.NetworkServiceEndpointStatus{
				State: v1.RUNNING,
			},
		})
		if err != nil {
			return nil, err
		}

		request.NetworkserviceEndpoint = &registry.NetworkServiceEndpoint{
			NetworkServiceName:        nseResponse.Spec.NetworkServiceName,
			Payload:                   networkService.Spec.Payload,
			NetworkServiceManagerName: nsm.GetObjectMeta().GetName(),
			EndpointName:              nseResponse.GetObjectMeta().GetName(),
			Labels:                    nseResponse.GetObjectMeta().GetLabels(),
			State:                     string(nseResponse.Status.State),
		}
	}
	logrus.Infof("Returned from RegisterNSE: %v", request)
	return request, nil

}

func (rs registryService) RemoveNSE(ctx context.Context, request *registry.RemoveNSERequest) (*empty.Empty, error) {
	if err := rs.clientset.Networkservicemesh().NetworkServiceEndpoints("default").Delete(request.EndpointName, &metav1.DeleteOptions{}); err != nil {
		return nil, err
	}
	return &empty.Empty{}, nil
}

func (rs registryService) FindNetworkService(ctx context.Context, request *registry.FindNetworkServiceRequest) (*registry.FindNetworkServiceResponse, error) {
	service, e := rs.clientset.Networkservicemesh().NetworkServices("default").Get(request.NetworkServiceName, metav1.GetOptions{})
	if e != nil {
		return nil, e
	}
	payload := service.Spec.Payload

	lo := metav1.ListOptions{}
	lo.LabelSelector = "networkservicename=" + request.NetworkServiceName
	endpointList, e := rs.clientset.Networkservicemesh().NetworkServiceEndpoints("default").List(lo)
	if e != nil {
		return nil, e
	}

	logrus.Println(len(endpointList.Items))
	NSEs := make([]*registry.NetworkServiceEndpoint, len(endpointList.Items))
	NSMs := make(map[string]*registry.NetworkServiceManager)
	for i, endpoint := range endpointList.Items {
		log.Println(endpoint.Name)
		NSEs[i] = &registry.NetworkServiceEndpoint{
			EndpointName:              endpoint.Name,
			NetworkServiceName:        endpoint.Spec.NetworkServiceName,
			NetworkServiceManagerName: endpoint.Spec.NsmName,
			Payload:                   payload,
			Labels:                    endpoint.ObjectMeta.Labels,
		}
		// TODO check in the NSMs map first before potentially looking up again
		manager, e := rs.clientset.Networkservicemesh().NetworkServiceManagers("default").Get(endpoint.Spec.NsmName, metav1.GetOptions{})
		if e != nil {
			return nil, e
		}
		NSMs[endpoint.Spec.NsmName] = &registry.NetworkServiceManager{
			Name: manager.ObjectMeta.Name,
			Url:  manager.Status.URL,
			LastSeen: &timestamp.Timestamp{
				Seconds: manager.Status.LastSeen.ProtoTime().Seconds,
				Nanos:   manager.Status.LastSeen.ProtoTime().Nanos,
			},
		}
	}

	response := &registry.FindNetworkServiceResponse{
		Payload: payload,
		NetworkService: &registry.NetworkService{
			Name:    service.ObjectMeta.Name,
			Payload: service.Spec.Payload,
		},
		NetworkServiceManagers:  NSMs,
		NetworkServiceEndpoints: NSEs,
	}
	return response, nil
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package tests_test

import (
	"context"
	"fmt"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	k8sv1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/kubevirt/tests/util"

	cd "kubevirt.io/kubevirt/tests/containerdisk"

	v1 "kubevirt.io/client-go/apis/core/v1"
	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/kubevirt/tests"
	"kubevirt.io/kubevirt/tests/flags"
)

var _ = Describe("[sig-compute]Subresource Api", func() {

	var err error
	var virtCli kubecli.KubevirtClient

	manual := v1.RunStrategyManual
	restartOnError := v1.RunStrategyRerunOnFailure

	BeforeEach(func() {
		virtCli, err = kubecli.GetKubevirtClient()
		util.PanicOnError(err)

		tests.BeforeTestCleanup()
	})

	Describe("[rfe_id:1195][crit:medium][vendor:cnv-qe@redhat.com][level:component] Rbac Authorization", func() {
		var resource string
		BeforeEach(func() {
			vm := tests.NewRandomVMWithEphemeralDisk(cd.ContainerDiskFor(cd.ContainerDiskCirros))
			resource = vm.Name
			vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
		})

		Context("with correct permissions", func() {
			It("[test_id:3170]should be allowed to access subresource endpoint", func() {
				testClientJob(virtCli, true, resource)
			}, 15)
		})
		Context("Without permissions", func() {
			It("[test_id:3171]should not be able to access subresource endpoint", func() {
				testClientJob(virtCli, false, resource)
			}, 15)
		})
	})

	Describe("[rfe_id:1195][crit:medium][vendor:cnv-qe@redhat.com][level:component] Rbac Authorization For Version Command", func() {
		resource := "version"

		Context("with authenticated user", func() {
			It("[test_id:3172]should be allowed to access subresource version endpoint", func() {
				testClientJob(virtCli, true, resource)
			})
		})
		Context("Without permissions", func() {
			It("[test_id:3173]should be able to access subresource version endpoint", func() {
				testClientJob(virtCli, false, resource)
			})
		})
	})

	Describe("[rfe_id:1177][crit:medium][vendor:cnv-qe@redhat.com][level:component] VirtualMachine subresource", func() {
		Context("with a restart endpoint", func() {
			It("[test_id:1304] should restart a VM", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).NotTo(HaveOccurred())

				tests.StartVirtualMachine(vm)
				vmi, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
				Expect(err).NotTo(HaveOccurred())
				Expect(vmi.Status.Phase).To(Equal(v1.Running))

				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).NotTo(HaveOccurred())

				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil || vmi.UID == newVMI.UID {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))
			})

			It("[test_id:1305][posneg:negative] should return an error when VM is not running", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).NotTo(HaveOccurred())

				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).To(HaveOccurred())
			})

			It("[test_id:2265][posneg:negative] should return an error when VM has not been found but VMI is running", func() {
				vmi := tests.NewRandomVMI()
				tests.RunVMIAndExpectLaunch(vmi, 60)

				err := virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vmi.Name)
				Expect(err).To(HaveOccurred())
			})
		})

		Context("With manual RunStrategy", func() {
			It("[test_id:3174]Should not restart when VM is not running", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm.Spec.RunStrategy = &manual
				vm.Spec.Running = nil

				By("Creating VM")
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).ToNot(HaveOccurred())

				By("Trying to start VM via Restart subresource")
				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).To(HaveOccurred())
			})

			It("[test_id:3175]Should restart when VM is running", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm.Spec.RunStrategy = &manual
				vm.Spec.Running = nil

				By("Creating VM")
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).ToNot(HaveOccurred())

				By("Starting VM via Start subresource")
				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Start(vm.Name, &v1.StartOptions{Paused: false})
				Expect(err).ToNot(HaveOccurred())

				By("Waiting for VMI to start")
				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))

				vmi, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
				Expect(err).NotTo(HaveOccurred())

				By("Restarting VM")
				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).ToNot(HaveOccurred())

				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil || vmi.UID == newVMI.UID {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))
			})
		})

		Context("With RunStrategy RerunOnFailure", func() {
			It("[test_id:3176]Should restart the VM", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm.Spec.RunStrategy = &restartOnError
				vm.Spec.Running = nil

				By("Creating VM")
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).ToNot(HaveOccurred())

				By("Waiting for VMI to start")
				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))

				vmi, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
				Expect(err).NotTo(HaveOccurred())

				By("Restarting VM")
				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).ToNot(HaveOccurred())

				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil || vmi.UID == newVMI.UID {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))
			})
		})
	})

	Describe("[rfe_id:1195][crit:medium][vendor:cnv-qe@redhat.com][level:component] the openapi spec for the subresources", func() {
		It("[test_id:3177]should be aggregated into the apiserver openapi spec", func() {
			Eventually(func() string {
				spec, err := virtCli.RestClient().Get().AbsPath("/openapi/v2").DoRaw(context.Background())
				Expect(err).ToNot(HaveOccurred())
				return string(spec)
				// The first item in the SubresourceGroupVersions array is the preferred version
			}, 60*time.Second, 1*time.Second).Should(ContainSubstring("subresources.kubevirt.io/" + v1.SubresourceGroupVersions[0].Version))
		})
	})
})

func testClientJob(virtCli kubecli.KubevirtClient, withServiceAccount bool, resource string) {
	namespace := util.NamespaceTestDefault
	expectedPhase := k8sv1.PodFailed
	name := "subresource-access-tester"
	job := &k8sv1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: name,
			Labels: map[string]string{
				v1.AppLabel: tests.SubresourceTestLabel,
			},
		},
		Spec: k8sv1.PodSpec{
			RestartPolicy: k8sv1.RestartPolicyNever,
			Containers: []k8sv1.Container{
				{
					Name:    name,
					Image:   fmt.Sprintf("%s/subresource-access-test:%s", flags.KubeVirtUtilityRepoPrefix, flags.KubeVirtUtilityVersionTag),
					Command: []string{"/subresource-access-test", "-n", namespace, resource},
				},
			},
		},
	}

	if withServiceAccount {
		job.Spec.ServiceAccountName = tests.SubresourceServiceAccountName
		expectedPhase = k8sv1.PodSucceeded
	} else if resource == "version" {
		expectedPhase = k8sv1.PodSucceeded
	}

	pod, err := virtCli.CoreV1().Pods(namespace).Create(context.Background(), job, metav1.CreateOptions{})
	Expect(err).ToNot(HaveOccurred())

	getStatus := func() k8sv1.PodPhase {
		pod, err := virtCli.CoreV1().Pods(namespace).Get(context.Background(), pod.Name, metav1.GetOptions{})
		Expect(err).ToNot(HaveOccurred())
		return pod.Status.Phase
	}

	Eventually(getStatus, 60, 0.5).Should(Equal(expectedPhase))
}
<commit_msg>Add freeze unfreeze functional tests<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package tests_test

import (
	"context"
	"fmt"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/kubevirt/tests/util"

	cd "kubevirt.io/kubevirt/tests/containerdisk"

	v1 "kubevirt.io/client-go/apis/core/v1"
	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/kubevirt/tests"
	"kubevirt.io/kubevirt/tests/flags"
)

var _ = Describe("[sig-compute]Subresource Api", func() {

	var err error
	var virtCli kubecli.KubevirtClient

	manual := v1.RunStrategyManual
	restartOnError := v1.RunStrategyRerunOnFailure

	BeforeEach(func() {
		virtCli, err = kubecli.GetKubevirtClient()
		util.PanicOnError(err)

		tests.BeforeTestCleanup()
	})

	Describe("[rfe_id:1195][crit:medium][vendor:cnv-qe@redhat.com][level:component] Rbac Authorization", func() {
		var resource string
		BeforeEach(func() {
			vm := tests.NewRandomVMWithEphemeralDisk(cd.ContainerDiskFor(cd.ContainerDiskCirros))
			resource = vm.Name
			vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
		})

		Context("with correct permissions", func() {
			It("[test_id:3170]should be allowed to access subresource endpoint", func() {
				testClientJob(virtCli, true, resource)
			}, 15)
		})
		Context("Without permissions", func() {
			It("[test_id:3171]should not be able to access subresource endpoint", func() {
				testClientJob(virtCli, false, resource)
			}, 15)
		})
	})

	Describe("[rfe_id:1195][crit:medium][vendor:cnv-qe@redhat.com][level:component] Rbac Authorization For Version Command", func() {
		resource := "version"

		Context("with authenticated user", func() {
			It("[test_id:3172]should be allowed to access subresource version endpoint", func() {
				testClientJob(virtCli, true, resource)
			})
		})
		Context("Without permissions", func() {
			It("[test_id:3173]should be able to access subresource version endpoint", func() {
				testClientJob(virtCli, false, resource)
			})
		})
	})

	Describe("[rfe_id:1177][crit:medium][vendor:cnv-qe@redhat.com][level:component] VirtualMachine subresource", func() {
		Context("with a restart endpoint", func() {
			It("[test_id:1304] should restart a VM", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).NotTo(HaveOccurred())

				tests.StartVirtualMachine(vm)
				vmi, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
				Expect(err).NotTo(HaveOccurred())
				Expect(vmi.Status.Phase).To(Equal(v1.Running))

				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).NotTo(HaveOccurred())

				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil || vmi.UID == newVMI.UID {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))
			})

			It("[test_id:1305][posneg:negative] should return an error when VM is not running", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).NotTo(HaveOccurred())

				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).To(HaveOccurred())
			})

			It("[test_id:2265][posneg:negative] should return an error when VM has not been found but VMI is running", func() {
				vmi := tests.NewRandomVMI()
				tests.RunVMIAndExpectLaunch(vmi, 60)

				err := virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vmi.Name)
				Expect(err).To(HaveOccurred())
			})
		})

		Context("With manual RunStrategy", func() {
			It("[test_id:3174]Should not restart when VM is not running", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm.Spec.RunStrategy = &manual
				vm.Spec.Running = nil

				By("Creating VM")
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).ToNot(HaveOccurred())

				By("Trying to start VM via Restart subresource")
				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).To(HaveOccurred())
			})

			It("[test_id:3175]Should restart when VM is running", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm.Spec.RunStrategy = &manual
				vm.Spec.Running = nil

				By("Creating VM")
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).ToNot(HaveOccurred())

				By("Starting VM via Start subresource")
				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Start(vm.Name, &v1.StartOptions{Paused: false})
				Expect(err).ToNot(HaveOccurred())

				By("Waiting for VMI to start")
				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))

				vmi, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
				Expect(err).NotTo(HaveOccurred())

				By("Restarting VM")
				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).ToNot(HaveOccurred())

				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil || vmi.UID == newVMI.UID {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))
			})
		})

		Context("With RunStrategy RerunOnFailure", func() {
			It("[test_id:3176]Should restart the VM", func() {
				vm := tests.NewRandomVirtualMachine(tests.NewRandomVMI(), false)
				vm.Spec.RunStrategy = &restartOnError
				vm.Spec.Running = nil

				By("Creating VM")
				vm, err := virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).ToNot(HaveOccurred())

				By("Waiting for VMI to start")
				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))

				vmi, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
				Expect(err).NotTo(HaveOccurred())

				By("Restarting VM")
				err = virtCli.VirtualMachine(util.NamespaceTestDefault).Restart(vm.Name)
				Expect(err).ToNot(HaveOccurred())

				Eventually(func() v1.VirtualMachineInstancePhase {
					newVMI, err := virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Get(vm.Name, &metav1.GetOptions{})
					if err != nil || vmi.UID == newVMI.UID {
						return v1.VmPhaseUnset
					}
					return newVMI.Status.Phase
				}, 90*time.Second, 1*time.Second).Should(Equal(v1.Running))
			})
		})
	})

	Describe("[rfe_id:1195][crit:medium][vendor:cnv-qe@redhat.com][level:component] the openapi spec for the subresources", func() {
		It("[test_id:3177]should be aggregated into the apiserver openapi spec", func() {
			Eventually(func() string {
				spec, err := virtCli.RestClient().Get().AbsPath("/openapi/v2").DoRaw(context.Background())
				Expect(err).ToNot(HaveOccurred())
				return string(spec)
				// The first item in the SubresourceGroupVersions array is the preferred version
			}, 60*time.Second, 1*time.Second).Should(ContainSubstring("subresources.kubevirt.io/" + v1.SubresourceGroupVersions[0].Version))
		})
	})

	Describe("VirtualMachineInstance subresource", func() {
		Context("Freeze Unfreeze should fail", func() {
			var vm *v1.VirtualMachine

			BeforeEach(func() {
				var err error
				vmiImage := cd.ContainerDiskFor(cd.ContainerDiskCirros)
				vmi := tests.NewRandomVMIWithEphemeralDiskAndUserdata(vmiImage, "#!/bin/bash\necho 'hello'\n")
				vm = tests.NewRandomVirtualMachine(vmi, true)
				vm, err = virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).ToNot(HaveOccurred())
				Eventually(func() bool {
					vmi, err = virtCli.VirtualMachineInstance(vm.Namespace).Get(vm.Name, &metav1.GetOptions{})
					if errors.IsNotFound(err) {
						return false
					}
					Expect(err).ToNot(HaveOccurred())
					return vmi.Status.Phase == v1.Running
				}, 180*time.Second, time.Second).Should(BeTrue())
				tests.WaitForSuccessfulVMIStartWithTimeout(vmi, 180)
			})

			It("Freeze without guest agent", func() {
				expectedErr := "Internal error occurred"
				err = virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Freeze(vm.Name, 0)
				Expect(err).ToNot(BeNil())
				Expect(err.Error()).To(ContainSubstring(expectedErr))
			})

			It("Unfreeze without guest agent", func() {
				expectedErr := "Internal error occurred"
				err = virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Unfreeze(vm.Name)
				Expect(err).ToNot(BeNil())
				Expect(err.Error()).To(ContainSubstring(expectedErr))
			})
		})

		Context("Freeze Unfreeze commands", func() {
			var vm *v1.VirtualMachine

			BeforeEach(func() {
				var err error
				vmi := tests.NewRandomFedoraVMIWithGuestAgent()
				vmi.Namespace = util.NamespaceTestDefault
				vm = tests.NewRandomVirtualMachine(vmi, true)
				vm, err = virtCli.VirtualMachine(util.NamespaceTestDefault).Create(vm)
				Expect(err).ToNot(HaveOccurred())
				Eventually(func() bool {
					vmi, err = virtCli.VirtualMachineInstance(vm.Namespace).Get(vm.Name, &metav1.GetOptions{})
					if errors.IsNotFound(err) {
						return false
					}
					Expect(err).ToNot(HaveOccurred())
					return vmi.Status.Phase == v1.Running
				}, 180*time.Second, time.Second).Should(BeTrue())
				tests.WaitForSuccessfulVMIStartWithTimeout(vmi, 300)
				tests.WaitAgentConnected(virtCli, vmi)
			})

			waitVMIFSFreezeStatus := func(expectedStatus string) {
				Eventually(func() bool {
					updatedVMI, err := virtCli.VirtualMachineInstance(vm.Namespace).Get(vm.Name, &metav1.GetOptions{})
					Expect(err).ToNot(HaveOccurred())
					return updatedVMI.Status.FSFreezeStatus == expectedStatus
				}, 30*time.Second, 2*time.Second).Should(BeTrue())
			}

			It("Freeze Unfreeze should succeed", func() {
				By("Freezing VMI")
				err = virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Freeze(vm.Name, 0)
				Expect(err).ToNot(HaveOccurred())

				waitVMIFSFreezeStatus("frozen")

				By("Unfreezing VMI")
				err = virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Unfreeze(vm.Name)
				Expect(err).ToNot(HaveOccurred())

				waitVMIFSFreezeStatus("")
			})

			It("Multi Freeze Unfreeze calls should succeed", func() {
				for i := 0; i < 5; i++ {
					By("Freezing VMI")
					err = virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Freeze(vm.Name, 0)
					Expect(err).ToNot(HaveOccurred())

					waitVMIFSFreezeStatus("frozen")
				}

				By("Unfreezing VMI")
				for i := 0; i < 5; i++ {
					err = virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Unfreeze(vm.Name)
					Expect(err).ToNot(HaveOccurred())

					waitVMIFSFreezeStatus("")
				}
			})

			It("Freeze without Unfreeze should trigger unfreeze after timeout", func() {
				By("Freezing VMI")
				unfreezeTimeout := 10 * time.Second
				err = virtCli.VirtualMachineInstance(util.NamespaceTestDefault).Freeze(vm.Name, unfreezeTimeout)
				Expect(err).ToNot(HaveOccurred())

				waitVMIFSFreezeStatus("frozen")

				By("Wait Unfreeze VMI to be triggered")
				waitVMIFSFreezeStatus("")
			})
		})
	})
})

func testClientJob(virtCli kubecli.KubevirtClient, withServiceAccount bool, resource string) {
	namespace := util.NamespaceTestDefault
	expectedPhase := k8sv1.PodFailed
	name := "subresource-access-tester"
	job := &k8sv1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: name,
			Labels: map[string]string{
				v1.AppLabel: tests.SubresourceTestLabel,
			},
		},
		Spec: k8sv1.PodSpec{
			RestartPolicy: k8sv1.RestartPolicyNever,
			Containers: []k8sv1.Container{
				{
					Name:    name,
					Image:   fmt.Sprintf("%s/subresource-access-test:%s", flags.KubeVirtUtilityRepoPrefix, flags.KubeVirtUtilityVersionTag),
					Command: []string{"/subresource-access-test", "-n", namespace, resource},
				},
			},
		},
	}

	if withServiceAccount {
		job.Spec.ServiceAccountName = tests.SubresourceServiceAccountName
		expectedPhase = k8sv1.PodSucceeded
	} else if resource == "version" {
		expectedPhase = k8sv1.PodSucceeded
	}

	pod, err := virtCli.CoreV1().Pods(namespace).Create(context.Background(), job, metav1.CreateOptions{})
	Expect(err).ToNot(HaveOccurred())

	getStatus := func() k8sv1.PodPhase {
		pod, err := virtCli.CoreV1().Pods(namespace).Get(context.Background(), pod.Name, metav1.GetOptions{})
		Expect(err).ToNot(HaveOccurred())
		return pod.Status.Phase
	}

	Eventually(getStatus, 60, 0.5).Should(Equal(expectedPhase))
}
<|endoftext|>