<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package dump

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
)

// logDumper gets all the nodes from a kubernetes cluster and dumps a well-known set of logs
type logDumper struct {
	sshClientFactory sshClientFactory

	artifactsDir string

	services     []string
	files        []string
	podSelectors []string
}

// NewLogDumper is the constructor for a logDumper
func NewLogDumper(sshConfig *ssh.ClientConfig, artifactsDir string) *logDumper {
	sshClientFactory := &sshClientFactoryImplementation{
		sshConfig: sshConfig,
	}

	d := &logDumper{
		sshClientFactory: sshClientFactory,
		artifactsDir:     artifactsDir,
	}

	d.services = []string{
		"node-problem-detector",
		"kubelet",
		"containerd",
		"docker",
		"kops-configuration",
		"protokube",
	}
	d.files = []string{
		"kube-apiserver",
		"kube-scheduler",
		"rescheduler",
		"kube-controller-manager",
		"etcd",
		"etcd-events",
		"etcd-cilium",
		"glbc",
		"cluster-autoscaler",
		"kube-addon-manager",
		"fluentd",
		"kube-proxy",
		"node-problem-detector",
		"cloud-init-output",
		"startupscript",
		"kern",
		"docker",
		"aws-routed-eni/ipamd",
		"aws-routed-eni/plugin",
	}
	d.podSelectors = []string{
		"k8s-app=external-dns",
		"k8s-app=dns-controller",
	}

	return d
}

// DumpAllNodes connects to every node from kubectl get nodes and dumps the logs.
// additionalIPs holds IP addresses of instances found by the deployment tool;
// if the IPs are not found from kubectl get nodes, then these will be dumped also.
// This allows for dumping log on nodes even if they don't register as a kubernetes
// node, or if a node fails to register, or if the whole cluster fails to start.
func (d *logDumper) DumpAllNodes(ctx context.Context, nodes corev1.NodeList, additionalIPs []string) error {
	var dumped []*corev1.Node

	for i := range nodes.Items {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		node := &nodes.Items[i]

		ip := ""
		for _, address := range node.Status.Addresses {
			if address.Type == "ExternalIP" {
				ip = address.Address
				break
			}
		}

		err := d.dumpNode(ctx, node.Name, ip)
		if err != nil {
			log.Printf("could not dump node %s (%s): %v", node.Name, ip, err)
		} else {
			dumped = append(dumped, node)
		}
	}

	notDumped := findInstancesNotDumped(additionalIPs, dumped)
	for _, ip := range notDumped {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		log.Printf("dumping node not registered in kubernetes: %s", ip)
		err := d.dumpNode(ctx, ip, ip)
		if err != nil {
			log.Printf("error dumping node %s: %v", ip, err)
		}
	}

	return nil
}

// findInstancesNotDumped returns ips from the slice that do not appear as any address of the nodes
func findInstancesNotDumped(ips []string, dumped []*corev1.Node) []string {
	var notDumped []string
	dumpedAddresses := make(map[string]bool)
	for _, node := range dumped {
		for _, address := range node.Status.Addresses {
			dumpedAddresses[address.Address] = true
		}
	}

	for _, ip := range ips {
		if !dumpedAddresses[ip] {
			notDumped = append(notDumped, ip)
		}
	}
	return notDumped
}

// DumpNode connects to a node and dumps the logs.
func (d *logDumper) dumpNode(ctx context.Context, name string, ip string) error {
	if ip == "" {
		return fmt.Errorf("could not find address for %v, ", name)
	}

	log.Printf("Dumping node %s", name)

	n, err := d.connectToNode(ctx, name, ip)
	if err != nil {
		return fmt.Errorf("could not connect: %v", err)
	}

	// As long as we connect to the node we will not return an error;
	// a failure to collect a log (or even any logs at all) is not
	// considered an error in dumping the node.
	// TODO(justinsb): clean up / rationalize
	errors := n.dump(ctx)
	for _, e := range errors {
		log.Printf("error dumping node %s: %v", name, e)
	}

	if err := n.Close(); err != nil {
		log.Printf("error closing connection: %v", err)
	}

	return nil
}

// sshClient is an interface abstracting *ssh.Client, which allows us to test it
type sshClient interface {
	io.Closer

	// ExecPiped runs the command, piping stdout & stderr
	ExecPiped(ctx context.Context, command string, stdout io.Writer, stderr io.Writer) error
}

// sshClientFactory is an interface abstracting to a node over SSH
type sshClientFactory interface {
	Dial(ctx context.Context, host string) (sshClient, error)
}

// logDumperNode holds state for a particular node we are dumping
type logDumperNode struct {
	client sshClient
	dumper *logDumper

	dir string
}

// connectToNode makes an SSH connection to the node and returns a logDumperNode
func (d *logDumper) connectToNode(ctx context.Context, nodeName string, host string) (*logDumperNode, error) {
	client, err := d.sshClientFactory.Dial(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("unable to SSH to %q: %v", host, err)
	}
	return &logDumperNode{
		client: client,
		dir:    filepath.Join(d.artifactsDir, nodeName),
		dumper: d,
	}, nil
}

// logDumperNode cleans up any state in the logDumperNode
func (n *logDumperNode) Close() error {
	return n.client.Close()
}

// dump captures the well-known set of logs
func (n *logDumperNode) dump(ctx context.Context) []error {
	if ctx.Err() != nil {
		return []error{ctx.Err()}
	}

	var errors []error

	// Capture kernel log
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise -k", filepath.Join(n.dir, "kern.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture full journal - needed so we can see e.g. disk mounts
	// This does duplicate the other files, but ensures we have all output
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise", filepath.Join(n.dir, "journal.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture logs from any systemd services in our list that are registered
	services, err := n.listSystemdUnits(ctx)
	if err != nil {
		errors = append(errors, fmt.Errorf("error listing systemd services: %v", err))
	}
	for _, s := range n.dumper.services {
		name := s + ".service"
		for _, service := range services {
			if service == name {
				if err := n.shellToFile(ctx, "sudo journalctl --output=cat -u "+name, filepath.Join(n.dir, s+".log")); err != nil {
					errors = append(errors, err)
				}
			}
		}
	}

	// Capture any file logs where the files exist
	fileList, err := n.findFiles(ctx, "/var/log")
	if err != nil {
		errors = append(errors, fmt.Errorf("error reading /var/log: %v", err))
	}
	for _, name := range n.dumper.files {
		prefix := "/var/log/" + name + ".log"
		for _, f := range fileList {
			if !strings.HasPrefix(f, prefix) {
				continue
			}
			if err := n.shellToFile(ctx, "sudo cat '"+strings.ReplaceAll(f, "'", "'\\''")+"'", filepath.Join(n.dir, strings.ReplaceAll(strings.TrimPrefix(f, "/var/log/"), "/", "_"))); err != nil {
				errors = append(errors, err)
			}
		}
	}

	for _, selector := range n.dumper.podSelectors {
		kv := strings.Split(selector, "=")
		logFile := fmt.Sprintf("%v.log", kv[len(kv)-1])
		if err := n.shellToFile(ctx, "if command -v kubectl &> /dev/null; then kubectl logs -n kube-system --all-containers -l \""+selector+"\"; fi", filepath.Join(n.dir, logFile)); err != nil {
			errors = append(errors, err)
		}
	}

	if err := n.shellToFile(ctx, "cat /etc/hosts", filepath.Join(n.dir, "etchosts")); err != nil {
		errors = append(errors, err)
	}

	return errors
}

// findFiles lists files under the specified directory (recursively)
func (n *logDumperNode) findFiles(ctx context.Context, dir string) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo find "+dir+" -print0", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing %q: %v", dir, err)
	}

	paths := []string{}
	for _, b := range bytes.Split(stdout.Bytes(), []byte{0}) {
		if len(b) == 0 {
			// Likely the last value
			continue
		}
		paths = append(paths, string(b))
	}
	return paths, nil
}

// listSystemdUnits returns the list of systemd units on the node
func (n *logDumperNode) listSystemdUnits(ctx context.Context) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo systemctl list-units -t service --no-pager --no-legend --all", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing systemd units: %v", err)
	}

	var services []string
	for _, line := range strings.Split(stdout.String(), "\n") {
		tokens := strings.Fields(line)
		if len(tokens) == 0 || tokens[0] == "" {
			continue
		}
		services = append(services, tokens[0])
	}
	return services, nil
}

// shellToFile executes a command and copies the output to a file
func (n *logDumperNode) shellToFile(ctx context.Context, command string, destPath string) error {
	if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil {
		log.Printf("unable to mkdir on %q: %v", filepath.Dir(destPath), err)
	}

	f, err := os.Create(destPath)
	if err != nil {
		return fmt.Errorf("error creating file %q: %v", destPath, err)
	}
	defer f.Close()

	if err := n.client.ExecPiped(ctx, command, f, f); err != nil {
		return fmt.Errorf("error executing command %q: %v", command, err)
	}

	return nil
}

// sshClientImplementation is the default implementation of sshClient, binding to a *ssh.Client
type sshClientImplementation struct {
	client *ssh.Client
}

var _ sshClient = &sshClientImplementation{}

// ExecPiped implements sshClientImplementation::ExecPiped
func (s *sshClientImplementation) ExecPiped(ctx context.Context, cmd string, stdout io.Writer, stderr io.Writer) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	finished := make(chan error)
	go func() {
		session, err := s.client.NewSession()
		if err != nil {
			finished <- fmt.Errorf("error creating ssh session: %v", err)
			return
		}
		defer session.Close()

		klog.V(2).Infof("running SSH command: %v", cmd)

		session.Stdout = stdout
		session.Stderr = stderr

		finished <- session.Run(cmd)
	}()

	select {
	case <-ctx.Done():
		log.Print("closing SSH tcp connection due to context completion")

		// terminate the TCP connection to force a disconnect - we assume everyone is using the same context.
		// We could make this better by sending a signal on the session, waiting and then closing the session,
		// and only if we still haven't succeeded then closing the TCP connection.  This is sufficient for our
		// current usage though - and hopefully that logic will be implemented in the SSH package itself.
		s.Close()

		<-finished // Wait for cancellation
		return ctx.Err()

	case err := <-finished:
		return err
	}
}

// Close implements sshClientImplementation::Close
func (s *sshClientImplementation) Close() error {
	return s.client.Close()
}

// sshClientFactoryImplementation is the default implementation of sshClientFactory
type sshClientFactoryImplementation struct {
	sshConfig *ssh.ClientConfig
}

var _ sshClientFactory = &sshClientFactoryImplementation{}

// Dial implements sshClientFactory::Dial
func (f *sshClientFactoryImplementation) Dial(ctx context.Context, host string) (sshClient, error) {
	addr := host + ":22"
	d := net.Dialer{
		Timeout: 15 * time.Second,
	}
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}

	// We have a TCP connection; we will force-close it to support context cancellation

	var client *ssh.Client
	finished := make(chan error)
	go func() {
		c, chans, reqs, err := ssh.NewClientConn(conn, addr, f.sshConfig)
		if err == nil {
			client = ssh.NewClient(c, chans, reqs)
		}
		finished <- err
	}()

	select {
	case <-ctx.Done():
		log.Print("cancelling SSH tcp connection due to context completion")
		conn.Close() // Close the TCP connection to force cancellation
		<-finished   // Wait for cancellation
		return nil, ctx.Err()
	case err := <-finished:
		if err != nil {
			return nil, err
		}
		return &sshClientImplementation{
			client: client,
		}, nil
	}
}
<commit_msg>Include sysctls in toolbox dump<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package dump

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
)

// logDumper gets all the nodes from a kubernetes cluster and dumps a well-known set of logs
type logDumper struct {
	sshClientFactory sshClientFactory

	artifactsDir string

	services     []string
	files        []string
	podSelectors []string
}

// NewLogDumper is the constructor for a logDumper
func NewLogDumper(sshConfig *ssh.ClientConfig, artifactsDir string) *logDumper {
	sshClientFactory := &sshClientFactoryImplementation{
		sshConfig: sshConfig,
	}

	d := &logDumper{
		sshClientFactory: sshClientFactory,
		artifactsDir:     artifactsDir,
	}

	d.services = []string{
		"node-problem-detector",
		"kubelet",
		"containerd",
		"docker",
		"kops-configuration",
		"protokube",
	}
	d.files = []string{
		"kube-apiserver",
		"kube-scheduler",
		"rescheduler",
		"kube-controller-manager",
		"etcd",
		"etcd-events",
		"etcd-cilium",
		"glbc",
		"cluster-autoscaler",
		"kube-addon-manager",
		"fluentd",
		"kube-proxy",
		"node-problem-detector",
		"cloud-init-output",
		"startupscript",
		"kern",
		"docker",
		"aws-routed-eni/ipamd",
		"aws-routed-eni/plugin",
	}
	d.podSelectors = []string{
		"k8s-app=external-dns",
		"k8s-app=dns-controller",
	}

	return d
}

// DumpAllNodes connects to every node from kubectl get nodes and dumps the logs.
// additionalIPs holds IP addresses of instances found by the deployment tool;
// if the IPs are not found from kubectl get nodes, then these will be dumped also.
// This allows for dumping log on nodes even if they don't register as a kubernetes
// node, or if a node fails to register, or if the whole cluster fails to start.
func (d *logDumper) DumpAllNodes(ctx context.Context, nodes corev1.NodeList, additionalIPs []string) error {
	var dumped []*corev1.Node

	for i := range nodes.Items {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		node := &nodes.Items[i]

		ip := ""
		for _, address := range node.Status.Addresses {
			if address.Type == "ExternalIP" {
				ip = address.Address
				break
			}
		}

		err := d.dumpNode(ctx, node.Name, ip)
		if err != nil {
			log.Printf("could not dump node %s (%s): %v", node.Name, ip, err)
		} else {
			dumped = append(dumped, node)
		}
	}

	notDumped := findInstancesNotDumped(additionalIPs, dumped)
	for _, ip := range notDumped {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		log.Printf("dumping node not registered in kubernetes: %s", ip)
		err := d.dumpNode(ctx, ip, ip)
		if err != nil {
			log.Printf("error dumping node %s: %v", ip, err)
		}
	}

	return nil
}

// findInstancesNotDumped returns ips from the slice that do not appear as any address of the nodes
func findInstancesNotDumped(ips []string, dumped []*corev1.Node) []string {
	var notDumped []string
	dumpedAddresses := make(map[string]bool)
	for _, node := range dumped {
		for _, address := range node.Status.Addresses {
			dumpedAddresses[address.Address] = true
		}
	}

	for _, ip := range ips {
		if !dumpedAddresses[ip] {
			notDumped = append(notDumped, ip)
		}
	}
	return notDumped
}

// DumpNode connects to a node and dumps the logs.
func (d *logDumper) dumpNode(ctx context.Context, name string, ip string) error {
	if ip == "" {
		return fmt.Errorf("could not find address for %v, ", name)
	}

	log.Printf("Dumping node %s", name)

	n, err := d.connectToNode(ctx, name, ip)
	if err != nil {
		return fmt.Errorf("could not connect: %v", err)
	}

	// As long as we connect to the node we will not return an error;
	// a failure to collect a log (or even any logs at all) is not
	// considered an error in dumping the node.
	// TODO(justinsb): clean up / rationalize
	errors := n.dump(ctx)
	for _, e := range errors {
		log.Printf("error dumping node %s: %v", name, e)
	}

	if err := n.Close(); err != nil {
		log.Printf("error closing connection: %v", err)
	}

	return nil
}

// sshClient is an interface abstracting *ssh.Client, which allows us to test it
type sshClient interface {
	io.Closer

	// ExecPiped runs the command, piping stdout & stderr
	ExecPiped(ctx context.Context, command string, stdout io.Writer, stderr io.Writer) error
}

// sshClientFactory is an interface abstracting to a node over SSH
type sshClientFactory interface {
	Dial(ctx context.Context, host string) (sshClient, error)
}

// logDumperNode holds state for a particular node we are dumping
type logDumperNode struct {
	client sshClient
	dumper *logDumper

	dir string
}

// connectToNode makes an SSH connection to the node and returns a logDumperNode
func (d *logDumper) connectToNode(ctx context.Context, nodeName string, host string) (*logDumperNode, error) {
	client, err := d.sshClientFactory.Dial(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("unable to SSH to %q: %v", host, err)
	}
	return &logDumperNode{
		client: client,
		dir:    filepath.Join(d.artifactsDir, nodeName),
		dumper: d,
	}, nil
}

// logDumperNode cleans up any state in the logDumperNode
func (n *logDumperNode) Close() error {
	return n.client.Close()
}

// dump captures the well-known set of logs
func (n *logDumperNode) dump(ctx context.Context) []error {
	if ctx.Err() != nil {
		return []error{ctx.Err()}
	}

	var errors []error

	// Capture kernel log
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise -k", filepath.Join(n.dir, "kern.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture full journal - needed so we can see e.g. disk mounts
	// This does duplicate the other files, but ensures we have all output
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise", filepath.Join(n.dir, "journal.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture logs from any systemd services in our list that are registered
	services, err := n.listSystemdUnits(ctx)
	if err != nil {
		errors = append(errors, fmt.Errorf("error listing systemd services: %v", err))
	}
	for _, s := range n.dumper.services {
		name := s + ".service"
		for _, service := range services {
			if service == name {
				if err := n.shellToFile(ctx, "sudo journalctl --output=cat -u "+name, filepath.Join(n.dir, s+".log")); err != nil {
					errors = append(errors, err)
				}
			}
		}
	}

	// Capture any file logs where the files exist
	fileList, err := n.findFiles(ctx, "/var/log")
	if err != nil {
		errors = append(errors, fmt.Errorf("error reading /var/log: %v", err))
	}
	for _, name := range n.dumper.files {
		prefix := "/var/log/" + name + ".log"
		for _, f := range fileList {
			if !strings.HasPrefix(f, prefix) {
				continue
			}
			if err := n.shellToFile(ctx, "sudo cat '"+strings.ReplaceAll(f, "'", "'\\''")+"'", filepath.Join(n.dir, strings.ReplaceAll(strings.TrimPrefix(f, "/var/log/"), "/", "_"))); err != nil {
				errors = append(errors, err)
			}
		}
	}

	for _, selector := range n.dumper.podSelectors {
		kv := strings.Split(selector, "=")
		logFile := fmt.Sprintf("%v.log", kv[len(kv)-1])
		if err := n.shellToFile(ctx, "if command -v kubectl &> /dev/null; then kubectl logs -n kube-system --all-containers -l \""+selector+"\"; fi", filepath.Join(n.dir, logFile)); err != nil {
			errors = append(errors, err)
		}
	}

	if err := n.shellToFile(ctx, "cat /etc/hosts", filepath.Join(n.dir, "etchosts")); err != nil {
		errors = append(errors, err)
	}
	if err := n.shellToFile(ctx, "sysctl -a", filepath.Join(n.dir, "sysctls")); err != nil {
		errors = append(errors, err)
	}

	return errors
}

// findFiles lists files under the specified directory (recursively)
func (n *logDumperNode) findFiles(ctx context.Context, dir string) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo find "+dir+" -print0", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing %q: %v", dir, err)
	}

	paths := []string{}
	for _, b := range bytes.Split(stdout.Bytes(), []byte{0}) {
		if len(b) == 0 {
			// Likely the last value
			continue
		}
		paths = append(paths, string(b))
	}
	return paths, nil
}

// listSystemdUnits returns the list of systemd units on the node
func (n *logDumperNode) listSystemdUnits(ctx context.Context) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo systemctl list-units -t service --no-pager --no-legend --all", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing systemd units: %v", err)
	}

	var services []string
	for _, line := range strings.Split(stdout.String(), "\n") {
		tokens := strings.Fields(line)
		if len(tokens) == 0 || tokens[0] == "" {
			continue
		}
		services = append(services, tokens[0])
	}
	return services, nil
}

// shellToFile executes a command and copies the output to a file
func (n *logDumperNode) shellToFile(ctx context.Context, command string, destPath string) error {
	if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil {
		log.Printf("unable to mkdir on %q: %v", filepath.Dir(destPath), err)
	}

	f, err := os.Create(destPath)
	if err != nil {
		return fmt.Errorf("error creating file %q: %v", destPath, err)
	}
	defer f.Close()

	if err := n.client.ExecPiped(ctx, command, f, f); err != nil {
		return fmt.Errorf("error executing command %q: %v", command, err)
	}

	return nil
}

// sshClientImplementation is the default implementation of sshClient, binding to a *ssh.Client
type sshClientImplementation struct {
	client *ssh.Client
}

var _ sshClient = &sshClientImplementation{}

// ExecPiped implements sshClientImplementation::ExecPiped
func (s *sshClientImplementation) ExecPiped(ctx context.Context, cmd string, stdout io.Writer, stderr io.Writer) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	finished := make(chan error)
	go func() {
		session, err := s.client.NewSession()
		if err != nil {
			finished <- fmt.Errorf("error creating ssh session: %v", err)
			return
		}
		defer session.Close()

		klog.V(2).Infof("running SSH command: %v", cmd)

		session.Stdout = stdout
		session.Stderr = stderr

		finished <- session.Run(cmd)
	}()

	select {
	case <-ctx.Done():
		log.Print("closing SSH tcp connection due to context completion")

		// terminate the TCP connection to force a disconnect - we assume everyone is using the same context.
		// We could make this better by sending a signal on the session, waiting and then closing the session,
		// and only if we still haven't succeeded then closing the TCP connection.  This is sufficient for our
		// current usage though - and hopefully that logic will be implemented in the SSH package itself.
		s.Close()

		<-finished // Wait for cancellation
		return ctx.Err()

	case err := <-finished:
		return err
	}
}

// Close implements sshClientImplementation::Close
func (s *sshClientImplementation) Close() error {
	return s.client.Close()
}

// sshClientFactoryImplementation is the default implementation of sshClientFactory
type sshClientFactoryImplementation struct {
	sshConfig *ssh.ClientConfig
}

var _ sshClientFactory = &sshClientFactoryImplementation{}

// Dial implements sshClientFactory::Dial
func (f *sshClientFactoryImplementation) Dial(ctx context.Context, host string) (sshClient, error) {
	addr := host + ":22"
	d := net.Dialer{
		Timeout: 15 * time.Second,
	}
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}

	// We have a TCP connection; we will force-close it to support context cancellation

	var client *ssh.Client
	finished := make(chan error)
	go func() {
		c, chans, reqs, err := ssh.NewClientConn(conn, addr, f.sshConfig)
		if err == nil {
			client = ssh.NewClient(c, chans, reqs)
		}
		finished <- err
	}()

	select {
	case <-ctx.Done():
		log.Print("cancelling SSH tcp connection due to context completion")
		conn.Close() // Close the TCP connection to force cancellation
		<-finished   // Wait for cancellation
		return nil, ctx.Err()
	case err := <-finished:
		if err != nil {
			return nil, err
		}
		return &sshClientImplementation{
			client: client,
		}, nil
	}
}
<|endoftext|>


<commit_before>package helm

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/Masterminds/semver"
	"sigs.k8s.io/yaml"
)

// LoadChartfile opens a Chartfile tree
func LoadChartfile(projectRoot string) (*Charts, error) {
	// make sure project root is valid
	abs, err := filepath.Abs(projectRoot)
	if err != nil {
		return nil, err
	}

	// open chartfile
	chartfile := filepath.Join(abs, Filename)
	data, err := ioutil.ReadFile(chartfile)
	if err != nil {
		return nil, err
	}

	// parse it
	c := Chartfile{
		Version:   Version,
		Directory: DefaultDir,
	}
	if err := yaml.UnmarshalStrict(data, &c); err != nil {
		return nil, err
	}

	for i, r := range c.Requires {
		if r.Chart == "" {
			return nil, fmt.Errorf("requirements[%v]: 'chart' must be set", i)
		}
	}

	// return Charts handle
	charts := &Charts{
		Manifest:    c,
		projectRoot: abs,

		// default to ExecHelm, but allow injecting from the outside
		Helm: ExecHelm{},
	}
	return charts, nil
}

// Charts exposes the central Chartfile management functions
type Charts struct {
	// Manifest are the chartfile.yaml contents. It holds data about the developers intentions
	Manifest Chartfile

	// projectRoot is the enclosing directory of chartfile.yaml
	projectRoot string

	// Helm is the helm implementation underneath. ExecHelm is the default, but
	// any implementation of the Helm interface may be used
	Helm Helm
}

// ChartDir returns the directory pulled charts are saved in
func (c Charts) ChartDir() string {
	return filepath.Join(c.projectRoot, c.Manifest.Directory)
}

// ManifestFile returns the full path to the chartfile.yaml
func (c Charts) ManifestFile() string {
	return filepath.Join(c.projectRoot, Filename)
}

// Vendor pulls all Charts specified in the manifest into the local charts
// directory. It fetches the repository index before doing so.
func (c Charts) Vendor() error {
	dir := c.ChartDir()
	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
		return err
	}

	log.Println("Syncing Repositories ...")
	if err := c.Helm.RepoUpdate(Opts{Repositories: c.Manifest.Repositories}); err != nil {
		return err
	}

	log.Println("Pulling Charts ...")
	for _, r := range c.Manifest.Requires {
		err := c.Helm.Pull(r.Chart, r.Version.String(), PullOpts{
			Destination: dir,
			Opts:        Opts{Repositories: c.Manifest.Repositories},
		})
		if err != nil {
			return err
		}

		log.Printf(" %s@%s", r.Chart, r.Version.String())
	}

	return nil
}

// Add adds every Chart in reqs to the Manifest after validation, and runs
// Vendor afterwards
func (c Charts) Add(reqs []string) error {
	log.Printf("Adding %v Charts ...", len(reqs))

	skip := func(s string, err error) {
		log.Printf(" Skipping %s: %s.", s, err)
	}

	// parse new charts, append in memory
	added := 0
	for _, s := range reqs {
		r, err := parseReq(s)
		if err != nil {
			skip(s, err)
			continue
		}

		if c.Manifest.Requires.Has(*r) {
			skip(s, fmt.Errorf("already exists"))
			continue
		}

		c.Manifest.Requires = append(c.Manifest.Requires, *r)
		added++
		log.Println(" OK:", s)
	}

	// write out
	if err := write(c.Manifest, c.ManifestFile()); err != nil {
		return err
	}

	// skipped some? fail then
	if added != len(reqs) {
		return fmt.Errorf("%v Charts were skipped. Please check above logs for details", len(reqs)-added)
	}

	// worked fine? vendor it
	log.Printf("Added %v Charts to helmfile.yaml. Vendoring ...", added)
	return c.Vendor()
}

func InitChartfile(path string) (*Charts, error) {
	c := Chartfile{
		Version: Version,
		Repositories: []Repo{{
			Name: "stable",
			URL:  "https://kubernetes-charts.storage.googleapis.com",
		}},
		Requires: make(Requirements, 0),
	}

	if err := write(c, path); err != nil {
		return nil, err
	}

	return LoadChartfile(filepath.Dir(path))
}

// write saves a Chartfile to dest
func write(c Chartfile, dest string) error {
	data, err := yaml.Marshal(c)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(dest, data, 0644)
}

var chartExp = regexp.MustCompile(`\w+\/.+@.+`)

// parseReq parses a requirement from a string of the format `repo/name@version`
func parseReq(s string) (*Requirement, error) {
	if !chartExp.MatchString(s) {
		return nil, fmt.Errorf("not of form 'repo/chart@version'")
	}

	elems := strings.Split(s, "@")
	chart := elems[0]
	ver, err := semver.NewVersion(elems[1])
	if errors.Is(err, semver.ErrInvalidSemVer) {
		return nil, fmt.Errorf("version is invalid")
	} else if err != nil {
		return nil, fmt.Errorf("version is invalid: %s", err)
	}

	return &Requirement{
		Chart:   chart,
		Version: *ver,
	}, nil
}
<commit_msg>fix(helm): Detect already pulled charts (#402)<commit_after>package helm

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/Masterminds/semver"
	"sigs.k8s.io/yaml"
)

// LoadChartfile opens a Chartfile tree
func LoadChartfile(projectRoot string) (*Charts, error) {
	// make sure project root is valid
	abs, err := filepath.Abs(projectRoot)
	if err != nil {
		return nil, err
	}

	// open chartfile
	chartfile := filepath.Join(abs, Filename)
	data, err := ioutil.ReadFile(chartfile)
	if err != nil {
		return nil, err
	}

	// parse it
	c := Chartfile{
		Version:   Version,
		Directory: DefaultDir,
	}
	if err := yaml.UnmarshalStrict(data, &c); err != nil {
		return nil, err
	}

	for i, r := range c.Requires {
		if r.Chart == "" {
			return nil, fmt.Errorf("requirements[%v]: 'chart' must be set", i)
		}
	}

	// return Charts handle
	charts := &Charts{
		Manifest:    c,
		projectRoot: abs,

		// default to ExecHelm, but allow injecting from the outside
		Helm: ExecHelm{},
	}
	return charts, nil
}

// Charts exposes the central Chartfile management functions
type Charts struct {
	// Manifest are the chartfile.yaml contents. It holds data about the developers intentions
	Manifest Chartfile

	// projectRoot is the enclosing directory of chartfile.yaml
	projectRoot string

	// Helm is the helm implementation underneath. ExecHelm is the default, but
	// any implementation of the Helm interface may be used
	Helm Helm
}

// ChartDir returns the directory pulled charts are saved in
func (c Charts) ChartDir() string {
	return filepath.Join(c.projectRoot, c.Manifest.Directory)
}

// ManifestFile returns the full path to the chartfile.yaml
func (c Charts) ManifestFile() string {
	return filepath.Join(c.projectRoot, Filename)
}

// Vendor pulls all Charts specified in the manifest into the local charts
// directory. It fetches the repository index before doing so.
func (c Charts) Vendor() error {
	dir := c.ChartDir()
	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
		return err
	}

	log.Println("Syncing Repositories ...")
	if err := c.Helm.RepoUpdate(Opts{Repositories: c.Manifest.Repositories}); err != nil {
		return err
	}

	log.Println("Pulling Charts ...")
	for _, r := range c.Manifest.Requires {
		chartName := parseReqName(r.Chart)
		chartPath := filepath.Join(dir, chartName)
		if _, err := os.Stat(chartPath); err == nil {
			log.Printf(" %s@%s exists", r.Chart, r.Version.String())
			continue
		}

		err := c.Helm.Pull(r.Chart, r.Version.String(), PullOpts{
			Destination: dir,
			Opts:        Opts{Repositories: c.Manifest.Repositories},
		})
		if err != nil {
			return err
		}

		log.Printf(" %s@%s downloaded", r.Chart, r.Version.String())
	}

	return nil
}

// Add adds every Chart in reqs to the Manifest after validation, and runs
// Vendor afterwards
func (c Charts) Add(reqs []string) error {
	log.Printf("Adding %v Charts ...", len(reqs))

	skip := func(s string, err error) {
		log.Printf(" Skipping %s: %s.", s, err)
	}

	// parse new charts, append in memory
	added := 0
	for _, s := range reqs {
		r, err := parseReq(s)
		if err != nil {
			skip(s, err)
			continue
		}

		if c.Manifest.Requires.Has(*r) {
			skip(s, fmt.Errorf("already exists"))
			continue
		}

		c.Manifest.Requires = append(c.Manifest.Requires, *r)
		added++
		log.Println(" OK:", s)
	}

	// write out
	if err := write(c.Manifest, c.ManifestFile()); err != nil {
		return err
	}

	// skipped some? fail then
	if added != len(reqs) {
		return fmt.Errorf("%v Charts were skipped. Please check above logs for details", len(reqs)-added)
	}

	// worked fine? vendor it
	log.Printf("Added %v Charts to helmfile.yaml. Vendoring ...", added)
	return c.Vendor()
}

func InitChartfile(path string) (*Charts, error) {
	c := Chartfile{
		Version: Version,
		Repositories: []Repo{{
			Name: "stable",
			URL:  "https://kubernetes-charts.storage.googleapis.com",
		}},
		Requires: make(Requirements, 0),
	}

	if err := write(c, path); err != nil {
		return nil, err
	}

	return LoadChartfile(filepath.Dir(path))
}

// write saves a Chartfile to dest
func write(c Chartfile, dest string) error {
	data, err := yaml.Marshal(c)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(dest, data, 0644)
}

var chartExp = regexp.MustCompile(`\w+\/.+@.+`)

// parseReq parses a requirement from a string of the format `repo/name@version`
func parseReq(s string) (*Requirement, error) {
	if !chartExp.MatchString(s) {
		return nil, fmt.Errorf("not of form 'repo/chart@version'")
	}

	elems := strings.Split(s, "@")
	chart := elems[0]
	ver, err := semver.NewVersion(elems[1])
	if errors.Is(err, semver.ErrInvalidSemVer) {
		return nil, fmt.Errorf("version is invalid")
	} else if err != nil {
		return nil, fmt.Errorf("version is invalid: %s", err)
	}

	return &Requirement{
		Chart:   chart,
		Version: *ver,
	}, nil
}

// parseReqName parses a name from a string of the format `repo/name`
func parseReqName(s string) string {
	elems := strings.Split(s, "/")
	name := elems[1]
	return name
}
<|endoftext|>


<commit_before>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package models

import (
	"encoding/json"
	"fmt"
	"path"

	"github.com/ngaut/zkhelper"

	"github.com/juju/errors"
)

type SlotStatus string

const (
	SLOT_STATUS_ONLINE      SlotStatus = "online"
	SLOT_STATUS_OFFLINE     SlotStatus = "offline"
	SLOT_STATUS_MIGRATE     SlotStatus = "migrate"
	SLOT_STATUS_PRE_MIGRATE SlotStatus = "pre_migrate"
)

var ErrSlotAlreadyExists = errors.New("slots already exists")
var ErrUnknownSlotStatus = errors.New("unknown slot status, slot status should be (online, offline, migrate, pre_migrate)")

type SlotMigrateStatus struct {
	From int `json:"from"`
	To   int `json:"to"`
}

type SlotMultiSetParam struct {
	From    int        `json:"from"`
	To      int        `json:"to"`
	Status  SlotStatus `json:"status"`
	GroupId int        `json:"group_id"`
}

type SlotState struct {
	Status        SlotStatus        `json:"status"`
	MigrateStatus SlotMigrateStatus `json:"migrate_status"`
	LastOpTs      string            `json:"last_op_ts"` // operation timestamp
}

type Slot struct {
	ProductName string    `json:"product_name"`
	Id          int       `json:"id"`
	GroupId     int       `json:"group_id"`
	State       SlotState `json:"state"`
}

func (s *Slot) String() string {
	b, _ := json.MarshalIndent(s, "", "  ")
	return string(b)
}

func NewSlot(productName string, id int) *Slot {
	return &Slot{
		ProductName: productName,
		Id:          id,
		GroupId:     INVALID_ID,
		State: SlotState{
			Status:   SLOT_STATUS_OFFLINE,
			LastOpTs: "0",
			MigrateStatus: SlotMigrateStatus{
				From: INVALID_ID,
				To:   INVALID_ID,
			},
		},
	}
}

func GetSlotPath(productName string, slotId int) string {
	return fmt.Sprintf("/zk/codis/db_%s/slots/slot_%d", productName, slotId)
}

func GetSlotBasePath(productName string) string {
	return fmt.Sprintf("/zk/codis/db_%s/slots", productName)
}

func GetSlot(zkConn zkhelper.Conn, productName string, id int) (*Slot, error) {
	zkPath := GetSlotPath(productName, id)
	data, _, err := zkConn.Get(zkPath)
	if err != nil {
		return nil, err
	}
	slot := Slot{}
	if err := json.Unmarshal(data, &slot); err != nil {
		return nil, err
	}
	return &slot, nil
}

func GetMigratingSlots(conn zkhelper.Conn, productName string) ([]Slot, error) {
	var migrateSlots []Slot = make([]Slot, 0)
	slots, err := Slots(conn, productName)
	if err != nil {
		return nil, err
	}
	for _, slot := range slots {
		if slot.State.Status == SLOT_STATUS_MIGRATE {
			migrateSlots = append(migrateSlots, slot)
		}
	}
	return migrateSlots, nil
}

func Slots(zkConn zkhelper.Conn, productName string) ([]Slot, error) {
	zkPath := GetSlotBasePath(productName)
	children, _, err := zkConn.Children(zkPath)
	if err != nil {
		return nil, errors.Trace(err)
	}

	var slots []Slot

	for _, p := range children {
		data, _, err := zkConn.Get(path.Join(zkPath, p))
		if err != nil {
			return nil, errors.Trace(err)
		}
		slot := Slot{}
		if err := json.Unmarshal(data, &slot); err != nil {
			return nil, errors.Trace(err)
		}
		slots = append(slots, slot)
	}

	return slots, nil
}

func NoGroupSlots(zkConn zkhelper.Conn, productName string) ([]Slot, error) {
	slots, err := Slots(zkConn, productName)
	if err != nil {
		return nil, errors.Trace(err)
	}
	var ret []Slot
	for _, slot := range slots {
		if slot.GroupId == INVALID_ID {
			ret = append(ret, slot)
		}
	}
	return ret, nil
}

func SetSlots(zkConn zkhelper.Conn, productName string, slots []Slot, groupId int, status SlotStatus) error {
	if status != SLOT_STATUS_OFFLINE && status != SLOT_STATUS_ONLINE {
		return errors.New("invalid status")
	}
	for _, s := range slots {
		s.GroupId = groupId
		s.State.Status = status
		data, err := json.Marshal(s)
		if err != nil {
			return errors.Trace(err)
		}

		zkPath := GetSlotPath(productName, s.Id)
		_, err = zkhelper.CreateOrUpdate(zkConn, zkPath, string(data), 0, zkhelper.DefaultFileACLs(), true)
		if err != nil {
			return errors.Trace(err)
		}
	}

	param := SlotMultiSetParam{
		From:    -1,
		To:      -1,
		GroupId: groupId,
		Status:  status,
	}

	err := NewAction(zkConn, productName, ACTION_TYPE_MULTI_SLOT_CHANGED, param, "", true)
	return errors.Trace(err)

}

func SetSlotRange(zkConn zkhelper.Conn, productName string, fromSlot, toSlot, groupId int, status SlotStatus) error {
	if status != SLOT_STATUS_OFFLINE && status != SLOT_STATUS_ONLINE {
		return errors.New("invalid status")
	}

	for i := fromSlot; i <= toSlot; i++ {
		s, err := GetSlot(zkConn, productName, i)
		if err != nil {
			return errors.Trace(err)
		}
		s.GroupId = groupId
		s.State.Status = status
		data, err := json.Marshal(s)
		if err != nil {
			return errors.Trace(err)
		}

		zkPath := GetSlotPath(productName, i)
		_, err = zkhelper.CreateOrUpdate(zkConn, zkPath, string(data), 0, zkhelper.DefaultFileACLs(), true)
		if err != nil {
			return errors.Trace(err)
		}
	}

	param := SlotMultiSetParam{
		From:    fromSlot,
		To:      toSlot,
		GroupId: groupId,
		Status:  status,
	}
	err := NewAction(zkConn, productName, ACTION_TYPE_MULTI_SLOT_CHANGED, param, "", true)
	return errors.Trace(err)
}

// danger operation !
func InitSlotSet(zkConn zkhelper.Conn, productName string, totalSlotNum int) error {
	for i := 0; i < totalSlotNum; i++ {
		slot := NewSlot(productName, i)
		if err := slot.Update(zkConn); err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func (s *Slot) SetMigrateStatus(zkConn zkhelper.Conn, fromGroup, toGroup int) error {
	if fromGroup < 0 || toGroup < 0 {
		return errors.Errorf("invalid group id, from %d, to %d", fromGroup, toGroup)
	}
	// wait until all proxy confirmed
	err := NewAction(zkConn, s.ProductName, ACTION_TYPE_SLOT_PREMIGRATE, s, "", true)
	if err != nil {
		return errors.Trace(err)
	}

	s.State.Status = SLOT_STATUS_MIGRATE
	s.State.MigrateStatus.From = fromGroup
	s.State.MigrateStatus.To = toGroup

	s.GroupId = toGroup

	return s.Update(zkConn)
}

func (s *Slot) Update(zkConn zkhelper.Conn) error {
	// status validation
	switch s.State.Status {
	case SLOT_STATUS_MIGRATE, SLOT_STATUS_OFFLINE,
		SLOT_STATUS_ONLINE, SLOT_STATUS_PRE_MIGRATE:
		{
			// valid status, OK
		}
	default:
		{
			return errors.Trace(ErrUnknownSlotStatus)
		}
	}
	var err error
	data, err := json.Marshal(s)
	if err != nil {
		return errors.Trace(err)
	}
	zkPath := GetSlotPath(s.ProductName, s.Id)
	_, err = zkhelper.CreateOrUpdate(zkConn, zkPath, string(data), 0, zkhelper.DefaultFileACLs(), true)
	if err != nil {
		return errors.Trace(err)
	}

	if s.State.Status == SLOT_STATUS_MIGRATE {
		err = NewAction(zkConn, s.ProductName, ACTION_TYPE_SLOT_MIGRATE, s, "", true)
	} else {
		err = NewAction(zkConn, s.ProductName, ACTION_TYPE_SLOT_CHANGED, s, "", true)
	}

	return errors.Trace(err)
}
<commit_msg>clean up<commit_after>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package models

import (
	"encoding/json"
	"fmt"
	"path"

	"github.com/ngaut/zkhelper"

	"github.com/juju/errors"
)

type SlotStatus string

const (
	SLOT_STATUS_ONLINE      SlotStatus = "online"
	SLOT_STATUS_OFFLINE     SlotStatus = "offline"
	SLOT_STATUS_MIGRATE     SlotStatus = "migrate"
	SLOT_STATUS_PRE_MIGRATE SlotStatus = "pre_migrate"
)

var ErrSlotAlreadyExists = errors.New("slots already exists")
var ErrUnknownSlotStatus = errors.New("unknown slot status, slot status should be (online, offline, migrate, pre_migrate)")

type SlotMigrateStatus struct {
	From int `json:"from"`
	To   int `json:"to"`
}

type SlotMultiSetParam struct {
	From    int        `json:"from"`
	To      int        `json:"to"`
	Status  SlotStatus `json:"status"`
	GroupId int        `json:"group_id"`
}

type SlotState struct {
	Status        SlotStatus        `json:"status"`
	MigrateStatus SlotMigrateStatus `json:"migrate_status"`
	LastOpTs      string            `json:"last_op_ts"` // operation timestamp
}

type Slot struct {
	ProductName string    `json:"product_name"`
	Id          int       `json:"id"`
	GroupId     int       `json:"group_id"`
	State       SlotState `json:"state"`
}

func (s *Slot) String() string {
	b, _ := json.MarshalIndent(s, "", "  ")
	return string(b)
}

func NewSlot(productName string, id int) *Slot {
	return &Slot{
		ProductName: productName,
		Id:          id,
		GroupId:     INVALID_ID,
		State: SlotState{
			Status:   SLOT_STATUS_OFFLINE,
			LastOpTs: "0",
			MigrateStatus: SlotMigrateStatus{
				From: INVALID_ID,
				To:   INVALID_ID,
			},
		},
	}
}

func GetSlotPath(productName string, slotId int) string {
	return fmt.Sprintf("/zk/codis/db_%s/slots/slot_%d", productName, slotId)
}

func GetSlotBasePath(productName string) string {
	return fmt.Sprintf("/zk/codis/db_%s/slots", productName)
}

func GetSlot(zkConn zkhelper.Conn, productName string, id int) (*Slot, error) {
	zkPath := GetSlotPath(productName, id)
	data, _, err := zkConn.Get(zkPath)
	if err != nil {
		return nil, err
	}

	slot := Slot{}
	if err := json.Unmarshal(data, &slot); err != nil {
		return nil, err
	}

	return &slot, nil
}

func GetMigratingSlots(conn zkhelper.Conn, productName string) ([]Slot, error) {
	var migrateSlots []Slot = make([]Slot, 0)
	slots, err := Slots(conn, productName)
	if err != nil {
		return nil, err
	}

	for _, slot := range slots {
		if slot.State.Status == SLOT_STATUS_MIGRATE {
			migrateSlots = append(migrateSlots, slot)
		}
	}

	return migrateSlots, nil
}

func Slots(zkConn zkhelper.Conn, productName string) ([]Slot, error) {
	zkPath := GetSlotBasePath(productName)
	children, _, err := zkConn.Children(zkPath)
	if err != nil {
		return nil, errors.Trace(err)
	}

	var slots []Slot
	for _, p := range children {
		data, _, err := zkConn.Get(path.Join(zkPath, p))
		if err != nil {
			return nil, errors.Trace(err)
		}
		slot := Slot{}
		if err := json.Unmarshal(data, &slot); err != nil {
			return nil, errors.Trace(err)
		}
		slots = append(slots, slot)
	}

	return slots, nil
}

func NoGroupSlots(zkConn zkhelper.Conn, productName string) ([]Slot, error) {
	slots, err := Slots(zkConn, productName)
	if err != nil {
		return nil, errors.Trace(err)
	}

	var ret []Slot
	for _, slot := range slots {
		if slot.GroupId == INVALID_ID {
			ret = append(ret, slot)
		}
	}
	return ret, nil
}

func SetSlots(zkConn zkhelper.Conn, productName string, slots []Slot, groupId int, status SlotStatus) error {
	if status != SLOT_STATUS_OFFLINE && status != SLOT_STATUS_ONLINE {
		return errors.New("invalid status")
	}

	for _, s := range slots {
		s.GroupId = groupId
		s.State.Status = status
		data, err := json.Marshal(s)
		if err != nil {
			return errors.Trace(err)
		}

		zkPath := GetSlotPath(productName, s.Id)
		_, err = zkhelper.CreateOrUpdate(zkConn, zkPath, string(data), 0, zkhelper.DefaultFileACLs(), true)
		if err != nil {
			return errors.Trace(err)
		}
	}

	param := SlotMultiSetParam{
		From:    -1,
		To:      -1,
		GroupId: groupId,
		Status:  status,
	}

	err := NewAction(zkConn, productName, ACTION_TYPE_MULTI_SLOT_CHANGED, param, "", true)
	return errors.Trace(err)

}

func SetSlotRange(zkConn zkhelper.Conn, productName string, fromSlot, toSlot, groupId int, status SlotStatus) error {
	if status != SLOT_STATUS_OFFLINE && status != SLOT_STATUS_ONLINE {
		return errors.New("invalid status")
	}

	for i := fromSlot; i <= toSlot; i++ {
		s, err := GetSlot(zkConn, productName, i)
		if err != nil {
			return errors.Trace(err)
		}
		s.GroupId = groupId
		s.State.Status = status
		data, err := json.Marshal(s)
		if err != nil {
			return errors.Trace(err)
		}

		zkPath := GetSlotPath(productName, i)
		_, err = zkhelper.CreateOrUpdate(zkConn, zkPath, string(data), 0, zkhelper.DefaultFileACLs(), true)
		if err != nil {
			return errors.Trace(err)
		}
	}

	param := SlotMultiSetParam{
		From:    fromSlot,
		To:      toSlot,
		GroupId: groupId,
		Status:  status,
	}
	err := NewAction(zkConn, productName, ACTION_TYPE_MULTI_SLOT_CHANGED, param, "", true)
	return errors.Trace(err)
}

// danger operation !
func InitSlotSet(zkConn zkhelper.Conn, productName string, totalSlotNum int) error {
	for i := 0; i < totalSlotNum; i++ {
		slot := NewSlot(productName, i)
		if err := slot.Update(zkConn); err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func (s *Slot) SetMigrateStatus(zkConn zkhelper.Conn, fromGroup, toGroup int) error {
	if fromGroup < 0 || toGroup < 0 {
		return errors.Errorf("invalid group id, from %d, to %d", fromGroup, toGroup)
	}
	// wait until all proxy confirmed
	err := NewAction(zkConn, s.ProductName, ACTION_TYPE_SLOT_PREMIGRATE, s, "", true)
	if err != nil {
		return errors.Trace(err)
	}

	s.State.Status = SLOT_STATUS_MIGRATE
	s.State.MigrateStatus.From = fromGroup
	s.State.MigrateStatus.To = toGroup

	s.GroupId = toGroup

	return s.Update(zkConn)
}

func (s *Slot) Update(zkConn zkhelper.Conn) error {
	// status validation
	switch s.State.Status {
	case SLOT_STATUS_MIGRATE, SLOT_STATUS_OFFLINE,
		SLOT_STATUS_ONLINE, SLOT_STATUS_PRE_MIGRATE:
		{
			// valid status, OK
		}
	default:
		{
			return errors.Trace(ErrUnknownSlotStatus)
		}
	}

	data, err := json.Marshal(s)
	if err != nil {
		return errors.Trace(err)
	}
	zkPath := GetSlotPath(s.ProductName, s.Id)
	_, err = zkhelper.CreateOrUpdate(zkConn, zkPath, string(data), 0, zkhelper.DefaultFileACLs(), true)
	if err != nil {
		return errors.Trace(err)
	}

	if s.State.Status == SLOT_STATUS_MIGRATE {
		err = NewAction(zkConn, s.ProductName, ACTION_TYPE_SLOT_MIGRATE, s, "", true)
	} else {
		err = NewAction(zkConn, s.ProductName, ACTION_TYPE_SLOT_CHANGED, s, "", true)
	}

	return errors.Trace(err)
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Kubernetes Authors.

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

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

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

package spdx

import (
	"bufio"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// OpenDoc opens a file, parses a SPDX tag-value file and returns a loaded
// spdx.Document object. This functions has the cyclomatic chec disabled as
// it spans specific cases for each of the tags it recognizes.
// nolint:gocyclo
func OpenDoc(path string) (*Document, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, errors.Wrapf(err, "opening document from %s", path)
	}
	defer file.Close()

	// Create a blank document
	doc := &Document{
		Packages:        map[string]*Package{},
		Files:           map[string]*File{},
		ExternalDocRefs: []ExternalDocumentRef{},
	}

	// Regexp to match the tag-value spdx expressions
	tagRegExp := regexp.MustCompile(`^([a-z0-9A-Z]+):\s+(.+)`)
	relationshioRegExp := regexp.MustCompile(`^*(\S+)\s+([_A-Z]+)\s+(\S+)`)

	// Scan the file, looking for tags
	scanner := bufio.NewScanner(file)
	i := 0 // Line counter
	var currentEntity *Entity
	var currentObject Object
	var value, tag, textValue string
	var captureMultiline bool
	objects := map[string]Object{}
	rels := []struct {
		Source       string
		Relationship string
		Peer         string
		ExtDoc       string
	}{}
	for scanner.Scan() {
		// If we are capturing text for a multiline value, read and add
		// the line to the buffer
		if captureMultiline {
			textValue += scanner.Text() + "\n"
			// If we closing tag is not here, continue to the next line
			if !strings.Contains(scanner.Text(), "</text>") {
				continue
			}

			// If closing tag found, remove it from value
			value = strings.ReplaceAll(textValue, "</text>", "")
			textValue = ""
		}

		// Check if line matches if we are not reading multiline values
		if !captureMultiline {
			match := tagRegExp.FindStringSubmatch(scanner.Text())
			if len(match) != 3 {
				continue
			}
			tag = match[1]
			value = match[2]

			// If it is a multiline value, start buffering it
			if strings.HasPrefix(value, "<text>") {
				textValue = strings.ReplaceAll(value, "<text>", "") + "\n"
				captureMultiline = true

				// It may be that the closing tag is right in the same
				// line. If so, capture and finish buffering
				if strings.Contains(scanner.Text(), "</text>") {
					value = strings.ReplaceAll(textValue, "</text>", "")
					textValue = ""
				} else {
					continue
				}
			}
		}

		captureMultiline = false

		switch tag {
		case "FileName", "PackageName":
			// Both FileName or PackageName signal the start of a new entity

			// If we have an object, we store it and continue
			if currentObject != nil {
				currentObject.SetEntity(currentEntity)

				if _, ok := objects[currentObject.SPDXID()]; ok {
					return nil, errors.Errorf("Duplicate SPDXID %s", currentObject.SPDXID())
				}

				objects[currentObject.SPDXID()] = currentObject
			}

			// Create the new entity:
			currentEntity = &Entity{}

			// And the new SPDX object:
			if tag == "FileName" {
				currentObject = &File{}
			}
			if tag == "PackageName" {
				currentObject = &Package{}
			}
			currentEntity.Name = value

		case "SPDXID":
			logrus.Debugf("Entity ID %s", value)
			if currentEntity == nil {
				doc.ID = value
			} else {
				currentEntity.ID = value
			}
		case "PackageLicenseConcluded", "LicenseConcluded":
			if value != NOASSERTION {
				currentEntity.LicenseConcluded = value
			}
		case "PackageCopyrightText", "FileCopyrightText":
			if value != NOASSERTION {
				currentEntity.CopyrightText = value
			}
			// Tags for packages
		case "FilesAnalyzed":
			currentObject.(*Package).FilesAnalyzed = value == "true"
		case "PackageVersion":
			currentObject.(*Package).Version = value
		case "PackageLicenseDeclared":
			currentObject.(*Package).LicenseDeclared = value
		case "PackageVerificationCode":
			currentObject.(*Package).VerificationCode = value
		case "PackageComment":
			currentObject.(*Package).Comment = value
		case "PackageFileName":
			currentObject.(*Package).FileName = value
		case "PackageLicenseInfoFromFiles":
			have := false
			// Check if we already have the license
			for _, licid := range currentObject.(*Package).LicenseInfoFromFiles {
				if licid == value {
					have = true
					break
				}
			}
			if have {
				currentObject.(*Package).LicenseInfoFromFiles = append(currentObject.(*Package).LicenseInfoFromFiles, value)
			}
		case "LicenseInfoInFile":
			if value != NONE {
				currentObject.(*File).LicenseInfoInFile = value
			}
		case "FileChecksum", "PackageChecksum":
			// Checksums are also tag/value -> algo/hash
			match := tagRegExp.FindStringSubmatch(value)
			if len(match) != 3 {
				return nil, errors.Errorf("invalid checksum tag syntax at line %d", i)
			}
			if currentEntity.Checksum == nil {
				currentEntity.Checksum = map[string]string{}
			}
			currentEntity.Checksum[match[1]] = match[2]
		case "Relationship":
			matches := relationshioRegExp.FindStringSubmatch(value)
			if len(matches) != 4 {
				return nil, errors.Errorf("invalid SPDX relationship on line %d: %s", i, value)
			}

			// Check if the relationship is external
			ext := ""
			if strings.HasPrefix(matches[3], "DocumentRef-") && strings.Contains(matches[3], ":") {
				parts := strings.Split(matches[3], ":")
				if len(parts) != 2 {
					return nil, errors.Wrapf(err, "Unable to parse external document reference %s", matches[3])
				}
				matches[3] = parts[0]
				ext = parts[1]
			}

			// Parse the relationship
			rels = append(rels, struct {
				Source       string
				Relationship string
				Peer         string
				ExtDoc       string
			}{
				matches[1], matches[2], matches[3], ext,
			})
		case "PackageDownloadLocation":
			currentEntity.DownloadLocation = value
			// Tags that apply top the doc
		case "Created":
			t, err := time.Parse("2006-01-02T15:04:05Z", value)
			if err != nil {
				return nil, errors.Wrapf(err, "parsing time string in file: %s", value)
			}
			doc.Created = t
		case "Creator":
			// Creator has a tag/value format inside
			match := tagRegExp.FindStringSubmatch(value)
			if len(match) != 3 {
				return nil, errors.Errorf("invalid creator tag syntax at line %d", i)
			}
			switch match[1] {
			case "Person":
				doc.Creator.Person = match[2]
			case "Tool":
				doc.Creator.Tool = append(doc.Creator.Tool, match[2])
			case "Organization":
				doc.Creator.Organization = match[2]
			default:
				return nil, errors.Errorf(
					"invalid creator tag '%s' syntax at line %d, valid values are 'Tool', 'Organization' or 'Person'",
					match[1], i,
				)
			}
		case "DataLicense":
			doc.DataLicense = value
		case "DocumentName":
			doc.Name = value
		case "DocumentNamespace":
			doc.Namespace = value
		case "SPDXVersion":
			doc.Version = value
		case "ExternalRef":
			if _, ok := currentObject.(*Package); ok {
				parts := strings.Split(value, " ")
				if len(parts) != 3 {
					return nil, errors.New("malformed external reference")
				}
				currentObject.(*Package).ExternalRefs = append(currentObject.(*Package).ExternalRefs, ExternalRef{
					Category: parts[0],
					Type:     parts[1],
					Locator:  parts[2],
				})
			} else {
				return nil, errors.New("external reference found outside of package")
			}
		case "LicenseListVersion":
			doc.LicenseListVersion = value
		default:
			logrus.Warnf("Unknown tag: %s", tag)
		}
		i++
	}

	// Add the last object from the doc
	currentObject.SetEntity(currentEntity)
	if _, ok := objects[currentObject.SPDXID()]; ok {
		return nil, errors.Errorf("Duplicate SPDXID %s", currentObject.SPDXID())
	}
	objects[currentObject.SPDXID()] = currentObject

	// If somehow the scanner returned an error. Kill it.
	if err := scanner.Err(); err != nil {
		return nil, errors.Wrap(err, "scanned through spdx file, but got an error")
	}

	// Now assign the relationships to the proper objects
	owned := map[string]struct{}{}
	for _, rdata := range rels {
		logrus.Debugf("Procesing %s %s %s", rdata.Source, rdata.Relationship, rdata.Peer)
		// If the source is the doc. Add them
		if rdata.Source == doc.ID {
			if p, ok := objects[rdata.Peer].(*Package); ok {
				logrus.Debugf("doc %s describes package %s", doc.ID, rdata.Peer)
				doc.Packages[rdata.Peer] = p
			}

			if f, ok := objects[rdata.Peer].(*File); ok {
				logrus.Debugf("doc %s describes file %s", doc.ID, rdata.Peer)
				doc.Files[(objects[rdata.Peer]).(*File).SPDXID()] = f
			}
			continue
		}

		// Check if the source object is defined
		if _, ok := objects[rdata.Source]; !ok {
			return nil, errors.Errorf("Unable to find source object with SPDXID %s", rdata.Source)
		}

		// Check that the peer exists
		if _, ok := objects[rdata.Peer]; !ok {
			// ... but only if it is not an external document
			if rdata.ExtDoc == "" {
				return nil, errors.Errorf("Unable to find peer object with SPDXID %s", rdata.Peer)
			}
		}

		if (objects[rdata.Source]).SPDXID() == "" {
			logrus.Fatalf("No ID in object %s:\n%+v", rdata.Source, objects[rdata.Source])
		}
		(objects[rdata.Source]).AddRelationship(&Relationship{
			FullRender:       false,
			PeerReference:    rdata.Peer,
			Type:             RelationshipType(rdata.Relationship),
			Peer:             objects[rdata.Peer],
			PeerExtReference: rdata.ExtDoc,
			// Comment:          "",
		})
		owned[rdata.Peer] = struct{}{}
	}

	// Now, finally any objects not referenced should be made
	// leafs of the document
	for id, obj := range objects {
		if _, ok := owned[id]; !ok {
			if p, ok := obj.(*Package); ok {
				doc.Packages[id] = p
			}

			if f, ok := obj.(*File); ok {
				doc.Files[id] = f
			}
		}
	}

	return doc, nil
}
<commit_msg>spdx: Make regular expressions global to avoid recompiling<commit_after>/*
Copyright 2021 The Kubernetes Authors.

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

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

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

package spdx

import (
	"bufio"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// Regexp to match the tag-value spdx expressions
var (
	tagRegExp          = regexp.MustCompile(`^([a-z0-9A-Z]+):\s+(.+)`)
	relationshioRegExp = regexp.MustCompile(`^*(\S+)\s+([_A-Z]+)\s+(\S+)`)
)

// OpenDoc opens a file, parses a SPDX tag-value file and returns a loaded
// spdx.Document object. This functions has the cyclomatic chec disabled as
// it spans specific cases for each of the tags it recognizes.
// nolint:gocyclo
func OpenDoc(path string) (*Document, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, errors.Wrapf(err, "opening document from %s", path)
	}
	defer file.Close()

	// Create a blank document
	doc := &Document{
		Packages:        map[string]*Package{},
		Files:           map[string]*File{},
		ExternalDocRefs: []ExternalDocumentRef{},
	}

	// Scan the file, looking for tags
	scanner := bufio.NewScanner(file)
	i := 0 // Line counter
	var currentEntity *Entity
	var currentObject Object
	var value, tag, textValue string
	var captureMultiline bool
	objects := map[string]Object{}
	rels := []struct {
		Source       string
		Relationship string
		Peer         string
		ExtDoc       string
	}{}
	for scanner.Scan() {
		// If we are capturing text for a multiline value, read and add
		// the line to the buffer
		if captureMultiline {
			textValue += scanner.Text() + "\n"
			// If we closing tag is not here, continue to the next line
			if !strings.Contains(scanner.Text(), "</text>") {
				continue
			}

			// If closing tag found, remove it from value
			value = strings.ReplaceAll(textValue, "</text>", "")
			textValue = ""
		}

		// Check if line matches if we are not reading multiline values
		if !captureMultiline {
			match := tagRegExp.FindStringSubmatch(scanner.Text())
			if len(match) != 3 {
				continue
			}
			tag = match[1]
			value = match[2]

			// If it is a multiline value, start buffering it
			if strings.HasPrefix(value, "<text>") {
				textValue = strings.ReplaceAll(value, "<text>", "") + "\n"
				captureMultiline = true

				// It may be that the closing tag is right in the same
				// line. If so, capture and finish buffering
				if strings.Contains(scanner.Text(), "</text>") {
					value = strings.ReplaceAll(textValue, "</text>", "")
					textValue = ""
				} else {
					continue
				}
			}
		}

		captureMultiline = false

		switch tag {
		case "FileName", "PackageName":
			// Both FileName or PackageName signal the start of a new entity

			// If we have an object, we store it and continue
			if currentObject != nil {
				currentObject.SetEntity(currentEntity)

				if _, ok := objects[currentObject.SPDXID()]; ok {
					return nil, errors.Errorf("Duplicate SPDXID %s", currentObject.SPDXID())
				}

				objects[currentObject.SPDXID()] = currentObject
			}

			// Create the new entity:
			currentEntity = &Entity{}

			// And the new SPDX object:
			if tag == "FileName" {
				currentObject = &File{}
			}
			if tag == "PackageName" {
				currentObject = &Package{}
			}
			currentEntity.Name = value

		case "SPDXID":
			logrus.Debugf("Entity ID %s", value)
			if currentEntity == nil {
				doc.ID = value
			} else {
				currentEntity.ID = value
			}
		case "PackageLicenseConcluded", "LicenseConcluded":
			if value != NOASSERTION {
				currentEntity.LicenseConcluded = value
			}
		case "PackageCopyrightText", "FileCopyrightText":
			if value != NOASSERTION {
				currentEntity.CopyrightText = value
			}
			// Tags for packages
		case "FilesAnalyzed":
			currentObject.(*Package).FilesAnalyzed = value == "true"
		case "PackageVersion":
			currentObject.(*Package).Version = value
		case "PackageLicenseDeclared":
			currentObject.(*Package).LicenseDeclared = value
		case "PackageVerificationCode":
			currentObject.(*Package).VerificationCode = value
		case "PackageComment":
			currentObject.(*Package).Comment = value
		case "PackageFileName":
			currentObject.(*Package).FileName = value
		case "PackageLicenseInfoFromFiles":
			have := false
			// Check if we already have the license
			for _, licid := range currentObject.(*Package).LicenseInfoFromFiles {
				if licid == value {
					have = true
					break
				}
			}
			if have {
				currentObject.(*Package).LicenseInfoFromFiles = append(currentObject.(*Package).LicenseInfoFromFiles, value)
			}
		case "LicenseInfoInFile":
			if value != NONE {
				currentObject.(*File).LicenseInfoInFile = value
			}
		case "FileChecksum", "PackageChecksum":
			// Checksums are also tag/value -> algo/hash
			match := tagRegExp.FindStringSubmatch(value)
			if len(match) != 3 {
				return nil, errors.Errorf("invalid checksum tag syntax at line %d", i)
			}
			if currentEntity.Checksum == nil {
				currentEntity.Checksum = map[string]string{}
			}
			currentEntity.Checksum[match[1]] = match[2]
		case "Relationship":
			matches := relationshioRegExp.FindStringSubmatch(value)
			if len(matches) != 4 {
				return nil, errors.Errorf("invalid SPDX relationship on line %d: %s", i, value)
			}

			// Check if the relationship is external
			ext := ""
			if strings.HasPrefix(matches[3], "DocumentRef-") && strings.Contains(matches[3], ":") {
				parts := strings.Split(matches[3], ":")
				if len(parts) != 2 {
					return nil, errors.Wrapf(err, "Unable to parse external document reference %s", matches[3])
				}
				matches[3] = parts[0]
				ext = parts[1]
			}

			// Parse the relationship
			rels = append(rels, struct {
				Source       string
				Relationship string
				Peer         string
				ExtDoc       string
			}{
				matches[1], matches[2], matches[3], ext,
			})
		case "PackageDownloadLocation":
			currentEntity.DownloadLocation = value
			// Tags that apply top the doc
		case "Created":
			t, err := time.Parse("2006-01-02T15:04:05Z", value)
			if err != nil {
				return nil, errors.Wrapf(err, "parsing time string in file: %s", value)
			}
			doc.Created = t
		case "Creator":
			// Creator has a tag/value format inside
			match := tagRegExp.FindStringSubmatch(value)
			if len(match) != 3 {
				return nil, errors.Errorf("invalid creator tag syntax at line %d", i)
			}
			switch match[1] {
			case "Person":
				doc.Creator.Person = match[2]
			case "Tool":
				doc.Creator.Tool = append(doc.Creator.Tool, match[2])
			case "Organization":
				doc.Creator.Organization = match[2]
			default:
				return nil, errors.Errorf(
					"invalid creator tag '%s' syntax at line %d, valid values are 'Tool', 'Organization' or 'Person'",
					match[1], i,
				)
			}
		case "DataLicense":
			doc.DataLicense = value
		case "DocumentName":
			doc.Name = value
		case "DocumentNamespace":
			doc.Namespace = value
		case "SPDXVersion":
			doc.Version = value
		case "ExternalRef":
			if _, ok := currentObject.(*Package); ok {
				parts := strings.Split(value, " ")
				if len(parts) != 3 {
					return nil, errors.New("malformed external reference")
				}
				currentObject.(*Package).ExternalRefs = append(currentObject.(*Package).ExternalRefs, ExternalRef{
					Category: parts[0],
					Type:     parts[1],
					Locator:  parts[2],
				})
			} else {
				return nil, errors.New("external reference found outside of package")
			}
		case "LicenseListVersion":
			doc.LicenseListVersion = value
		default:
			logrus.Warnf("Unknown tag: %s", tag)
		}
		i++
	}

	// Add the last object from the doc
	currentObject.SetEntity(currentEntity)
	if _, ok := objects[currentObject.SPDXID()]; ok {
		return nil, errors.Errorf("Duplicate SPDXID %s", currentObject.SPDXID())
	}
	objects[currentObject.SPDXID()] = currentObject

	// If somehow the scanner returned an error. Kill it.
	if err := scanner.Err(); err != nil {
		return nil, errors.Wrap(err, "scanned through spdx file, but got an error")
	}

	// Now assign the relationships to the proper objects
	owned := map[string]struct{}{}
	for _, rdata := range rels {
		logrus.Debugf("Procesing %s %s %s", rdata.Source, rdata.Relationship, rdata.Peer)
		// If the source is the doc. Add them
		if rdata.Source == doc.ID {
			if p, ok := objects[rdata.Peer].(*Package); ok {
				logrus.Debugf("doc %s describes package %s", doc.ID, rdata.Peer)
				doc.Packages[rdata.Peer] = p
			}

			if f, ok := objects[rdata.Peer].(*File); ok {
				logrus.Debugf("doc %s describes file %s", doc.ID, rdata.Peer)
				doc.Files[(objects[rdata.Peer]).(*File).SPDXID()] = f
			}
			continue
		}

		// Check if the source object is defined
		if _, ok := objects[rdata.Source]; !ok {
			return nil, errors.Errorf("Unable to find source object with SPDXID %s", rdata.Source)
		}

		// Check that the peer exists
		if _, ok := objects[rdata.Peer]; !ok {
			// ... but only if it is not an external document
			if rdata.ExtDoc == "" {
				return nil, errors.Errorf("Unable to find peer object with SPDXID %s", rdata.Peer)
			}
		}

		if (objects[rdata.Source]).SPDXID() == "" {
			logrus.Fatalf("No ID in object %s:\n%+v", rdata.Source, objects[rdata.Source])
		}
		(objects[rdata.Source]).AddRelationship(&Relationship{
			FullRender:       false,
			PeerReference:    rdata.Peer,
			Type:             RelationshipType(rdata.Relationship),
			Peer:             objects[rdata.Peer],
			PeerExtReference: rdata.ExtDoc,
			// Comment:          "",
		})
		owned[rdata.Peer] = struct{}{}
	}

	// Now, finally any objects not referenced should be made
	// leafs of the document
	for id, obj := range objects {
		if _, ok := owned[id]; !ok {
			if p, ok := obj.(*Package); ok {
				doc.Packages[id] = p
			}

			if f, ok := obj.(*File); ok {
				doc.Files[id] = f
			}
		}
	}

	return doc, nil
}
<|endoftext|>


<commit_before>package store

import (
	"os"
	"path"
	"strings"
)

type Event struct {
	Type uint
	Seqn uint64
	Path string
	Value string
}

const (
	Set = uint(1<<iota)
	Del
	Add
	Rem
)

var (
	BadPathError = os.NewError("bad path")
)

type Store struct {
	applyCh chan apply
	reqCh chan req
	watchCh chan watch
	watches map[string][]watch
	todo map[uint64]apply
}

type apply struct {
	seqn uint64
	op uint
	k string
	v string
}

type req struct {
	k string
	ch chan reply
}

type reply struct {
	v string
	ok bool
}

type watch struct {
	ch chan Event
	mask uint
	k string
}

func NewStore() *Store {
	s := &Store{
		applyCh: make(chan apply),
		reqCh: make(chan req),
		watchCh: make(chan watch),
		todo: make(map[uint64]apply),
		watches: make(map[string][]watch),
	}
	go s.process()
	return s
}

func checkPath(k string) os.Error {
	switch {
	case len(k) < 1,
	     k[0] != '/',
	     strings.Count(k, "=") > 0,
	     strings.Count(k, " ") > 0:
		return BadPathError
	}
	return nil
}

func EncodeSet(path, v string) (mutation string, err os.Error) {
	if err = checkPath(path); err != nil {
		return
	}
	return path + "=" + v, nil
}

func EncodeDel(path string) (mutation string, err os.Error) {
	if err := checkPath(path); err != nil {
		return
	}
	return path, nil
}

func decode(mutation string) (op uint, path, v string, err os.Error) {
	parts := strings.Split(mutation, "=", 2)
	if err = checkPath(parts[0]); err != nil {
		return
	}
	switch len(parts) {
	case 1:
		return Del, parts[0], "", nil
	case 2:
		return Set, parts[0], parts[1], nil
	}
	panic("can't happen")
}

func (s *Store) notify(ev uint, seqn uint64, k, v string) {
	for _, w := range s.watches[k] {
		if w.mask & ev != 0 {
			w.ch <- Event{ev, seqn, k, v}
		}
	}
}

func append(ws *[]watch, w watch) {
	l := len(*ws)
	if l + 1 > cap(*ws) {
		ns := make([]watch, (l + 1)*2)
		copy(ns, *ws)
		*ws = ns
	}
	*ws = (*ws)[0:l + 1]
	(*ws)[l] = w
}

func (s *Store) process() {
	next := uint64(1)
	values := make(map[string]string)
	for {
		select {
		case a := <-s.applyCh:
			if a.seqn >= next {
				s.todo[a.seqn] = a
			}
			for t, ok := s.todo[next]; ok; t, ok = s.todo[next] {
				switch t.op {
				case Set:
					go s.notify(Set, a.seqn, t.k, t.v)
					if _, ok := values[t.k]; !ok {
						dirname, basename := path.Split(t.k)
						go s.notify(Add, a.seqn, dirname, basename)
					}
					values[t.k] = t.v
				case Del:
					go s.notify(Del, a.seqn, t.k, t.v)
					if _, ok := values[t.k]; ok {
						dirname, basename := path.Split(t.k)
						go s.notify(Rem, a.seqn, dirname, basename)
					}
					values[t.k] = "", false
				}
				s.todo[next] = apply{}, false
				next++
			}
		case r := <-s.reqCh:
			v, ok := values[r.k]
			r.ch <- reply{v, ok}
		case w := <-s.watchCh:
			watches := s.watches[w.k]
			append(&watches, w)
			s.watches[w.k] = watches
		}
	}
}

func (s *Store) Apply(seqn uint64, mutation string) {
	op, path, v, err := decode(mutation)
	if err != nil {
		return
	}
	s.applyCh <- apply{seqn, op, path, v}
}

// For a missing path, `ok == false`. Otherwise, it is `true`.
func (s *Store) Lookup(path string) (v string, ok bool) {
	ch := make(chan reply)
	s.reqCh <- req{path, ch}
	rep := <-ch
	return rep.v, rep.ok
}

// `mask` is one or more of `Set`, `Del`, `Add`, and `Rem`, bitwise OR-ed
// together.
func (s *Store) Watch(path string, mask uint) (events chan Event) {
	ch := make(chan Event)
	s.watchCh <- watch{ch, mask, path}
	return ch
}
<commit_msg>refactor<commit_after>package store

import (
	"os"
	"path"
	"strings"
)

type Event struct {
	Type uint
	Seqn uint64
	Path string
	Value string
}

const (
	Set = uint(1<<iota)
	Del
	Add
	Rem
)

var (
	BadPathError = os.NewError("bad path")
)

type Store struct {
	applyCh chan apply
	reqCh chan req
	watchCh chan watch
	watches map[string][]watch
	todo map[uint64]apply
}

type apply struct {
	seqn uint64
	op uint
	k string
	v string
}

type req struct {
	k string
	ch chan reply
}

type reply struct {
	v string
	ok bool
}

type watch struct {
	ch chan Event
	mask uint
	k string
}

func NewStore() *Store {
	s := &Store{
		applyCh: make(chan apply),
		reqCh: make(chan req),
		watchCh: make(chan watch),
		todo: make(map[uint64]apply),
		watches: make(map[string][]watch),
	}
	go s.process()
	return s
}

func checkPath(k string) os.Error {
	switch {
	case len(k) < 1,
	     k[0] != '/',
	     strings.Count(k, "=") > 0,
	     strings.Count(k, " ") > 0:
		return BadPathError
	}
	return nil
}

func EncodeSet(path, v string) (mutation string, err os.Error) {
	if err = checkPath(path); err != nil {
		return
	}
	return path + "=" + v, nil
}

func EncodeDel(path string) (mutation string, err os.Error) {
	if err := checkPath(path); err != nil {
		return
	}
	return path, nil
}

func decode(mutation string) (op uint, path, v string, err os.Error) {
	parts := strings.Split(mutation, "=", 2)
	if err = checkPath(parts[0]); err != nil {
		return
	}
	switch len(parts) {
	case 1:
		return Del, parts[0], "", nil
	case 2:
		return Set, parts[0], parts[1], nil
	}
	panic("can't happen")
}

func (s *Store) notify(ev uint, seqn uint64, k, v string) {
	for _, w := range s.watches[k] {
		if w.mask & ev != 0 {
			w.ch <- Event{ev, seqn, k, v}
		}
	}
}

func append(ws *[]watch, w watch) {
	l := len(*ws)
	if l + 1 > cap(*ws) {
		ns := make([]watch, (l + 1)*2)
		copy(ns, *ws)
		*ws = ns
	}
	*ws = (*ws)[0:l + 1]
	(*ws)[l] = w
}

func (s *Store) process() {
	next := uint64(1)
	values := make(map[string]string)
	for {
		select {
		case a := <-s.applyCh:
			if a.seqn >= next {
				s.todo[a.seqn] = a
			}
			for t, ok := s.todo[next]; ok; t, ok = s.todo[next] {
				go s.notify(t.op, a.seqn, t.k, t.v)
				switch t.op {
				case Set:
					if _, ok := values[t.k]; !ok {
						dirname, basename := path.Split(t.k)
						go s.notify(Add, a.seqn, dirname, basename)
					}
					values[t.k] = t.v
				case Del:
					if _, ok := values[t.k]; ok {
						dirname, basename := path.Split(t.k)
						go s.notify(Rem, a.seqn, dirname, basename)
					}
					values[t.k] = "", false
				}
				s.todo[next] = apply{}, false
				next++
			}
		case r := <-s.reqCh:
			v, ok := values[r.k]
			r.ch <- reply{v, ok}
		case w := <-s.watchCh:
			watches := s.watches[w.k]
			append(&watches, w)
			s.watches[w.k] = watches
		}
	}
}

func (s *Store) Apply(seqn uint64, mutation string) {
	op, path, v, err := decode(mutation)
	if err != nil {
		return
	}
	s.applyCh <- apply{seqn, op, path, v}
}

// For a missing path, `ok == false`. Otherwise, it is `true`.
func (s *Store) Lookup(path string) (v string, ok bool) {
	ch := make(chan reply)
	s.reqCh <- req{path, ch}
	rep := <-ch
	return rep.v, rep.ok
}

// `mask` is one or more of `Set`, `Del`, `Add`, and `Rem`, bitwise OR-ed
// together.
func (s *Store) Watch(path string, mask uint) (events chan Event) {
	ch := make(chan Event)
	s.watchCh <- watch{ch, mask, path}
	return ch
}
<|endoftext|>


<commit_before>//Copyright 2014 Mahendra Kathirvel. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package simplesession

import (
	"bytes"
	"crypto/rand"
	"crypto/sha1"
	"encoding/gob"
	"encoding/hex"
	"errors"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"
)

const (
	sidLength = 32
	sfileDir  = "/tmp"
)

func New(res http.ResponseWriter, option *Option) (*SimpleSession, error) {
	id, err := generateId()
	if err != nil {
		return nil, err
	}

	fpath := filepath.Join(sfileDir, "gosession_"+id)

	ss := &SimpleSession{
		name:   "GSESSIONID",
		id:     id,
		fpath:  fpath,
		option: option,
		data:   make(map[string]interface{}),
	}

	cookie := &http.Cookie{
		Name:     ss.name,
		Value:    ss.id,
		Path:     ss.option.Path,
		Domain:   ss.option.Domain,
		MaxAge:   ss.option.MaxAge,
		Secure:   ss.option.Secure,
		HttpOnly: ss.option.HttpOnly,
	}
	http.SetCookie(res, cookie)

	return ss, nil
}

type SimpleSession struct {
	name   string
	id     string
	fpath  string
	option *Option
	data   map[string]interface{}
}

func (ss *SimpleSession) Name() string {
	return ss.name
}

func (ss *SimpleSession) Id() string {
	return ss.id
}

func (ss *SimpleSession) FilePath() string {
	return ss.fpath
}

func (ss *SimpleSession) Set(key string, val interface{}) {
	ss.data[key] = val
}

func (ss *SimpleSession) Get(key string) interface{} {
	if val, found := ss.data[key]; found {
		return val
	}
	return nil
}

func (ss *SimpleSession) Del(key string) {
	if _, found := ss.data[key]; found {
		delete(ss.data, key)
	}
}

func (ss *SimpleSession) Write(res http.ResponseWriter, req *http.Request) error {
	serialized, err := serialize(ss.data)
	if err != nil {
		return err
	}

	var fmutex = &sync.RWMutex{}

	fmutex.Lock()
	fp, err := os.OpenFile(ss.fpath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
	if err != nil {
		return errors.New("simplesession: could not open session file")
	}

	if err = os.Chown(ss.fpath, os.Getuid(), os.Getgid()); err != nil {
		return errors.New("simplesession: could not change file owner")
	}
	defer fp.Close()

	if _, err = fp.Write(serialized); err != nil {
		return errors.New("simplesession: could not write session file")
	}
	defer fmutex.Unlock()

	return nil
}

// Read reads stored session
func Read(req *http.Request) (*SimpleSession, error) {
	cke, err := req.Cookie("GSESSIONID")
	if err != nil {
		return nil, errors.New("simplesession: session not set")
	}

	option := &Option{
		Path:     cke.Path,
		Domain:   cke.Domain,
		MaxAge:   cke.MaxAge,
		Secure:   cke.Secure,
		HttpOnly: cke.HttpOnly,
	}

	fpath := filepath.Join(sfileDir, "gosession_"+cke.Value)

	fp, err := os.OpenFile(fpath, os.O_RDONLY, 0400)
	if err != nil {
		return nil, errors.New("simplesession: could not open session file")
	}
	defer fp.Close()

	buf := make([]byte, 128)
	var serialized []byte
	for {
		_, err = fp.Read(buf)
		if err != nil || err == io.EOF {
			break
		}
		serialized = append(serialized, buf[0:]...)
	}

	data := make(map[string]interface{})
	if err = unserialize(serialized, data); err != nil {
		return nil, err
	}

	ss := &SimpleSession{
		name:   cke.Name,
		id:     cke.Value,
		fpath:  fpath,
		option: option,
		data:   data,
	}

	return ss, nil
}

func (ss *SimpleSession) Destroy(res http.ResponseWriter) error {

	cookie := &http.Cookie{
		Name:   ss.name,
		Value:  "",
		Path:   ss.option.Path,
		Domain: ss.option.Domain,
		MaxAge: -1,
	}
	http.SetCookie(res, cookie)

	if err := os.Remove(ss.fpath); err != nil {
		return errors.New("simplesession: no session file found")
	}

	return nil
}

// generateId generates a session id
func generateId() (string, error) {

	hash := sha1.New()
	io.WriteString(hash, strconv.FormatInt(time.Now().Unix(), 10))

	uran := make([]byte, 1024)
	if _, err := io.ReadFull(rand.Reader, uran); err != nil {
		return "", errors.New("simplesession: could not generate random key")
	}

	hash.Write(uran)
	id := hex.EncodeToString(hash.Sum(nil))

	fpath := filepath.Join(sfileDir, "gosession_"+id)
	if _, err := os.Stat(fpath); err == nil {
		generateId()
	}

	return id[0:sidLength], nil
}

// serialize encodes a value using binary.
func serialize(src map[string]interface{}) ([]byte, error) {
	buf := new(bytes.Buffer)
	if err := gob.NewEncoder(buf).Encode(src); err != nil {
		return nil, errors.New("simplesession: error occured while serializing session")
	}
	return buf.Bytes(), nil
}

// unserialize decodes a value using binary.
func unserialize(src []byte, dst map[string]interface{}) error {
	buf := bytes.NewReader(src)
	if err := gob.NewDecoder(buf).Decode(&dst); err != nil {
		return errors.New("simplesession: error occured while unserializing session")
	}
	return nil
}

type Option struct {
	Path     string
	Domain   string
	MaxAge   int
	Secure   bool
	HttpOnly bool
}
<commit_msg>Add more cookie option<commit_after>//Copyright 2014 Mahendra Kathirvel. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package simplesession

import (
	"bytes"
	"crypto/rand"
	"crypto/sha1"
	"encoding/gob"
	"encoding/hex"
	"errors"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"
)

const (
	sidLength = 32
	sfileDir  = "/tmp"
)

func New(res http.ResponseWriter, option *Option) (*SimpleSession, error) {
	id, err := generateId()
	if err != nil {
		return nil, err
	}

	fpath := filepath.Join(sfileDir, "gosession_"+id)

	ss := &SimpleSession{
		name:   "GSESSIONID",
		id:     id,
		fpath:  fpath,
		option: option,
		data:   make(map[string]interface{}),
	}

	cookie := &http.Cookie{
		Name:     ss.name,
		Value:    ss.id,
		Path:     ss.option.Path,
		Domain:   ss.option.Domain,
		MaxAge:   ss.option.MaxAge,
		Secure:   ss.option.Secure,
		HttpOnly: ss.option.HttpOnly,
	}
	http.SetCookie(res, cookie)

	return ss, nil
}

type SimpleSession struct {
	name   string
	id     string
	fpath  string
	option *Option
	data   map[string]interface{}
}

func (ss *SimpleSession) Name() string {
	return ss.name
}

func (ss *SimpleSession) Id() string {
	return ss.id
}

func (ss *SimpleSession) FilePath() string {
	return ss.fpath
}

func (ss *SimpleSession) Set(key string, val interface{}) {
	ss.data[key] = val
}

func (ss *SimpleSession) Get(key string) interface{} {
	if val, ok := ss.data[key]; ok {
		return val
	}
	return nil
}

func (ss *SimpleSession) Del(key string) {
	if _, ok := ss.data[key]; ok {
		delete(ss.data, key)
	}
}

func (ss *SimpleSession) Write(res http.ResponseWriter, req *http.Request) error {
	serialized, err := serialize(ss.data)
	if err != nil {
		return err
	}

	var fmutex = &sync.RWMutex{}

	fmutex.Lock()
	fp, err := os.OpenFile(ss.fpath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
	if err != nil {
		return errors.New("simplesession: could not open session file")
	}

	if err = os.Chown(ss.fpath, os.Getuid(), os.Getgid()); err != nil {
		return errors.New("simplesession: could not change file owner")
	}
	defer fp.Close()

	if _, err = fp.Write(serialized); err != nil {
		return errors.New("simplesession: could not write session file")
	}
	defer fmutex.Unlock()

	return nil
}

// Read reads stored session
func Read(req *http.Request) (*SimpleSession, error) {
	cke, err := req.Cookie("GSESSIONID")
	if err != nil {
		return nil, errors.New("simplesession: session not set")
	}

	option := &Option{
		Path:     cke.Path,
		Domain:   cke.Domain,
		MaxAge:   cke.MaxAge,
		Secure:   cke.Secure,
		HttpOnly: cke.HttpOnly,
	}

	fpath := filepath.Join(sfileDir, "gosession_"+cke.Value)

	fp, err := os.OpenFile(fpath, os.O_RDONLY, 0400)
	if err != nil {
		return nil, errors.New("simplesession: could not open session file")
	}
	defer fp.Close()

	buf := make([]byte, 128)
	var serialized []byte
	for {
		_, err = fp.Read(buf)
		if err != nil || err == io.EOF {
			break
		}
		serialized = append(serialized, buf[0:]...)
	}

	data := make(map[string]interface{})
	if err = unserialize(serialized, data); err != nil {
		return nil, err
	}

	ss := &SimpleSession{
		name:   cke.Name,
		id:     cke.Value,
		fpath:  fpath,
		option: option,
		data:   data,
	}

	return ss, nil
}

func (ss *SimpleSession) Destroy(res http.ResponseWriter) error {

	cookie := &http.Cookie{
		Name:    ss.name,
		Value:   "",
		Path:    ss.option.Path,
		Domain:  ss.option.Domain,
		Expires: time.Unix(1, 0),
		//MaxAge:   -1,
		Secure:   false,
		HttpOnly: true,
	}
	http.SetCookie(res, cookie)

	if err := os.Remove(ss.fpath); err != nil {
		return errors.New("simplesession: no session file found")
	}

	return nil
}

// generateId generates a session id
func generateId() (string, error) {

	hash := sha1.New()
	io.WriteString(hash, strconv.FormatInt(time.Now().Unix(), 10))

	uran := make([]byte, 1024)
	if _, err := io.ReadFull(rand.Reader, uran); err != nil {
		return "", errors.New("simplesession: could not generate random key")
	}

	hash.Write(uran)
	id := hex.EncodeToString(hash.Sum(nil))

	fpath := filepath.Join(sfileDir, "gosession_"+id)
	if _, err := os.Stat(fpath); err == nil {
		generateId()
	}

	return id[0:sidLength], nil
}

// serialize encodes a value using binary.
func serialize(src map[string]interface{}) ([]byte, error) {
	buf := new(bytes.Buffer)
	if err := gob.NewEncoder(buf).Encode(src); err != nil {
		return nil, errors.New("simplesession: error occured while serializing session")
	}
	return buf.Bytes(), nil
}

// unserialize decodes a value using binary.
func unserialize(src []byte, dst map[string]interface{}) error {
	buf := bytes.NewReader(src)
	if err := gob.NewDecoder(buf).Decode(&dst); err != nil {
		return errors.New("simplesession: error occured while unserializing session")
	}
	return nil
}

type Option struct {
	Path     string
	Domain   string
	MaxAge   int
	Secure   bool
	HttpOnly bool
}
<|endoftext|>


<commit_before>// Copyright 2013 Google. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

package plugin

import (
	"bytes"
	"strings"
	"testing"

	"github.com/golang/protobuf/proto"
	descriptor "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/protoc-gen-go/generator"
)

func TestGenerateRPCStubs(t *testing.T) {
	cases := []struct {
		Service *descriptor.ServiceDescriptorProto
		Output  string
	}{
		{
			Service: &descriptor.ServiceDescriptorProto{
				Name: proto.String("math"),
				Method: []*descriptor.MethodDescriptorProto{
					{
						Name:       proto.String("Sqrt"),
						InputType:  proto.String("SqrtInput"),
						OutputType: proto.String("SqrtOutput"),
					},
					{
						Name:       proto.String("Add"),
						InputType:  proto.String("AddInput"),
						OutputType: proto.String("AddOutput"),
					},
				},
			},
			Output: `
// internal wrapper for type-safe RPC calling
type rpcMathClient struct {
	*rpc.Client
}
func (this rpcMathClient) Sqrt(in *SqrtInput, out *SqrtOutput) error {
	return this.Call("Math.Sqrt", in, out)
}
func (this rpcMathClient) Add(in *AddInput, out *AddOutput) error {
	return this.Call("Math.Add", in, out)
}

// NewMathClient returns an *rpc.Client wrapper for calling the methods of
// Math remotely.
func NewMathClient(conn net.Conn) Math {
	return rpcMathClient{rpc.NewClientWithCodec(codec.NewClientCodec(conn))}
}

// ServeMath serves the given Math backend implementation on conn.
func ServeMath(conn net.Conn, backend Math) error {
	srv := rpc.NewServer()
	if err := srv.RegisterName("Math", backend); err != nil {
		return err
	}
	srv.ServeCodec(codec.NewServerCodec(conn))
	return nil
}

// DialMath returns a Math for calling the Math servince at addr (TCP).
func DialMath(addr string) (Math, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}
	return NewMathClient(conn), nil
}

// ListenAndServeMath serves the given Math backend implementation
// on all connections accepted as a result of listening on addr (TCP).
func ListenAndServeMath(addr string, backend Math) error {
	clients, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	srv := rpc.NewServer()
	if err := srv.RegisterName("Math", backend); err != nil {
		return err
	}
	for {
		conn, err := clients.Accept()
		if err != nil {
			return err
		}
		go srv.ServeCodec(codec.NewServerCodec(conn))
	}
	panic("unreachable")
}
`,
		},
	}

	for _, c := range cases {
		buf := new(bytes.Buffer)
		p := Plugin{compileGen: fakeCompileGen{&generator.Generator{Buffer: buf}}}
		p.GenerateRPCStubs(c.Service)
		if got, want := buf.String(), strings.TrimSpace(c.Output)+"\n"; got != want {
			t.Fail()
			t.Logf("GenerateRPCStubs")
			t.Logf("  Input: %s", c.Service)
			t.Logf("  Got:\n%s", got)
			t.Logf("  Want:\n%s", want)
		}
	}
}
<commit_msg>Updated test fixture.<commit_after>// Copyright 2013 Google. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

package plugin

import (
	"bytes"
	"strings"
	"testing"

	"github.com/golang/protobuf/proto"
	descriptor "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/protoc-gen-go/generator"
)

func TestGenerateRPCStubs(t *testing.T) {
	cases := []struct {
		Service *descriptor.ServiceDescriptorProto
		Output  string
	}{
		{
			Service: &descriptor.ServiceDescriptorProto{
				Name: proto.String("math"),
				Method: []*descriptor.MethodDescriptorProto{
					{
						Name:       proto.String("Sqrt"),
						InputType:  proto.String("SqrtInput"),
						OutputType: proto.String("SqrtOutput"),
					},
					{
						Name:       proto.String("Add"),
						InputType:  proto.String("AddInput"),
						OutputType: proto.String("AddOutput"),
					},
				},
			},
			Output: `
// internal wrapper for type-safe RPC calling
type rpcMathClient struct {
	*client.Client
}
func (c rpcMathClient) Sqrt(in *SqrtInput, out *SqrtOutput) error {
	return c.Call("Math.Sqrt", in, out)
}
func (c rpcMathClient) Add(in *AddInput, out *AddOutput) error {
	return c.Call("Math.Add", in, out)
}

// NewMathClient returns an *rpc.Client wrapper for calling the methods of
// Math remotely.
func NewMathClient(addr string) Math {
	return rpcMathClient{client.NewClient(addr)}
}

// ServeMath serves the given Math backend implementation on conn.
func ServeMath(conn net.Conn, backend Math) error {
	srv := rpc.NewServer()
	if err := srv.RegisterName("Math", backend); err != nil {
		return err
	}
	srv.ServeCodec(codec.NewServerCodec(conn))
	return nil
}

// ListenAndServeMath serves the given Math backend implementation
// on all connections accepted as a result of listening on addr (TCP).
func ListenAndServeMath(addr string, backend Math) error {
	clients, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	srv := rpc.NewServer()
	if err := srv.RegisterName("Math", backend); err != nil {
		return err
	}
	for {
		conn, err := clients.Accept()
		if err != nil {
			return err
		}
		go srv.ServeCodec(codec.NewServerCodec(conn))
	}
	panic("unreachable")
}
`,
		},
	}

	for _, c := range cases {
		buf := new(bytes.Buffer)
		p := Plugin{compileGen: fakeCompileGen{&generator.Generator{Buffer: buf}}}
		p.GenerateRPCStubs(c.Service)
		if got, want := buf.String(), strings.TrimSpace(c.Output)+"\n"; got != want {
			t.Fail()
			t.Logf("GenerateRPCStubs")
			t.Logf("  Input: %s", c.Service)
			t.Logf("  Got:\n%s", got)
			t.Logf("  Want:\n%s", want)
		}
	}
}
<|endoftext|>