<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"
	"text/template"
	"time"

	"github.com/google/go-github/github"
	"golang.org/x/crypto/ssh"
)

var termTmpl = template.Must(template.New("termTmpl").Parse(strings.Replace(`
    +---------------------------------------------------------------------+
    |                                                                     |
    |             _o/ Hello {{ .Name }}!
    |                                                                     |
    |                                                                     |
    |  Did you know that ssh sends all your public keys to any server     |
    |  it tries to authenticate to?                                       |
    |                                                                     |
    |  That's how we know you are @{{ .User }} on GitHub!
    |                                                                     |
    |  Ah, maybe what you didn't know is that GitHub publishes all users' |
    |  ssh public keys and Ben (benjojo.co.uk) grabbed them all.          |
    |                                                                     |
    |  That's pretty handy at times :) for example your key is at         |
    |  https://github.com/{{ .User }}.keys
    |                                                                     |
    |                                                                     |
    |  P.S. This whole thingy is Open Source! (And written in Go!)        |
    |  https://github.com/FiloSottile/whosthere                           |
    |                                                                     |
    |  -- @FiloSottile (https://twitter.com/FiloSottile)                  |
    |                                                                     |
    +---------------------------------------------------------------------+

`, "\n", "\n\r", -1)))

var failedMsg = []byte(strings.Replace(`
    +---------------------------------------------------------------------+
    |                                                                     |
    |             _o/ Hello!                                              |
    |                                                                     |
    |                                                                     |
    |  Did you know that ssh sends all your public keys to any server     |
    |  it tries to authenticate to? You can see yours echoed below.       |
    |                                                                     |
    |  We tried to use that to find your GitHub username, but we          |
    |  couldn't :( maybe you don't even have GitHub ssh keys, do you?     |
    |                                                                     |
    |  By the way, did you know that GitHub publishes all users'          |
    |  ssh public keys and Ben (benjojo.co.uk) grabbed them all?          |
    |                                                                     |
    |  That's pretty handy at times :) But not this time :(               |
    |                                                                     |
    |                                                                     |
    |  P.S. This whole thingy is Open Source! (And written in Go!)        |
    |  https://github.com/FiloSottile/whosthere                           |
    |                                                                     |
    |  -- @FiloSottile (https://twitter.com/FiloSottile)                  |
    |                                                                     |
    +---------------------------------------------------------------------+

`, "\n", "\n\r", -1))

var agentMsg = []byte(strings.Replace(`
                      ***** WARNING ***** WARNING *****

         You have SSH agent forwarding turned (universally?) on. That
        is a VERY BAD idea. For example right now I have access to your
        agent and I can use your keys however I want as long as you are
       connected. I'm a good guy and I won't do anything, but ANY SERVER
        YOU LOG IN TO AND ANYONE WITH ROOT ON THOSE SERVERS CAN LOGIN AS
                                 YOU ANYWHERE.

                       Read more:  http://git.io/vO2A6
`, "\n", "\n\r", -1))

var x11Msg = []byte(strings.Replace(`
                      ***** WARNING ***** WARNING *****

            You have X11 forwarding turned (universally?) on. That
        is a VERY BAD idea. For example right now I have access to your
          X11 server and I can access your desktop as long as you are
       connected. I'm a good guy and I won't do anything, but ANY SERVER
         YOU LOG IN TO AND ANYONE WITH ROOT ON THOSE SERVERS CAN SNIFF
                  YOUR KEYSTROKES AND ACCESS YOUR WINDOWS.

     Read more:  http://www.hackinglinuxexposed.com/articles/20040705.html
`, "\n", "\n\r", -1))

type sessionInfo struct {
	User string
	Keys []ssh.PublicKey
}

type Server struct {
	githubClient *github.Client
	sshConfig    *ssh.ServerConfig
	sqlQuery     *sql.Stmt

	mu          sync.RWMutex
	sessionInfo map[string]sessionInfo
}

func (s *Server) PublicKeyCallback(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	s.mu.Lock()
	si := s.sessionInfo[string(conn.SessionID())]
	si.User = conn.User()
	si.Keys = append(si.Keys, key)
	s.sessionInfo[string(conn.SessionID())] = si
	s.mu.Unlock()

	// Never succeed a key, or we might not see the next. See KeyboardInteractiveCallback.
	return nil, errors.New("")
}

func (s *Server) KeyboardInteractiveCallback(ssh.ConnMetadata, ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
	// keyboard-interactive is tried when all public keys failed, and
	// since it's server-driven we can just pass without user
	// interaction to let the user in once we got all the public keys
	return nil, nil
}

type logEntry struct {
	Timestamp     string
	Username      string
	ChannelTypes  []string
	RequestTypes  []string
	Error         string
	KeysOffered   []string
	GitHub        string
	ClientVersion string
}

func (s *Server) Handle(nConn net.Conn) {
	le := &logEntry{Timestamp: time.Now().Format(time.RFC3339)}
	defer json.NewEncoder(os.Stdout).Encode(le)

	conn, chans, reqs, err := ssh.NewServerConn(nConn, s.sshConfig)
	if err != nil {
		le.Error = "Handshake failed: " + err.Error()
		return
	}
	defer func() {
		s.mu.Lock()
		delete(s.sessionInfo, string(conn.SessionID()))
		s.mu.Unlock()
		conn.Close()
	}()
	go func(in <-chan *ssh.Request) {
		for req := range in {
			le.RequestTypes = append(le.RequestTypes, req.Type)
			if req.WantReply {
				req.Reply(false, nil)
			}
		}
	}(reqs)

	s.mu.RLock()
	si := s.sessionInfo[string(conn.SessionID())]
	s.mu.RUnlock()

	le.Username = conn.User()
	le.ClientVersion = fmt.Sprintf("%x", conn.ClientVersion())
	for _, key := range si.Keys {
		le.KeysOffered = append(le.KeysOffered, string(ssh.MarshalAuthorizedKey(key)))
	}

	for newChannel := range chans {
		le.ChannelTypes = append(le.ChannelTypes, newChannel.ChannelType())

		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		channel, requests, err := newChannel.Accept()
		if err != nil {
			le.Error = "Channel accept failed: " + err.Error()
			continue
		}

		agentFwd, x11 := false, false
		reqLock := &sync.Mutex{}
		reqLock.Lock()
		timeout := time.AfterFunc(30*time.Second, func() { reqLock.Unlock() })

		go func(in <-chan *ssh.Request) {
			for req := range in {
				le.RequestTypes = append(le.RequestTypes, req.Type)
				ok := false
				switch req.Type {
				case "shell":
					fallthrough
				case "pty-req":
					ok = true

					// "auth-agent-req@openssh.com" and "x11-req" always arrive
					// before the "pty-req", so we can go ahead now
					if timeout.Stop() {
						reqLock.Unlock()
					}

				case "auth-agent-req@openssh.com":
					agentFwd = true
				case "x11-req":
					x11 = true
				}

				if req.WantReply {
					req.Reply(ok, nil)
				}
			}
		}(requests)

		reqLock.Lock()
		if agentFwd {
			channel.Write(agentMsg)
		}
		if x11 {
			channel.Write(x11Msg)
		}

		user, err := s.findUser(si.Keys)
		if err != nil {
			le.Error = "findUser failed: " + err.Error()
			channel.Close()
			continue
		}

		if user == "" {
			channel.Write(failedMsg)
			for _, key := range si.Keys {
				channel.Write(ssh.MarshalAuthorizedKey(key))
				channel.Write([]byte("\r"))
			}
			channel.Write([]byte("\n\r"))
			channel.Close()
			continue
		}

		le.GitHub = user
		name, err := s.getUserName(user)
		if err != nil {
			le.Error = "getUserName failed: " + err.Error()
			channel.Close()
			continue
		}

		termTmpl.Execute(channel, struct{ Name, User string }{name, user})

		channel.Close()
	}
}
<commit_msg>Close the connection after servicing a session<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"
	"text/template"
	"time"

	"github.com/google/go-github/github"
	"golang.org/x/crypto/ssh"
)

var termTmpl = template.Must(template.New("termTmpl").Parse(strings.Replace(`
    +---------------------------------------------------------------------+
    |                                                                     |
    |             _o/ Hello {{ .Name }}!
    |                                                                     |
    |                                                                     |
    |  Did you know that ssh sends all your public keys to any server     |
    |  it tries to authenticate to?                                       |
    |                                                                     |
    |  That's how we know you are @{{ .User }} on GitHub!
    |                                                                     |
    |  Ah, maybe what you didn't know is that GitHub publishes all users' |
    |  ssh public keys and Ben (benjojo.co.uk) grabbed them all.          |
    |                                                                     |
    |  That's pretty handy at times :) for example your key is at         |
    |  https://github.com/{{ .User }}.keys
    |                                                                     |
    |                                                                     |
    |  P.S. This whole thingy is Open Source! (And written in Go!)        |
    |  https://github.com/FiloSottile/whosthere                           |
    |                                                                     |
    |  -- @FiloSottile (https://twitter.com/FiloSottile)                  |
    |                                                                     |
    +---------------------------------------------------------------------+

`, "\n", "\n\r", -1)))

var failedMsg = []byte(strings.Replace(`
    +---------------------------------------------------------------------+
    |                                                                     |
    |             _o/ Hello!                                              |
    |                                                                     |
    |                                                                     |
    |  Did you know that ssh sends all your public keys to any server     |
    |  it tries to authenticate to? You can see yours echoed below.       |
    |                                                                     |
    |  We tried to use that to find your GitHub username, but we          |
    |  couldn't :( maybe you don't even have GitHub ssh keys, do you?     |
    |                                                                     |
    |  By the way, did you know that GitHub publishes all users'          |
    |  ssh public keys and Ben (benjojo.co.uk) grabbed them all?          |
    |                                                                     |
    |  That's pretty handy at times :) But not this time :(               |
    |                                                                     |
    |                                                                     |
    |  P.S. This whole thingy is Open Source! (And written in Go!)        |
    |  https://github.com/FiloSottile/whosthere                           |
    |                                                                     |
    |  -- @FiloSottile (https://twitter.com/FiloSottile)                  |
    |                                                                     |
    +---------------------------------------------------------------------+

`, "\n", "\n\r", -1))

var agentMsg = []byte(strings.Replace(`
                      ***** WARNING ***** WARNING *****

         You have SSH agent forwarding turned (universally?) on. That
        is a VERY BAD idea. For example right now I have access to your
        agent and I can use your keys however I want as long as you are
       connected. I'm a good guy and I won't do anything, but ANY SERVER
        YOU LOG IN TO AND ANYONE WITH ROOT ON THOSE SERVERS CAN LOGIN AS
                                 YOU ANYWHERE.

                       Read more:  http://git.io/vO2A6
`, "\n", "\n\r", -1))

var x11Msg = []byte(strings.Replace(`
                      ***** WARNING ***** WARNING *****

            You have X11 forwarding turned (universally?) on. That
        is a VERY BAD idea. For example right now I have access to your
          X11 server and I can access your desktop as long as you are
       connected. I'm a good guy and I won't do anything, but ANY SERVER
         YOU LOG IN TO AND ANYONE WITH ROOT ON THOSE SERVERS CAN SNIFF
                  YOUR KEYSTROKES AND ACCESS YOUR WINDOWS.

     Read more:  http://www.hackinglinuxexposed.com/articles/20040705.html
`, "\n", "\n\r", -1))

type sessionInfo struct {
	User string
	Keys []ssh.PublicKey
}

type Server struct {
	githubClient *github.Client
	sshConfig    *ssh.ServerConfig
	sqlQuery     *sql.Stmt

	mu          sync.RWMutex
	sessionInfo map[string]sessionInfo
}

func (s *Server) PublicKeyCallback(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	s.mu.Lock()
	si := s.sessionInfo[string(conn.SessionID())]
	si.User = conn.User()
	si.Keys = append(si.Keys, key)
	s.sessionInfo[string(conn.SessionID())] = si
	s.mu.Unlock()

	// Never succeed a key, or we might not see the next. See KeyboardInteractiveCallback.
	return nil, errors.New("")
}

func (s *Server) KeyboardInteractiveCallback(ssh.ConnMetadata, ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
	// keyboard-interactive is tried when all public keys failed, and
	// since it's server-driven we can just pass without user
	// interaction to let the user in once we got all the public keys
	return nil, nil
}

type logEntry struct {
	Timestamp     string
	Username      string
	RequestTypes  []string
	Error         string
	KeysOffered   []string
	GitHub        string
	ClientVersion string
}

func (s *Server) Handle(nConn net.Conn) {
	le := &logEntry{Timestamp: time.Now().Format(time.RFC3339)}
	defer json.NewEncoder(os.Stdout).Encode(le)

	conn, chans, reqs, err := ssh.NewServerConn(nConn, s.sshConfig)
	if err != nil {
		le.Error = "Handshake failed: " + err.Error()
		nConn.Close()
		return
	}
	defer func() {
		s.mu.Lock()
		delete(s.sessionInfo, string(conn.SessionID()))
		s.mu.Unlock()
		conn.Close()
	}()
	go func(in <-chan *ssh.Request) {
		for req := range in {
			le.RequestTypes = append(le.RequestTypes, req.Type)
			if req.WantReply {
				req.Reply(false, nil)
			}
		}
	}(reqs)

	s.mu.RLock()
	si := s.sessionInfo[string(conn.SessionID())]
	s.mu.RUnlock()

	le.Username = conn.User()
	le.ClientVersion = fmt.Sprintf("%x", conn.ClientVersion())
	for _, key := range si.Keys {
		le.KeysOffered = append(le.KeysOffered, string(ssh.MarshalAuthorizedKey(key)))
	}

	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		channel, requests, err := newChannel.Accept()
		if err != nil {
			le.Error = "Channel accept failed: " + err.Error()
			return
		}
		defer channel.Close()

		agentFwd, x11 := false, false
		reqLock := &sync.Mutex{}
		reqLock.Lock()
		timeout := time.AfterFunc(30*time.Second, func() { reqLock.Unlock() })

		go func(in <-chan *ssh.Request) {
			for req := range in {
				le.RequestTypes = append(le.RequestTypes, req.Type)
				ok := false
				switch req.Type {
				case "shell":
					fallthrough
				case "pty-req":
					ok = true

					// "auth-agent-req@openssh.com" and "x11-req" always arrive
					// before the "pty-req", so we can go ahead now
					if timeout.Stop() {
						reqLock.Unlock()
					}

				case "auth-agent-req@openssh.com":
					agentFwd = true
				case "x11-req":
					x11 = true
				}

				if req.WantReply {
					req.Reply(ok, nil)
				}
			}
		}(requests)

		reqLock.Lock()
		if agentFwd {
			channel.Write(agentMsg)
		}
		if x11 {
			channel.Write(x11Msg)
		}

		user, err := s.findUser(si.Keys)
		if err != nil {
			le.Error = "findUser failed: " + err.Error()
			return
		}

		if user == "" {
			channel.Write(failedMsg)
			for _, key := range si.Keys {
				channel.Write(ssh.MarshalAuthorizedKey(key))
				channel.Write([]byte("\r"))
			}
			channel.Write([]byte("\n\r"))
			return
		}

		le.GitHub = user
		name, err := s.getUserName(user)
		if err != nil {
			le.Error = "getUserName failed: " + err.Error()
			return
		}

		termTmpl.Execute(channel, struct{ Name, User string }{name, user})
		return
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	uuid "github.com/nu7hatch/gouuid"
	"io/ioutil"
	"net/http"
	"os/user"
	"strings"
	"time"
)

// var authSite = "https://koding.com"
var authSite = "http://localhost:3020"

type Register struct {
	ID         string
	PublicKey  string
	PrivateKey string
	Username   string
}

func NewRegister() *Register {
	id, _ := uuid.NewV4()
	return &Register{
		ID: "machineID-" + id.String(),
	}
}

func (r *Register) Definition() string {
	return "Registers the user's machine to kontrol"
}

func (r *Register) Exec() error {
	err := r.Register()
	if err != nil {
		return errors.New(fmt.Sprint("\nregister error: %s\n", err.Error()))
	}
	return nil
}

func (r *Register) Register() error {
	err := r.getConfig()
	if err != nil {
		return err
	}

	registerUrl := fmt.Sprintf("%s/-/auth/register/%s/%s", authSite, r.ID, r.PublicKey)
	checkUrl := fmt.Sprintf("%s/-/auth/check/%s", authSite, r.PublicKey)

	fmt.Printf("Please open the following url for authentication:\n\n")
	// PrintBox(registerUrl)
	fmt.Println(registerUrl)
	fmt.Printf("\nwaiting . ")

	// check the result every two seconds
	ticker := time.NewTicker(time.Second * 2).C

	// wait for three minutes, if not successfull abort it
	timeout := time.After(time.Minute * 3)

	for {
		select {
		case <-ticker:
			resp, err := http.Get(checkUrl)
			if err != nil {
				return err
			}

			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}

			resp.Body.Close()
			fmt.Printf(". ")

			if resp.StatusCode == 200 {
				type Result struct {
					Result string `json:"result"`
				}

				res := Result{}

				err := json.Unmarshal(bytes.TrimSpace(body), &res)
				if err != nil {
					return err
				}

				fmt.Println(res.Result)
				return nil
			}
		case <-timeout:
			return errors.New("timeout")
		}
	}

	return nil
}

func PrintBox(msg string) {
	fmt.Printf("\n\n")
	fmt.Println("╔" + strings.Repeat("═", len(msg)+2) + "╗")
	fmt.Println("║" + strings.Repeat(" ", len(msg)+2) + "║")
	fmt.Println("║ " + msg + " ║")
	fmt.Println("║" + strings.Repeat(" ", len(msg)+2) + "║")
	fmt.Println("╚" + strings.Repeat("═", len(msg)+2) + "╝")
	fmt.Printf("\n")
}

func (r *Register) getConfig() error {
	var err error
	// we except to read all of them, if one fails it should be abort
	r.PublicKey, err = getKey("public")
	if err != nil {
		return err
	}

	r.PrivateKey, err = getKey("private")
	if err != nil {
		return err
	}
	return nil
}

func getKey(key string) (string, error) {
	usr, err := user.Current()
	if err != nil {
		return "", err
	}

	var keyfile string
	switch key {
	case "public":
		keyfile = usr.HomeDir + "/.kd/koding.key.pub"
	case "private":
		keyfile = usr.HomeDir + "/.kd/koding.key"
	default:
		return "", fmt.Errorf("key is not recognized '%s'\n", key)
	}

	file, err := ioutil.ReadFile(keyfile)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(file)), nil
}
<commit_msg>do not eat errors<commit_after>package cli

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	uuid "github.com/nu7hatch/gouuid"
	"io/ioutil"
	"net/http"
	"os/user"
	"strings"
	"time"
)

// var authSite = "https://koding.com"
var authSite = "http://localhost:3020"

type Register struct {
	ID         string
	PublicKey  string
	PrivateKey string
	Username   string
}

func NewRegister() *Register {
	id, err := uuid.NewV4()
	if err != nil {
		panic(err)
	}

	return &Register{
		ID: "machineID-" + id.String(),
	}
}

func (r *Register) Definition() string {
	return "Registers the user's machine to kontrol"
}

func (r *Register) Exec() error {
	err := r.Register()
	if err != nil {
		return errors.New(fmt.Sprint("\nregister error: %s\n", err.Error()))
	}
	return nil
}

func (r *Register) Register() error {
	err := r.getConfig()
	if err != nil {
		return err
	}

	registerUrl := fmt.Sprintf("%s/-/auth/register/%s/%s", authSite, r.ID, r.PublicKey)
	checkUrl := fmt.Sprintf("%s/-/auth/check/%s", authSite, r.PublicKey)

	fmt.Printf("Please open the following url for authentication:\n\n")
	// PrintBox(registerUrl)
	fmt.Println(registerUrl)
	fmt.Printf("\nwaiting . ")

	// check the result every two seconds
	ticker := time.NewTicker(time.Second * 2).C

	// wait for three minutes, if not successfull abort it
	timeout := time.After(time.Minute * 3)

	for {
		select {
		case <-ticker:
			resp, err := http.Get(checkUrl)
			if err != nil {
				return err
			}

			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}

			resp.Body.Close()
			fmt.Printf(". ")

			if resp.StatusCode == 200 {
				type Result struct {
					Result string `json:"result"`
				}

				res := Result{}

				err := json.Unmarshal(bytes.TrimSpace(body), &res)
				if err != nil {
					return err
				}

				fmt.Println(res.Result)
				return nil
			}
		case <-timeout:
			return errors.New("timeout")
		}
	}

	return nil
}

func PrintBox(msg string) {
	fmt.Printf("\n\n")
	fmt.Println("╔" + strings.Repeat("═", len(msg)+2) + "╗")
	fmt.Println("║" + strings.Repeat(" ", len(msg)+2) + "║")
	fmt.Println("║ " + msg + " ║")
	fmt.Println("║" + strings.Repeat(" ", len(msg)+2) + "║")
	fmt.Println("╚" + strings.Repeat("═", len(msg)+2) + "╝")
	fmt.Printf("\n")
}

func (r *Register) getConfig() error {
	var err error
	// we except to read all of them, if one fails it should be abort
	r.PublicKey, err = getKey("public")
	if err != nil {
		return err
	}

	r.PrivateKey, err = getKey("private")
	if err != nil {
		return err
	}
	return nil
}

func getKey(key string) (string, error) {
	usr, err := user.Current()
	if err != nil {
		return "", err
	}

	var keyfile string
	switch key {
	case "public":
		keyfile = usr.HomeDir + "/.kd/koding.key.pub"
	case "private":
		keyfile = usr.HomeDir + "/.kd/koding.key"
	default:
		return "", fmt.Errorf("key is not recognized '%s'\n", key)
	}

	file, err := ioutil.ReadFile(keyfile)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(file)), nil
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"fmt"

	apiapps "k8s.io/api/apps/v1beta1"
	apibatch "k8s.io/api/batch/v1beta1"
	apiv1 "k8s.io/api/core/v1"
	apiext "k8s.io/api/extensions/v1beta1"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/weaveworks/flux"
	fhr_v1alpha2 "github.com/weaveworks/flux/apis/helm.integrations.flux.weave.works/v1alpha2"
	"github.com/weaveworks/flux/cluster"
	kresource "github.com/weaveworks/flux/cluster/kubernetes/resource"
	"github.com/weaveworks/flux/image"
	"github.com/weaveworks/flux/resource"
)

// AntecedentAnnotation is an annotation on a resource indicating that
// the cause of that resource (indirectly, via a Helm release) is a
// FluxHelmRelease. We use this rather than the `OwnerReference` type
// built into Kubernetes so that there are no garbage-collection
// implications. The value is expected to be a serialised
// `flux.ResourceID`.
const AntecedentAnnotation = "flux.weave.works/antecedent"

/////////////////////////////////////////////////////////////////////////////
// Kind registry

type resourceKind interface {
	getPodController(c *Cluster, namespace, name string) (podController, error)
	getPodControllers(c *Cluster, namespace string) ([]podController, error)
}

var (
	resourceKinds = make(map[string]resourceKind)
)

func init() {
	resourceKinds["cronjob"] = &cronJobKind{}
	resourceKinds["daemonset"] = &daemonSetKind{}
	resourceKinds["deployment"] = &deploymentKind{}
	resourceKinds["statefulset"] = &statefulSetKind{}
	resourceKinds["fluxhelmrelease"] = &fluxHelmReleaseKind{}
}

type podController struct {
	k8sObject
	apiVersion  string
	kind        string
	name        string
	status      string
	podTemplate apiv1.PodTemplateSpec
}

func (pc podController) toClusterController(resourceID flux.ResourceID) cluster.Controller {
	var clusterContainers []resource.Container
	var excuse string
	for _, container := range pc.podTemplate.Spec.Containers {
		ref, err := image.ParseRef(container.Image)
		if err != nil {
			clusterContainers = nil
			excuse = err.Error()
			break
		}
		clusterContainers = append(clusterContainers, resource.Container{Name: container.Name, Image: ref})
	}

	var antecedent flux.ResourceID
	if ante, ok := pc.GetAnnotations()[AntecedentAnnotation]; ok {
		id, err := flux.ParseResourceID(ante)
		if err == nil {
			antecedent = id
		}
	}

	return cluster.Controller{
		ID:         resourceID,
		Status:     pc.status,
		Antecedent: antecedent,
		Labels:     pc.GetLabels(),
		Containers: cluster.ContainersOrExcuse{Containers: clusterContainers, Excuse: excuse},
	}
}

/////////////////////////////////////////////////////////////////////////////
// extensions/v1beta1 Deployment

type deploymentKind struct{}

func (dk *deploymentKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	deployment, err := c.client.Deployments(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeDeploymentPodController(deployment), nil
}

func (dk *deploymentKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	deployments, err := c.client.Deployments(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for i := range deployments.Items {
		podControllers = append(podControllers, makeDeploymentPodController(&deployments.Items[i]))
	}

	return podControllers, nil
}

func makeDeploymentPodController(deployment *apiext.Deployment) podController {
	var status string
	objectMeta, deploymentStatus := deployment.ObjectMeta, deployment.Status

	if deploymentStatus.ObservedGeneration >= objectMeta.Generation {
		// the definition has been updated; now let's see about the replicas
		updated, wanted := deploymentStatus.UpdatedReplicas, *deployment.Spec.Replicas
		if updated == wanted {
			status = StatusReady
		} else {
			status = fmt.Sprintf("%d out of %d updated", updated, wanted)
		}
	} else {
		status = StatusUpdating
	}

	return podController{
		apiVersion:  "extensions/v1beta1",
		kind:        "Deployment",
		name:        deployment.ObjectMeta.Name,
		status:      status,
		podTemplate: deployment.Spec.Template,
		k8sObject:   deployment}
}

/////////////////////////////////////////////////////////////////////////////
// extensions/v1beta daemonset

type daemonSetKind struct{}

func (dk *daemonSetKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	daemonSet, err := c.client.DaemonSets(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeDaemonSetPodController(daemonSet), nil
}

func (dk *daemonSetKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	daemonSets, err := c.client.DaemonSets(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for i, _ := range daemonSets.Items {
		podControllers = append(podControllers, makeDaemonSetPodController(&daemonSets.Items[i]))
	}

	return podControllers, nil
}

func makeDaemonSetPodController(daemonSet *apiext.DaemonSet) podController {
	var status string
	objectMeta, daemonSetStatus := daemonSet.ObjectMeta, daemonSet.Status
	if daemonSetStatus.ObservedGeneration >= objectMeta.Generation {
		// the definition has been updated; now let's see about the replicas
		updated, wanted := daemonSetStatus.UpdatedNumberScheduled, daemonSetStatus.DesiredNumberScheduled
		if updated == wanted {
			status = StatusReady
		} else {
			status = fmt.Sprintf("%d out of %d updated", updated, wanted)
		}
	} else {
		status = StatusUpdating
	}

	return podController{
		apiVersion:  "extensions/v1beta1",
		kind:        "DaemonSet",
		name:        daemonSet.ObjectMeta.Name,
		status:      status,
		podTemplate: daemonSet.Spec.Template,
		k8sObject:   daemonSet}
}

/////////////////////////////////////////////////////////////////////////////
// apps/v1beta1 StatefulSet

type statefulSetKind struct{}

func (dk *statefulSetKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	statefulSet, err := c.client.StatefulSets(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeStatefulSetPodController(statefulSet), nil
}

func (dk *statefulSetKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	statefulSets, err := c.client.StatefulSets(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for i, _ := range statefulSets.Items {
		podControllers = append(podControllers, makeStatefulSetPodController(&statefulSets.Items[i]))
	}

	return podControllers, nil
}

func makeStatefulSetPodController(statefulSet *apiapps.StatefulSet) podController {
	var status string
	objectMeta, statefulSetStatus := statefulSet.ObjectMeta, statefulSet.Status
	// The type of ObservedGeneration is *int64, unlike other controllers.
	if statefulSetStatus.ObservedGeneration != nil && *statefulSetStatus.ObservedGeneration >= objectMeta.Generation {
		// the definition has been updated; now let's see about the replicas
		updated, wanted := statefulSetStatus.UpdatedReplicas, *statefulSet.Spec.Replicas
		if updated == wanted {
			status = StatusReady
		} else {
			status = fmt.Sprintf("%d out of %d updated", updated, wanted)
		}
	} else {
		status = StatusUpdating
	}

	return podController{
		apiVersion:  "apps/v1beta1",
		kind:        "StatefulSet",
		name:        statefulSet.ObjectMeta.Name,
		status:      status,
		podTemplate: statefulSet.Spec.Template,
		k8sObject:   statefulSet}
}

/////////////////////////////////////////////////////////////////////////////
// batch/v1beta1 CronJob

type cronJobKind struct{}

func (dk *cronJobKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	cronJob, err := c.client.CronJobs(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeCronJobPodController(cronJob), nil
}

func (dk *cronJobKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	cronJobs, err := c.client.CronJobs(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for i, _ := range cronJobs.Items {
		podControllers = append(podControllers, makeCronJobPodController(&cronJobs.Items[i]))
	}

	return podControllers, nil
}

func makeCronJobPodController(cronJob *apibatch.CronJob) podController {
	return podController{
		apiVersion:  "batch/v1beta1",
		kind:        "CronJob",
		name:        cronJob.ObjectMeta.Name,
		status:      StatusReady,
		podTemplate: cronJob.Spec.JobTemplate.Spec.Template,
		k8sObject:   cronJob}
}

/////////////////////////////////////////////////////////////////////////////
// helm.integrations.flux.weave.works/v1alpha2 FluxHelmRelease

type fluxHelmReleaseKind struct{}

func (fhr *fluxHelmReleaseKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	fluxHelmRelease, err := c.client.HelmV1alpha2().FluxHelmReleases(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeFluxHelmReleasePodController(fluxHelmRelease), nil
}

func (fhr *fluxHelmReleaseKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	fluxHelmReleases, err := c.client.HelmV1alpha2().FluxHelmReleases(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for _, f := range fluxHelmReleases.Items {
		podControllers = append(podControllers, makeFluxHelmReleasePodController(&f))
	}

	return podControllers, nil
}

func makeFluxHelmReleasePodController(fluxHelmRelease *fhr_v1alpha2.FluxHelmRelease) podController {
	containers := createK8sFHRContainers(fluxHelmRelease.Spec)

	podTemplate := apiv1.PodTemplateSpec{
		ObjectMeta: fluxHelmRelease.ObjectMeta,
		Spec: apiv1.PodSpec{
			Containers:       containers,
			ImagePullSecrets: []apiv1.LocalObjectReference{},
		},
	}

	return podController{
		apiVersion:  "helm.integrations.flux.weave.works/v1alpha2",
		kind:        "FluxHelmRelease",
		name:        fluxHelmRelease.ObjectMeta.Name,
		status:      fluxHelmRelease.Status.ReleaseStatus,
		podTemplate: podTemplate,
		k8sObject:   fluxHelmRelease,
	}
}

// createK8sContainers creates a list of k8s containers by
// interpreting the FluxHelmRelease resource. The interpretation is
// analogous to that in cluster/kubernetes/resource/fluxhelmrelease.go
func createK8sFHRContainers(spec fhr_v1alpha2.FluxHelmReleaseSpec) []apiv1.Container {
	var containers []apiv1.Container
	_ = kresource.FindFluxHelmReleaseContainers(spec.Values, func(name string, image image.Ref, _ kresource.ImageSetter) error {
		containers = append(containers, apiv1.Container{
			Name:  name,
			Image: image.String(),
		})
		return nil
	})
	return containers
}
<commit_msg>Include initContainers in containers considered<commit_after>package kubernetes

import (
	"fmt"

	apiapps "k8s.io/api/apps/v1beta1"
	apibatch "k8s.io/api/batch/v1beta1"
	apiv1 "k8s.io/api/core/v1"
	apiext "k8s.io/api/extensions/v1beta1"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/weaveworks/flux"
	fhr_v1alpha2 "github.com/weaveworks/flux/apis/helm.integrations.flux.weave.works/v1alpha2"
	"github.com/weaveworks/flux/cluster"
	kresource "github.com/weaveworks/flux/cluster/kubernetes/resource"
	"github.com/weaveworks/flux/image"
	"github.com/weaveworks/flux/resource"
)

// AntecedentAnnotation is an annotation on a resource indicating that
// the cause of that resource (indirectly, via a Helm release) is a
// FluxHelmRelease. We use this rather than the `OwnerReference` type
// built into Kubernetes so that there are no garbage-collection
// implications. The value is expected to be a serialised
// `flux.ResourceID`.
const AntecedentAnnotation = "flux.weave.works/antecedent"

/////////////////////////////////////////////////////////////////////////////
// Kind registry

type resourceKind interface {
	getPodController(c *Cluster, namespace, name string) (podController, error)
	getPodControllers(c *Cluster, namespace string) ([]podController, error)
}

var (
	resourceKinds = make(map[string]resourceKind)
)

func init() {
	resourceKinds["cronjob"] = &cronJobKind{}
	resourceKinds["daemonset"] = &daemonSetKind{}
	resourceKinds["deployment"] = &deploymentKind{}
	resourceKinds["statefulset"] = &statefulSetKind{}
	resourceKinds["fluxhelmrelease"] = &fluxHelmReleaseKind{}
}

type podController struct {
	k8sObject
	apiVersion  string
	kind        string
	name        string
	status      string
	podTemplate apiv1.PodTemplateSpec
}

func (pc podController) toClusterController(resourceID flux.ResourceID) cluster.Controller {
	var clusterContainers []resource.Container
	var excuse string
	for _, container := range pc.podTemplate.Spec.Containers {
		ref, err := image.ParseRef(container.Image)
		if err != nil {
			clusterContainers = nil
			excuse = err.Error()
			break
		}
		clusterContainers = append(clusterContainers, resource.Container{Name: container.Name, Image: ref})
	}
	for _, container := range pc.podTemplate.Spec.Containers {
		ref, err := image.ParseRef(container.Image)
		if err != nil {
			clusterContainers = nil
			excuse = err.Error()
			break
		}
		clusterContainers = append(clusterContainers, resource.Container{Name: container.Name, Image: ref})
	}

	var antecedent flux.ResourceID
	if ante, ok := pc.GetAnnotations()[AntecedentAnnotation]; ok {
		id, err := flux.ParseResourceID(ante)
		if err == nil {
			antecedent = id
		}
	}

	return cluster.Controller{
		ID:         resourceID,
		Status:     pc.status,
		Antecedent: antecedent,
		Labels:     pc.GetLabels(),
		Containers: cluster.ContainersOrExcuse{Containers: clusterContainers, Excuse: excuse},
	}
}

/////////////////////////////////////////////////////////////////////////////
// extensions/v1beta1 Deployment

type deploymentKind struct{}

func (dk *deploymentKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	deployment, err := c.client.Deployments(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeDeploymentPodController(deployment), nil
}

func (dk *deploymentKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	deployments, err := c.client.Deployments(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for i := range deployments.Items {
		podControllers = append(podControllers, makeDeploymentPodController(&deployments.Items[i]))
	}

	return podControllers, nil
}

func makeDeploymentPodController(deployment *apiext.Deployment) podController {
	var status string
	objectMeta, deploymentStatus := deployment.ObjectMeta, deployment.Status

	if deploymentStatus.ObservedGeneration >= objectMeta.Generation {
		// the definition has been updated; now let's see about the replicas
		updated, wanted := deploymentStatus.UpdatedReplicas, *deployment.Spec.Replicas
		if updated == wanted {
			status = StatusReady
		} else {
			status = fmt.Sprintf("%d out of %d updated", updated, wanted)
		}
	} else {
		status = StatusUpdating
	}

	return podController{
		apiVersion:  "extensions/v1beta1",
		kind:        "Deployment",
		name:        deployment.ObjectMeta.Name,
		status:      status,
		podTemplate: deployment.Spec.Template,
		k8sObject:   deployment}
}

/////////////////////////////////////////////////////////////////////////////
// extensions/v1beta daemonset

type daemonSetKind struct{}

func (dk *daemonSetKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	daemonSet, err := c.client.DaemonSets(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeDaemonSetPodController(daemonSet), nil
}

func (dk *daemonSetKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	daemonSets, err := c.client.DaemonSets(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for i, _ := range daemonSets.Items {
		podControllers = append(podControllers, makeDaemonSetPodController(&daemonSets.Items[i]))
	}

	return podControllers, nil
}

func makeDaemonSetPodController(daemonSet *apiext.DaemonSet) podController {
	var status string
	objectMeta, daemonSetStatus := daemonSet.ObjectMeta, daemonSet.Status
	if daemonSetStatus.ObservedGeneration >= objectMeta.Generation {
		// the definition has been updated; now let's see about the replicas
		updated, wanted := daemonSetStatus.UpdatedNumberScheduled, daemonSetStatus.DesiredNumberScheduled
		if updated == wanted {
			status = StatusReady
		} else {
			status = fmt.Sprintf("%d out of %d updated", updated, wanted)
		}
	} else {
		status = StatusUpdating
	}

	return podController{
		apiVersion:  "extensions/v1beta1",
		kind:        "DaemonSet",
		name:        daemonSet.ObjectMeta.Name,
		status:      status,
		podTemplate: daemonSet.Spec.Template,
		k8sObject:   daemonSet}
}

/////////////////////////////////////////////////////////////////////////////
// apps/v1beta1 StatefulSet

type statefulSetKind struct{}

func (dk *statefulSetKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	statefulSet, err := c.client.StatefulSets(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeStatefulSetPodController(statefulSet), nil
}

func (dk *statefulSetKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	statefulSets, err := c.client.StatefulSets(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for i, _ := range statefulSets.Items {
		podControllers = append(podControllers, makeStatefulSetPodController(&statefulSets.Items[i]))
	}

	return podControllers, nil
}

func makeStatefulSetPodController(statefulSet *apiapps.StatefulSet) podController {
	var status string
	objectMeta, statefulSetStatus := statefulSet.ObjectMeta, statefulSet.Status
	// The type of ObservedGeneration is *int64, unlike other controllers.
	if statefulSetStatus.ObservedGeneration != nil && *statefulSetStatus.ObservedGeneration >= objectMeta.Generation {
		// the definition has been updated; now let's see about the replicas
		updated, wanted := statefulSetStatus.UpdatedReplicas, *statefulSet.Spec.Replicas
		if updated == wanted {
			status = StatusReady
		} else {
			status = fmt.Sprintf("%d out of %d updated", updated, wanted)
		}
	} else {
		status = StatusUpdating
	}

	return podController{
		apiVersion:  "apps/v1beta1",
		kind:        "StatefulSet",
		name:        statefulSet.ObjectMeta.Name,
		status:      status,
		podTemplate: statefulSet.Spec.Template,
		k8sObject:   statefulSet}
}

/////////////////////////////////////////////////////////////////////////////
// batch/v1beta1 CronJob

type cronJobKind struct{}

func (dk *cronJobKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	cronJob, err := c.client.CronJobs(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeCronJobPodController(cronJob), nil
}

func (dk *cronJobKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	cronJobs, err := c.client.CronJobs(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for i, _ := range cronJobs.Items {
		podControllers = append(podControllers, makeCronJobPodController(&cronJobs.Items[i]))
	}

	return podControllers, nil
}

func makeCronJobPodController(cronJob *apibatch.CronJob) podController {
	return podController{
		apiVersion:  "batch/v1beta1",
		kind:        "CronJob",
		name:        cronJob.ObjectMeta.Name,
		status:      StatusReady,
		podTemplate: cronJob.Spec.JobTemplate.Spec.Template,
		k8sObject:   cronJob}
}

/////////////////////////////////////////////////////////////////////////////
// helm.integrations.flux.weave.works/v1alpha2 FluxHelmRelease

type fluxHelmReleaseKind struct{}

func (fhr *fluxHelmReleaseKind) getPodController(c *Cluster, namespace, name string) (podController, error) {
	fluxHelmRelease, err := c.client.HelmV1alpha2().FluxHelmReleases(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		return podController{}, err
	}

	return makeFluxHelmReleasePodController(fluxHelmRelease), nil
}

func (fhr *fluxHelmReleaseKind) getPodControllers(c *Cluster, namespace string) ([]podController, error) {
	fluxHelmReleases, err := c.client.HelmV1alpha2().FluxHelmReleases(namespace).List(meta_v1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var podControllers []podController
	for _, f := range fluxHelmReleases.Items {
		podControllers = append(podControllers, makeFluxHelmReleasePodController(&f))
	}

	return podControllers, nil
}

func makeFluxHelmReleasePodController(fluxHelmRelease *fhr_v1alpha2.FluxHelmRelease) podController {
	containers := createK8sFHRContainers(fluxHelmRelease.Spec)

	podTemplate := apiv1.PodTemplateSpec{
		ObjectMeta: fluxHelmRelease.ObjectMeta,
		Spec: apiv1.PodSpec{
			Containers:       containers,
			ImagePullSecrets: []apiv1.LocalObjectReference{},
		},
	}

	return podController{
		apiVersion:  "helm.integrations.flux.weave.works/v1alpha2",
		kind:        "FluxHelmRelease",
		name:        fluxHelmRelease.ObjectMeta.Name,
		status:      fluxHelmRelease.Status.ReleaseStatus,
		podTemplate: podTemplate,
		k8sObject:   fluxHelmRelease,
	}
}

// createK8sContainers creates a list of k8s containers by
// interpreting the FluxHelmRelease resource. The interpretation is
// analogous to that in cluster/kubernetes/resource/fluxhelmrelease.go
func createK8sFHRContainers(spec fhr_v1alpha2.FluxHelmReleaseSpec) []apiv1.Container {
	var containers []apiv1.Container
	_ = kresource.FindFluxHelmReleaseContainers(spec.Values, func(name string, image image.Ref, _ kresource.ImageSetter) error {
		containers = append(containers, apiv1.Container{
			Name:  name,
			Image: image.String(),
		})
		return nil
	})
	return containers
}
<|endoftext|>


<commit_before>/* Copyright 2020 Google Inc.

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

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

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

package config

import (
	"errors"
	"os"
	"path"
	"testing"

	"github.com/kylelemons/godebug/pretty"
	"github.com/spf13/viper"
)

type testCase struct {
	name          string
	targetName    string
	targetAddress string
	targetCA      string
	targetCAKey   string
	config        map[string]Target
	want          result
}

type result struct {
	targets map[string]Target
	err     error
}

func TestSetTarget(t *testing.T) {
	viper.SetConfigFile("/tmp/config.yml")
	tests := generateTargetTestCases()
	for _, test := range tests {
		viper.Reset()
		t.Run(test.name, func(t *testing.T) {
			var targets map[string]Target
			viper.Set("targets.devices", test.config)
			err := SetTarget(test.targetName, test.targetAddress, test.targetCA, test.targetCAKey, true)
			viper.UnmarshalKey("targets.devices", &targets)
			got := result{
				targets: targets,
				err:     err,
			}
			if diff := pretty.Compare(test.want, got); diff != "" {
				t.Errorf("prepareTarget(%s, %s, %s, %s): (-want +got)\n%s", test.targetName, test.targetAddress, test.targetCA, test.targetCAKey, diff)
			}
		})
	}
}

func TestPrepareTarget(t *testing.T) {
	tests := generateTargetTestCases()
	for _, test := range tests {
		viper.Reset()
		t.Run(test.name, func(t *testing.T) {
			var targets map[string]Target
			viper.Set("targets.devices", test.config)
			err := prepareTarget(test.targetName, test.targetAddress, test.targetCA, test.targetCAKey, true)
			viper.UnmarshalKey("targets.devices", &targets)
			got := result{
				targets: targets,
				err:     err,
			}
			if diff := pretty.Compare(test.want, got); diff != "" {
				t.Errorf("prepareTarget(%s, %s, %s, %s): (-want +got)\n%s", test.targetName, test.targetAddress, test.targetCA, test.targetCAKey, diff)
			}
		})
	}
}

func generateTargetTestCases() []testCase {
	dir, _ := os.Getwd()
	certPath := path.Join(dir, "ca.crt")
	certKeyPath := path.Join(dir, "ca.key")
	history := map[string]Target{"myhost.com": {
		Address: "localhost:9339",
		Ca:      certPath,
		CaKey:   certKeyPath,
	}}
	tests := []testCase{
		{
			name: "No targets in history, no target specified",
			want: result{
				err:     errors.New("No targets in history and no target specified"),
				targets: map[string]Target{},
			},
		},
		{
			name:   "No target specified",
			config: history,
			want: result{
				targets: history,
			},
		},
		{
			name:       "Non-existent target",
			config:     map[string]Target{},
			targetName: "nonexistenttarget",
			want: result{
				err:     errors.New("Target not found"),
				targets: history,
			},
		},
		{
			name:          "Add new target",
			targetName:    "myhost.com",
			targetAddress: "localhost:9339",
			targetCA:      certPath,
			targetCAKey:   certKeyPath,
			want: result{
				targets: history,
			},
		},
		{
			name:          "Update existing target's address",
			targetName:    "myhost.com",
			targetAddress: "newhost:9340",
			config:        history,
			want: result{
				targets: map[string]Target{
					"myhost.com": {
						Address: "newhost:9340",
						Ca:      certPath,
						CaKey:   certKeyPath,
					},
				},
			},
		},
		{
			name:       "Update existing target's ca",
			targetName: "myhost.com",
			targetCA:   "newca.crt",
			config:     history,
			want: result{
				targets: map[string]Target{
					"myhost.com": {
						Address: "localhost:9339",
						Ca:      path.Join(dir, "newca.crt"),
						CaKey:   certKeyPath,
					},
				},
			},
		},
		{
			name:        "Update existing target's ca_key",
			targetName:  "myhost.com",
			targetCAKey: "newca.key",
			config:      history,
			want: result{
				targets: map[string]Target{
					"myhost.com": {
						Address: "localhost:9339",
						Ca:      certPath,
						CaKey:   path.Join(dir, "newca.key"),
					},
				},
			},
		},
	}
	return tests
}
<commit_msg>Fixed viper reset config issue<commit_after>/* Copyright 2020 Google Inc.

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

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

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

package config

import (
	"errors"
	"os"
	"path"
	"testing"

	"github.com/kylelemons/godebug/pretty"
	"github.com/spf13/viper"
)

type testCase struct {
	name          string
	targetName    string
	targetAddress string
	targetCA      string
	targetCAKey   string
	config        map[string]Target
	want          result
}

type result struct {
	targets map[string]Target
	err     error
}

func TestSetTarget(t *testing.T) {
	tests := generateTargetTestCases()
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			viper.Reset()
			viper.SetConfigFile("/tmp/config.yml")
			var targets map[string]Target
			viper.Set("targets.devices", test.config)
			err := SetTarget(test.targetName, test.targetAddress, test.targetCA, test.targetCAKey, true)
			viper.UnmarshalKey("targets.devices", &targets)
			got := result{
				targets: targets,
				err:     err,
			}
			if diff := pretty.Compare(test.want, got); diff != "" {
				t.Errorf("prepareTarget(%s, %s, %s, %s): (-want +got)\n%s", test.targetName, test.targetAddress, test.targetCA, test.targetCAKey, diff)
			}
		})
	}
}

func TestPrepareTarget(t *testing.T) {
	tests := generateTargetTestCases()
	for _, test := range tests {
		viper.Reset()
		t.Run(test.name, func(t *testing.T) {
			var targets map[string]Target
			viper.Set("targets.devices", test.config)
			err := prepareTarget(test.targetName, test.targetAddress, test.targetCA, test.targetCAKey, true)
			viper.UnmarshalKey("targets.devices", &targets)
			got := result{
				targets: targets,
				err:     err,
			}
			if diff := pretty.Compare(test.want, got); diff != "" {
				t.Errorf("prepareTarget(%s, %s, %s, %s): (-want +got)\n%s", test.targetName, test.targetAddress, test.targetCA, test.targetCAKey, diff)
			}
		})
	}
}

func generateTargetTestCases() []testCase {
	dir, _ := os.Getwd()
	certPath := path.Join(dir, "ca.crt")
	certKeyPath := path.Join(dir, "ca.key")
	history := map[string]Target{"myhost.com": {
		Address: "localhost:9339",
		Ca:      certPath,
		CaKey:   certKeyPath,
	}}
	tests := []testCase{
		{
			name: "No targets in history, no target specified",
			want: result{
				err:     errors.New("No targets in history and no target specified"),
				targets: map[string]Target{},
			},
		},
		{
			name:   "No target specified",
			config: history,
			want: result{
				targets: history,
			},
		},
		{
			name:       "Non-existent target",
			config:     map[string]Target{},
			targetName: "nonexistenttarget",
			want: result{
				err:     errors.New("Target not found"),
				targets: nil,
			},
		},
		{
			name:          "Add new target",
			targetName:    "myhost.com",
			targetAddress: "localhost:9339",
			targetCA:      certPath,
			targetCAKey:   certKeyPath,
			want: result{
				targets: history,
			},
		},
		{
			name:          "Update existing target's address",
			targetName:    "myhost.com",
			targetAddress: "newhost:9340",
			config:        history,
			want: result{
				targets: map[string]Target{
					"myhost.com": {
						Address: "newhost:9340",
						Ca:      certPath,
						CaKey:   certKeyPath,
					},
				},
			},
		},
		{
			name:       "Update existing target's ca",
			targetName: "myhost.com",
			targetCA:   "newca.crt",
			config:     history,
			want: result{
				targets: map[string]Target{
					"myhost.com": {
						Address: "localhost:9339",
						Ca:      path.Join(dir, "newca.crt"),
						CaKey:   certKeyPath,
					},
				},
			},
		},
		{
			name:        "Update existing target's ca_key",
			targetName:  "myhost.com",
			targetCAKey: "newca.key",
			config:      history,
			want: result{
				targets: map[string]Target{
					"myhost.com": {
						Address: "localhost:9339",
						Ca:      certPath,
						CaKey:   path.Join(dir, "newca.key"),
					},
				},
			},
		},
	}
	return tests
}
<|endoftext|>


<commit_before>package ovn

import (
	"fmt"

	util "github.com/openvswitch/ovn-kubernetes/go-controller/pkg/util"
	"github.com/sirupsen/logrus"
	kapi "k8s.io/api/core/v1"
)

type lbEndpoints struct {
	IPs  []string
	Port int32
}

// AddEndpoints adds endpoints and creates corresponding resources in OVN
func (ovn *Controller) AddEndpoints(ep *kapi.Endpoints) error {
	// get service
	// TODO: cache the service
	svc, err := ovn.kube.GetService(ep.Namespace, ep.Name)
	if err != nil {
		// This is not necessarily an error. For e.g when there are endpoints
		// without a corresponding service.
		logrus.Debugf("no service found for endpoint %s in namespace %s",
			ep.Name, ep.Namespace)
		return nil
	}
	if !util.IsServiceIPSet(svc) {
		logrus.Debugf("Skipping service %s due to clusterIP = %q",
			svc.Name, svc.Spec.ClusterIP)
		return nil
	}
	tcpPortMap := make(map[string]lbEndpoints)
	udpPortMap := make(map[string]lbEndpoints)
	for _, s := range ep.Subsets {
		for _, ip := range s.Addresses {
			for _, port := range s.Ports {
				var ips []string
				var portMap map[string]lbEndpoints
				if port.Protocol == kapi.ProtocolUDP {
					portMap = udpPortMap
				} else if port.Protocol == kapi.ProtocolTCP {
					portMap = tcpPortMap
				}
				if lbEps, ok := portMap[port.Name]; ok {
					ips = lbEps.IPs
				} else {
					ips = make([]string, 0)
				}
				ips = append(ips, ip.IP)
				portMap[port.Name] = lbEndpoints{IPs: ips, Port: port.Port}
			}
		}
	}

	logrus.Debugf("Tcp table: %v\nUdp table: %v", tcpPortMap, udpPortMap)

	for svcPortName, lbEps := range tcpPortMap {
		ips := lbEps.IPs
		targetPort := lbEps.Port
		for _, svcPort := range svc.Spec.Ports {
			if svcPort.Protocol == kapi.ProtocolTCP && svcPort.Name == svcPortName {
				if svc.Spec.Type == kapi.ServiceTypeNodePort && ovn.nodePortEnable {
					logrus.Debugf("Creating Gateways IP for NodePort: %d, %v", svcPort.NodePort, ips)
					err = ovn.createGatewaysVIP(string(svcPort.Protocol), svcPort.NodePort, targetPort, ips)
					if err != nil {
						logrus.Errorf("Error in creating Node Port for svc %s, node port: %d - %v\n", svc.Name, svcPort.NodePort, err)
						continue
					}
				}
				if svc.Spec.Type == kapi.ServiceTypeClusterIP || svc.Spec.Type == kapi.ServiceTypeNodePort {
					var loadBalancer string
					loadBalancer, err = ovn.getLoadBalancer(svcPort.Protocol)
					if err != nil {
						logrus.Errorf("Failed to get loadbalancer for %s (%v)",
							svcPort.Protocol, err)
						continue
					}
					err = ovn.createLoadBalancerVIP(loadBalancer,
						svc.Spec.ClusterIP, svcPort.Port, ips, targetPort)
					if err != nil {
						logrus.Errorf("Error in creating Cluster IP for svc %s, target port: %d - %v\n", svc.Name, targetPort, err)
						continue
					}
				}
				ovn.handleExternalIPs(svc, svcPort, ips, targetPort)
			}
		}
	}
	for svcPortName, lbEps := range udpPortMap {
		ips := lbEps.IPs
		targetPort := lbEps.Port
		for _, svcPort := range svc.Spec.Ports {
			if svcPort.Protocol == kapi.ProtocolUDP && svcPort.Name == svcPortName {
				if svc.Spec.Type == kapi.ServiceTypeNodePort && ovn.nodePortEnable {
					err = ovn.createGatewaysVIP(string(svcPort.Protocol), svcPort.NodePort, targetPort, ips)
					if err != nil {
						logrus.Errorf("Error in creating Node Port for svc %s, node port: %d - %v\n", svc.Name, svcPort.NodePort, err)
						continue
					}
				} else if svc.Spec.Type == kapi.ServiceTypeNodePort || svc.Spec.Type == kapi.ServiceTypeClusterIP {
					var loadBalancer string
					loadBalancer, err = ovn.getLoadBalancer(svcPort.Protocol)
					if err != nil {
						logrus.Errorf("Failed to get loadbalancer for %s (%v)",
							svcPort.Protocol, err)
						continue
					}
					err = ovn.createLoadBalancerVIP(loadBalancer,
						svc.Spec.ClusterIP, svcPort.Port, ips, targetPort)
					if err != nil {
						logrus.Errorf("Error in creating Cluster IP for svc %s, target port: %d - %v\n", svc.Name, targetPort, err)
						continue
					}
				}
				ovn.handleExternalIPs(svc, svcPort, ips, targetPort)
			}
		}
	}
	return nil
}

func (ovn *Controller) handleExternalIPs(svc *kapi.Service, svcPort kapi.ServicePort, ips []string, targetPort int32) {
	logrus.Debugf("handling external IPs for svc %v", svc.Name)
	if len(svc.Spec.ExternalIPs) == 0 {
		return
	}
	for _, extIP := range svc.Spec.ExternalIPs {
		lb := ovn.getDefaultGatewayLoadBalancer(svcPort.Protocol)
		if lb == "" {
			logrus.Warningf("No default gateway found for protocol %s\n\tNote: 'nodeport' flag needs to be enabled for default gateway", svcPort.Protocol)
			continue
		}
		err := ovn.createLoadBalancerVIP(lb, extIP, svcPort.Port, ips, targetPort)
		if err != nil {
			logrus.Errorf("Error in creating external IP for service: %s, externalIP: %s", svc.Name, extIP)
		}
	}
}

func (ovn *Controller) deleteEndpoints(ep *kapi.Endpoints) error {
	svc, err := ovn.kube.GetService(ep.Namespace, ep.Name)
	if err != nil {
		// This is not necessarily an error. For e.g when a service is deleted,
		// you will get endpoint delete event and the call to fetch service
		// will fail.
		logrus.Debugf("no service found for endpoint %s in namespace %s",
			ep.Name, ep.Namespace)
		return nil
	}
	if !util.IsServiceIPSet(svc) {
		return nil
	}
	for _, svcPort := range svc.Spec.Ports {
		var lb string
		lb, err = ovn.getLoadBalancer(svcPort.Protocol)
		if err != nil {
			logrus.Errorf("Failed to get load-balancer for %s (%v)",
				lb, err)
			continue
		}
		key := fmt.Sprintf("\"%s:%d\"", svc.Spec.ClusterIP, svcPort.Port)
		_, stderr, err := util.RunOVNNbctl("remove", "load_balancer", lb,
			"vips", key)
		if err != nil {
			logrus.Errorf("Error in deleting endpoints for lb %s, "+
				"stderr: %q (%v)", lb, stderr, err)
		}
	}
	return nil
}
<commit_msg>loadbalancer: create OVN LB rules for ClusterIP for UDP protocol<commit_after>package ovn

import (
	"fmt"

	util "github.com/openvswitch/ovn-kubernetes/go-controller/pkg/util"
	"github.com/sirupsen/logrus"
	kapi "k8s.io/api/core/v1"
)

type lbEndpoints struct {
	IPs  []string
	Port int32
}

// AddEndpoints adds endpoints and creates corresponding resources in OVN
func (ovn *Controller) AddEndpoints(ep *kapi.Endpoints) error {
	// get service
	// TODO: cache the service
	svc, err := ovn.kube.GetService(ep.Namespace, ep.Name)
	if err != nil {
		// This is not necessarily an error. For e.g when there are endpoints
		// without a corresponding service.
		logrus.Debugf("no service found for endpoint %s in namespace %s",
			ep.Name, ep.Namespace)
		return nil
	}
	if !util.IsServiceIPSet(svc) {
		logrus.Debugf("Skipping service %s due to clusterIP = %q",
			svc.Name, svc.Spec.ClusterIP)
		return nil
	}
	tcpPortMap := make(map[string]lbEndpoints)
	udpPortMap := make(map[string]lbEndpoints)
	for _, s := range ep.Subsets {
		for _, ip := range s.Addresses {
			for _, port := range s.Ports {
				var ips []string
				var portMap map[string]lbEndpoints
				if port.Protocol == kapi.ProtocolUDP {
					portMap = udpPortMap
				} else if port.Protocol == kapi.ProtocolTCP {
					portMap = tcpPortMap
				}
				if lbEps, ok := portMap[port.Name]; ok {
					ips = lbEps.IPs
				} else {
					ips = make([]string, 0)
				}
				ips = append(ips, ip.IP)
				portMap[port.Name] = lbEndpoints{IPs: ips, Port: port.Port}
			}
		}
	}

	logrus.Debugf("Tcp table: %v\nUdp table: %v", tcpPortMap, udpPortMap)

	for svcPortName, lbEps := range tcpPortMap {
		ips := lbEps.IPs
		targetPort := lbEps.Port
		for _, svcPort := range svc.Spec.Ports {
			if svcPort.Protocol == kapi.ProtocolTCP && svcPort.Name == svcPortName {
				if svc.Spec.Type == kapi.ServiceTypeNodePort && ovn.nodePortEnable {
					logrus.Debugf("Creating Gateways IP for NodePort: %d, %v", svcPort.NodePort, ips)
					err = ovn.createGatewaysVIP(string(svcPort.Protocol), svcPort.NodePort, targetPort, ips)
					if err != nil {
						logrus.Errorf("Error in creating Node Port for svc %s, node port: %d - %v\n", svc.Name, svcPort.NodePort, err)
						continue
					}
				}
				if svc.Spec.Type == kapi.ServiceTypeClusterIP || svc.Spec.Type == kapi.ServiceTypeNodePort {
					var loadBalancer string
					loadBalancer, err = ovn.getLoadBalancer(svcPort.Protocol)
					if err != nil {
						logrus.Errorf("Failed to get loadbalancer for %s (%v)",
							svcPort.Protocol, err)
						continue
					}
					err = ovn.createLoadBalancerVIP(loadBalancer,
						svc.Spec.ClusterIP, svcPort.Port, ips, targetPort)
					if err != nil {
						logrus.Errorf("Error in creating Cluster IP for svc %s, target port: %d - %v\n", svc.Name, targetPort, err)
						continue
					}
				}
				ovn.handleExternalIPs(svc, svcPort, ips, targetPort)
			}
		}
	}
	for svcPortName, lbEps := range udpPortMap {
		ips := lbEps.IPs
		targetPort := lbEps.Port
		for _, svcPort := range svc.Spec.Ports {
			if svcPort.Protocol == kapi.ProtocolUDP && svcPort.Name == svcPortName {
				if svc.Spec.Type == kapi.ServiceTypeNodePort && ovn.nodePortEnable {
					err = ovn.createGatewaysVIP(string(svcPort.Protocol), svcPort.NodePort, targetPort, ips)
					if err != nil {
						logrus.Errorf("Error in creating Node Port for svc %s, node port: %d - %v\n", svc.Name, svcPort.NodePort, err)
						continue
					}
				}
				if svc.Spec.Type == kapi.ServiceTypeNodePort || svc.Spec.Type == kapi.ServiceTypeClusterIP {
					var loadBalancer string
					loadBalancer, err = ovn.getLoadBalancer(svcPort.Protocol)
					if err != nil {
						logrus.Errorf("Failed to get loadbalancer for %s (%v)",
							svcPort.Protocol, err)
						continue
					}
					err = ovn.createLoadBalancerVIP(loadBalancer,
						svc.Spec.ClusterIP, svcPort.Port, ips, targetPort)
					if err != nil {
						logrus.Errorf("Error in creating Cluster IP for svc %s, target port: %d - %v\n", svc.Name, targetPort, err)
						continue
					}
				}
				ovn.handleExternalIPs(svc, svcPort, ips, targetPort)
			}
		}
	}
	return nil
}

func (ovn *Controller) handleExternalIPs(svc *kapi.Service, svcPort kapi.ServicePort, ips []string, targetPort int32) {
	logrus.Debugf("handling external IPs for svc %v", svc.Name)
	if len(svc.Spec.ExternalIPs) == 0 {
		return
	}
	for _, extIP := range svc.Spec.ExternalIPs {
		lb := ovn.getDefaultGatewayLoadBalancer(svcPort.Protocol)
		if lb == "" {
			logrus.Warningf("No default gateway found for protocol %s\n\tNote: 'nodeport' flag needs to be enabled for default gateway", svcPort.Protocol)
			continue
		}
		err := ovn.createLoadBalancerVIP(lb, extIP, svcPort.Port, ips, targetPort)
		if err != nil {
			logrus.Errorf("Error in creating external IP for service: %s, externalIP: %s", svc.Name, extIP)
		}
	}
}

func (ovn *Controller) deleteEndpoints(ep *kapi.Endpoints) error {
	svc, err := ovn.kube.GetService(ep.Namespace, ep.Name)
	if err != nil {
		// This is not necessarily an error. For e.g when a service is deleted,
		// you will get endpoint delete event and the call to fetch service
		// will fail.
		logrus.Debugf("no service found for endpoint %s in namespace %s",
			ep.Name, ep.Namespace)
		return nil
	}
	if !util.IsServiceIPSet(svc) {
		return nil
	}
	for _, svcPort := range svc.Spec.Ports {
		var lb string
		lb, err = ovn.getLoadBalancer(svcPort.Protocol)
		if err != nil {
			logrus.Errorf("Failed to get load-balancer for %s (%v)",
				lb, err)
			continue
		}
		key := fmt.Sprintf("\"%s:%d\"", svc.Spec.ClusterIP, svcPort.Port)
		_, stderr, err := util.RunOVNNbctl("remove", "load_balancer", lb,
			"vips", key)
		if err != nil {
			logrus.Errorf("Error in deleting endpoints for lb %s, "+
				"stderr: %q (%v)", lb, stderr, err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Sia uses Reed-Solomon coding for error correction. This package has no
// method for error detection however, so error detection must be performed
// elsewhere.
//
// We use the repository 'Siacoin/longhair' to handle the erasure coding.
// As far as I'm aware, it's the fastest library that's open source.
// It is a fork of 'catid/longhair', and we inted to merge all changes from
// the original.
//
// Longhair is a c++ library. Here, it is cast to a C library and then called
// using cgo.
package erasure

// #include "erasure.c"
import "C"

import (
	"bytes"
	"common"
	"common/crypto"
	"fmt"
	"unsafe"
)

// EncodeRing takes a Sector and encodes it as a Ring: a set of common.QuorumSize Segments that include redundancy.
// k is the number of non-redundant segments, and b is the size of each segment. b is calculated from k.
// The erasure-coding algorithm requires that the original data must be k*b in size, so it is padded here as needed.
//
// The return value is a Ring.
// The first k Segments of the Ring are the original data split up.
// The remaining Segments are newly generated redundant data.
func EncodeRing(sec *common.Sector, k int) (ring *common.Ring, segs [common.QuorumSize]common.Segment, err error) {
	// check for legal size of k
	if k <= 0 || k >= common.QuorumSize {
		err = fmt.Errorf("k must be greater than 0 and smaller than %v", common.QuorumSize)
		return
	}

	// calculate b
	b := len(sec.Data) / k
	if b%64 != 0 {
		b += 64 - (b % 64)
	}

	// check for legal size of b
	if b < common.MinSegmentSize || b > common.MaxSegmentSize {
		err = fmt.Errorf("b must be greater than %v and smaller than %v", common.MinSegmentSize, common.MaxSegmentSize)
		return
	}

	// store encoding parameters in ring
	ring = common.NewRing(k, b, len(sec.Data))

	// pad data as needed
	padding := k*b - len(sec.Data)
	paddedData := append(sec.Data, bytes.Repeat([]byte{0x00}, padding)...)

	// call the encoding function
	m := common.QuorumSize - k
	redundantChunk := C.encodeRedundancy(C.int(k), C.int(m), C.int(b), (*C.char)(unsafe.Pointer(&paddedData[0])))
	redundantBytes := C.GoBytes(unsafe.Pointer(redundantChunk), C.int(m*b))

	// split paddedData into ring
	for i := 0; i < k; i++ {
		segs[i] = common.Segment{
			paddedData[i*b : (i+1)*b],
			uint8(i),
		}
		ring.SegHashes[i], err = crypto.CalculateHash(paddedData[i*b : (i+1)*b])
		if err != nil {
			return
		}
	}

	// split redundantString into ring
	for i := k; i < common.QuorumSize; i++ {
		segs[i] = common.Segment{
			redundantBytes[(i-k)*b : (i-k+1)*b],
			uint8(i),
		}
		ring.SegHashes[i], err = crypto.CalculateHash(redundantBytes[(i-k)*b : (i-k+1)*b])
		if err != nil {
			return
		}
	}

	// free the memory allocated by the C file
	C.free(unsafe.Pointer(redundantChunk))

	return
}

// RebuildSector takes a Ring and returns a Sector containing the original data.
// The Ring's k value must be equal to the number of non-redundant segments when the file was originally built.
// Because recovery is just a bunch of matrix operations, there is no way to tell if the data has been corrupted
// or if an incorrect value of k has been chosen. This error checking must happen before calling RebuildSector.
// Each Segment's Data must have the correct Index from when it was encoded.
func RebuildSector(ring *common.Ring, segs []common.Segment) (sec *common.Sector, err error) {
	k, b := ring.GetRedundancy(), ring.GetBytesPerSegment()
	if k == 0 && b == 0 {
		err = fmt.Errorf("could not rebuild using uninitialized encoding parameters")
		return
	}

	// check for legal size of k
	m := common.QuorumSize - k
	if k > common.QuorumSize || k < 1 {
		err = fmt.Errorf("k must be greater than 0 but smaller than %v", common.QuorumSize)
		return
	}

	// check for legal size of b
	if b < common.MinSegmentSize || b > common.MaxSegmentSize {
		err = fmt.Errorf("b must be greater than %v and smaller than %v", common.MinSegmentSize, common.MaxSegmentSize)
		return
	}

	// check for correct number of segments
	if len(segs) < k {
		err = fmt.Errorf("insufficient segments: expected at least %v, got %v", k, len(segs))
		return
	}

	// move all data into a single slice
	var segmentData []byte
	var segmentIndices []uint8
	for i := 0; i < k; i++ {
		// verify that each segment is the correct length
		if len(segs[i].Data) != b {
			err = fmt.Errorf("at least 1 Segment's Data field is the wrong length")
			return
		}

		segmentData = append(segmentData, segs[i].Data...)
		segmentIndices = append(segmentIndices, segs[i].Index)

	}
	// call the recovery function
	C.recoverData(C.int(k), C.int(m), C.int(b), (*C.uchar)(unsafe.Pointer(&segmentData[0])), (*C.uchar)(unsafe.Pointer(&segmentIndices[0])))

	// remove padding introduced by EncodeRing()
	sec, err = common.NewSector(segmentData[:ring.GetLength()])
	return
}
<commit_msg>update erasure for new Ring<commit_after>// Sia uses Reed-Solomon coding for error correction. This package has no
// method for error detection however, so error detection must be performed
// elsewhere.
//
// We use the repository 'Siacoin/longhair' to handle the erasure coding.
// As far as I'm aware, it's the fastest library that's open source.
// It is a fork of 'catid/longhair', and we inted to merge all changes from
// the original.
//
// Longhair is a c++ library. Here, it is cast to a C library and then called
// using cgo.
package erasure

// #include "erasure.c"
import "C"

import (
	"bytes"
	"common"
	"fmt"
	"unsafe"
)

// EncodeRing takes a Sector and encodes it as a Ring: a set of common.QuorumSize Segments that include redundancy.
// k is the number of non-redundant segments, and b is the size of each segment. b is calculated from k.
// The erasure-coding algorithm requires that the original data must be k*b in size, so it is padded here as needed.
//
// The return value is a Ring.
// The first k Segments of the Ring are the original data split up.
// The remaining Segments are newly generated redundant data.
func EncodeRing(sec *common.Sector, params *common.EncodingParams) (ring [common.QuorumSize]common.Segment, err error) {
	k, b, length := params.GetValues()

	// check for legal size of k
	if k <= 0 || k >= common.QuorumSize {
		err = fmt.Errorf("k must be greater than 0 and smaller than %v", common.QuorumSize)
		return
	}

	// check for legal size of b
	if b < common.MinSegmentSize || b > common.MaxSegmentSize {
		err = fmt.Errorf("b must be greater than %v and smaller than %v", common.MinSegmentSize, common.MaxSegmentSize)
		return
	}

	// check for legal size of length
	if length != len(sec.Data) {
		err = fmt.Errorf("length mismatch: sector length %v != parameter length %v", len(sec.Data), length)
		return
	} else if length > common.MaxSegmentSize*common.QuorumSize {
		err = fmt.Errorf("length must be smaller than %v", common.MaxSegmentSize*common.QuorumSize)
	}

	// pad data as needed
	padding := k*b - len(sec.Data)
	paddedData := append(sec.Data, bytes.Repeat([]byte{0x00}, padding)...)

	// call the encoding function
	m := common.QuorumSize - k
	redundantChunk := C.encodeRedundancy(C.int(k), C.int(m), C.int(b), (*C.char)(unsafe.Pointer(&paddedData[0])))
	redundantBytes := C.GoBytes(unsafe.Pointer(redundantChunk), C.int(m*b))

	// split paddedData into ring
	for i := 0; i < k; i++ {
		ring[i] = common.Segment{
			paddedData[i*b : (i+1)*b],
			uint8(i),
		}
	}

	// split redundantString into ring
	for i := k; i < common.QuorumSize; i++ {
		ring[i] = common.Segment{
			redundantBytes[(i-k)*b : (i-k+1)*b],
			uint8(i),
		}
	}

	// free the memory allocated by the C file
	C.free(unsafe.Pointer(redundantChunk))

	return
}

// RebuildSector takes a Ring and returns a Sector containing the original data.
// The Ring's k value must be equal to the number of non-redundant segments when the file was originally built.
// Because recovery is just a bunch of matrix operations, there is no way to tell if the data has been corrupted
// or if an incorrect value of k has been chosen. This error checking must happen before calling RebuildSector.
// Each Segment's Data must have the correct Index from when it was encoded.
func RebuildSector(ring []common.Segment, params *common.EncodingParams) (sec *common.Sector, err error) {
	k, b, length := params.GetValues()
	if k == 0 && b == 0 {
		err = fmt.Errorf("could not rebuild using uninitialized encoding parameters")
		return
	}

	// check for legal size of k
	if k > common.QuorumSize || k < 1 {
		err = fmt.Errorf("k must be greater than 0 but smaller than %v", common.QuorumSize)
		return
	}

	// check for legal size of b
	if b < common.MinSegmentSize || b > common.MaxSegmentSize {
		err = fmt.Errorf("b must be greater than %v and smaller than %v", common.MinSegmentSize, common.MaxSegmentSize)
		return
	}

	// check for legal size of length
	if length > common.MaxSegmentSize*common.QuorumSize {
		err = fmt.Errorf("length must be smaller than %v", common.MaxSegmentSize*common.QuorumSize)
	}

	// check for correct number of segments
	if len(ring) < k {
		err = fmt.Errorf("insufficient segments: expected at least %v, got %v", k, len(ring))
		return
	}

	// move all data into a single slice
	var segmentData []byte
	var segmentIndices []uint8
	for i := 0; i < k; i++ {
		// verify that each segment is the correct length
		// TODO: skip bad segments and continue rebuilding if possible
		if len(ring[i].Data) != b {
			err = fmt.Errorf("at least 1 Segment's Data field is the wrong length")
			return
		}

		segmentData = append(segmentData, ring[i].Data...)
		segmentIndices = append(segmentIndices, ring[i].Index)

	}
	// call the recovery function
	C.recoverData(C.int(k), C.int(common.QuorumSize-k), C.int(b), (*C.uchar)(unsafe.Pointer(&segmentData[0])), (*C.uchar)(unsafe.Pointer(&segmentIndices[0])))

	// remove padding introduced by EncodeRing()
	sec, err = common.NewSector(segmentData[:length])
	return
}
<|endoftext|>


<commit_before>package models

import (
	"fmt"
	"socialapi/request"
	"strings"

	"github.com/koding/bongo"
)

type Account struct {
	// unique id of the account
	Id int64 `json:"id,string"`

	// old id of the account, which is originally
	// perisisted in mongo
	// mongo ids has 24 char
	OldId string `json:"oldId"      sql:"NOT NULL;UNIQUE;TYPE:VARCHAR(24);"`

	// IsTroll
	IsTroll bool `json:"isTroll"`

	// unique account nicknames
	Nick string `json:"nick"        sql:"NOT NULL;UNIQUE;TYPE:VARCHAR(25);"`
}

func (a *Account) FetchOrCreate() error {
	if a.OldId == "" {
		return ErrOldIdIsNotSet
	}

	if strings.Contains(a.Nick, "guest-") {
		return ErrGuestsAreNotAllowed
	}

	selector := map[string]interface{}{
		"old_id": a.OldId,
	}

	err := a.One(bongo.NewQS(selector))
	// if we dont get any error
	// it means we found the record in our db
	if err == nil {
		return nil
	}

	// first check if the err is not found err
	if err == bongo.RecordNotFound {
		if a.Nick == "" {
			return ErrNickIsNotSet
		}

		if err := a.Create(); err != nil {
			return err
		}

		return nil
	}

	return err
}

func (a *Account) FetchChannels(q *request.Query) ([]Channel, error) {
	cp := NewChannelParticipant()
	// fetch channel ids
	cids, err := cp.FetchParticipatedChannelIds(a, q)
	if err != nil {
		return nil, err
	}

	// fetch channels by their ids
	channels, err := NewChannel().FetchByIds(cids)
	if err != nil {
		return nil, err
	}

	return channels, nil
}

func (a *Account) Follow(targetId int64) (*ChannelParticipant, error) {
	c, err := a.FetchChannel(Channel_TYPE_FOLLOWERS)
	if err == nil {
		return c.AddParticipant(targetId)
	}

	if err == bongo.RecordNotFound {
		c, err := a.CreateFollowingFeedChannel()
		if err != nil {
			return nil, err
		}
		return c.AddParticipant(targetId)
	}

	return nil, err
}

func (a *Account) Unfollow(targetId int64) (*Account, error) {
	c, err := a.FetchChannel(Channel_TYPE_FOLLOWERS)
	if err != nil {
		return nil, err
	}

	return a, c.RemoveParticipant(targetId)
}

func (a *Account) FetchFollowerIds(q *request.Query) ([]int64, error) {
	followerIds := make([]int64, 0)
	if a.Id == 0 {
		return nil, ErrAccountIdIsNotSet
	}

	c, err := a.FetchChannel(Channel_TYPE_FOLLOWERS)
	if err != nil {
		return followerIds, err
	}

	participants, err := c.FetchParticipantIds(q)
	if err != nil {
		return followerIds, err
	}

	return participants, nil
}

// FetchChannel fetchs the channel of the account
//
// Channel_TYPE_GROUP as parameter returns error , in the tests!!!!
// TO-DO, other types dont return error
//
// Tests are done
func (a *Account) FetchChannel(channelType string) (*Channel, error) {
	if a.Id == 0 {
		return nil, ErrAccountIdIsNotSet
	}

	c := NewChannel()
	selector := map[string]interface{}{
		"creator_id":    a.Id,
		"type_constant": channelType,
	}

	if err := c.One(bongo.NewQS(selector)); err != nil {
		return nil, err
	}

	return c, nil
}

func (a *Account) MarkAsTroll() error {
	if a.Id == 0 {
		return ErrAccountIdIsNotSet
	}

	if err := a.ById(a.Id); err != nil {
		return err
	}

	// do not try to mark twice
	if a.IsTroll {
		return fmt.Errorf("account is already a troll %d", a.Id)
	}

	a.IsTroll = true
	if err := a.Update(); err != nil {
		return err
	}

	if err := bongo.B.PublishEvent("marked_as_troll", a); err != nil {
		return err
	}

	return nil
}

func (a *Account) UnMarkAsTroll() error {
	if a.Id == 0 {
		return ErrAccountIdIsNotSet
	}

	if err := a.ById(a.Id); err != nil {
		return err
	}

	// do not try to un-mark twice
	if !a.IsTroll {
		return fmt.Errorf("account is not a troll %d", a.Id)
	}

	a.IsTroll = false
	if err := a.Update(); err != nil {
		return err
	}

	if err := bongo.B.PublishEvent("unmarked_as_troll", a); err != nil {
		return err
	}

	return nil
}

// Tests are done.
func (a *Account) CreateFollowingFeedChannel() (*Channel, error) {
	if a.Id == 0 {
		return nil, ErrAccountIdIsNotSet
	}

	c := NewChannel()
	c.CreatorId = a.Id
	c.Name = fmt.Sprintf("%d-FollowingFeedChannel", a.Id)
	c.GroupName = Channel_KODING_NAME
	c.Purpose = "Following Feed for Me"
	c.TypeConstant = Channel_TYPE_FOLLOWERS
	if err := c.Create(); err != nil {
		return nil, err
	}

	return c, nil
}

func (a *Account) FetchFollowerChannelIds(q *request.Query) ([]int64, error) {

	followerIds, err := a.FetchFollowerIds(q)
	if err != nil {
		return nil, err
	}

	cp := NewChannelParticipant()
	var channelIds []int64
	res := bongo.B.DB.
		Table(cp.TableName()).
		Where(
		"creator_id IN (?) and type_constant = ?",
		followerIds,
		Channel_TYPE_FOLLOWINGFEED,
	).Find(&channelIds)

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

	return channelIds, nil
}

// FetchAccountById gives all information about account by id of account
//
// Tests are done.
func FetchAccountById(accountId int64) (*Account, error) {
	a := NewAccount()
	if err := a.ById(accountId); err != nil {
		return nil, err
	}

	return a, nil
}

// FetchOldIdsByAccountIds takes slice as parameter
//
// Tests are done
func FetchOldIdsByAccountIds(accountIds []int64) ([]string, error) {
	oldIds := make([]string, 0)

	if len(accountIds) == 0 {
		return oldIds, nil
	}

	var accounts []Account
	account := Account{}

	err := bongo.B.FetchByIds(account, &accounts, accountIds)
	if err != nil {
		return oldIds, nil
	}

	for _, account := range accounts {
		// The append built-in function appends elements to the end of a slice
		oldIds = append(oldIds, account.OldId)
	}

	return oldIds, nil
}
<commit_msg>social/tests: Functions are controlled and required comments are added<commit_after>package models

import (
	"fmt"
	"socialapi/request"
	"strings"

	"github.com/koding/bongo"
)

type Account struct {
	// unique id of the account
	Id int64 `json:"id,string"`

	// old id of the account, which is originally
	// perisisted in mongo
	// mongo ids has 24 char
	OldId string `json:"oldId"      sql:"NOT NULL;UNIQUE;TYPE:VARCHAR(24);"`

	// IsTroll
	IsTroll bool `json:"isTroll"`

	// unique account nicknames
	Nick string `json:"nick"        sql:"NOT NULL;UNIQUE;TYPE:VARCHAR(25);"`
}

// Tests are done.
func (a *Account) FetchOrCreate() error {
	if a.OldId == "" {
		return ErrOldIdIsNotSet
	}

	if strings.Contains(a.Nick, "guest-") {
		return ErrGuestsAreNotAllowed
	}

	selector := map[string]interface{}{
		"old_id": a.OldId,
	}

	err := a.One(bongo.NewQS(selector))
	// if we dont get any error
	// it means we found the record in our db
	if err == nil {
		return nil
	}

	// first check if the err is not found err
	if err == bongo.RecordNotFound {
		if a.Nick == "" {
			return ErrNickIsNotSet
		}

		if err := a.Create(); err != nil {
			return err
		}

		return nil
	}

	return err
}

func (a *Account) FetchChannels(q *request.Query) ([]Channel, error) {
	cp := NewChannelParticipant()
	// fetch channel ids
	cids, err := cp.FetchParticipatedChannelIds(a, q)
	if err != nil {
		return nil, err
	}

	// fetch channels by their ids
	channels, err := NewChannel().FetchByIds(cids)
	if err != nil {
		return nil, err
	}

	return channels, nil
}

// TO-DO
// Control functions and remove ?
func (a *Account) Follow(targetId int64) (*ChannelParticipant, error) {
	c, err := a.FetchChannel(Channel_TYPE_FOLLOWERS)
	if err == nil {
		return c.AddParticipant(targetId)
	}

	if err == bongo.RecordNotFound {
		c, err := a.CreateFollowingFeedChannel()
		if err != nil {
			return nil, err
		}
		return c.AddParticipant(targetId)
	}

	return nil, err
}

// TO-DO
// Control functions and remove ?
func (a *Account) Unfollow(targetId int64) (*Account, error) {
	c, err := a.FetchChannel(Channel_TYPE_FOLLOWERS)
	if err != nil {
		return nil, err
	}

	return a, c.RemoveParticipant(targetId)
}

func (a *Account) FetchFollowerIds(q *request.Query) ([]int64, error) {
	followerIds := make([]int64, 0)
	if a.Id == 0 {
		return nil, ErrAccountIdIsNotSet
	}

	c, err := a.FetchChannel(Channel_TYPE_FOLLOWERS)
	if err != nil {
		return followerIds, err
	}

	participants, err := c.FetchParticipantIds(q)
	if err != nil {
		return followerIds, err
	}

	return participants, nil
}

// FetchChannel fetchs the channel of the account
//
// Channel_TYPE_GROUP as parameter returns error , in the tests!!!!
// TO-DO, other types dont return error
//
// Tests are done
func (a *Account) FetchChannel(channelType string) (*Channel, error) {
	if a.Id == 0 {
		return nil, ErrAccountIdIsNotSet
	}

	c := NewChannel()
	selector := map[string]interface{}{
		"creator_id":    a.Id,
		"type_constant": channelType,
	}

	if err := c.One(bongo.NewQS(selector)); err != nil {
		return nil, err
	}

	return c, nil
}

func (a *Account) MarkAsTroll() error {
	if a.Id == 0 {
		return ErrAccountIdIsNotSet
	}

	if err := a.ById(a.Id); err != nil {
		return err
	}

	// do not try to mark twice
	if a.IsTroll {
		return fmt.Errorf("account is already a troll %d", a.Id)
	}

	a.IsTroll = true
	if err := a.Update(); err != nil {
		return err
	}

	if err := bongo.B.PublishEvent("marked_as_troll", a); err != nil {
		return err
	}

	return nil
}

func (a *Account) UnMarkAsTroll() error {
	if a.Id == 0 {
		return ErrAccountIdIsNotSet
	}

	if err := a.ById(a.Id); err != nil {
		return err
	}

	// do not try to un-mark twice
	if !a.IsTroll {
		return fmt.Errorf("account is not a troll %d", a.Id)
	}

	a.IsTroll = false
	if err := a.Update(); err != nil {
		return err
	}

	if err := bongo.B.PublishEvent("unmarked_as_troll", a); err != nil {
		return err
	}

	return nil
}

// Tests are done.
func (a *Account) CreateFollowingFeedChannel() (*Channel, error) {
	if a.Id == 0 {
		return nil, ErrAccountIdIsNotSet
	}

	c := NewChannel()
	c.CreatorId = a.Id
	c.Name = fmt.Sprintf("%d-FollowingFeedChannel", a.Id)
	c.GroupName = Channel_KODING_NAME
	c.Purpose = "Following Feed for Me"
	c.TypeConstant = Channel_TYPE_FOLLOWERS
	if err := c.Create(); err != nil {
		return nil, err
	}

	return c, nil
}

func (a *Account) FetchFollowerChannelIds(q *request.Query) ([]int64, error) {

	followerIds, err := a.FetchFollowerIds(q)
	if err != nil {
		return nil, err
	}

	cp := NewChannelParticipant()
	var channelIds []int64
	res := bongo.B.DB.
		Table(cp.TableName()).
		Where(
		"creator_id IN (?) and type_constant = ?",
		followerIds,
		Channel_TYPE_FOLLOWINGFEED,
	).Find(&channelIds)

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

	return channelIds, nil
}

// FetchAccountById gives all information about account by id of account
//
// Tests are done.
func FetchAccountById(accountId int64) (*Account, error) {
	a := NewAccount()
	if err := a.ById(accountId); err != nil {
		return nil, err
	}

	return a, nil
}

// FetchOldIdsByAccountIds takes slice as parameter
//
// Tests are done
func FetchOldIdsByAccountIds(accountIds []int64) ([]string, error) {
	oldIds := make([]string, 0)

	if len(accountIds) == 0 {
		return oldIds, nil
	}

	var accounts []Account
	account := Account{}

	err := bongo.B.FetchByIds(account, &accounts, accountIds)
	if err != nil {
		return oldIds, nil
	}

	for _, account := range accounts {
		// The append built-in function appends elements to the end of a slice
		oldIds = append(oldIds, account.OldId)
	}

	return oldIds, nil
}
<|endoftext|>