<commit_before>package dev

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/puma/puma-dev/linebuffer"
	"github.com/puma/puma-dev/watch"
	"github.com/vektra/errors"
	"gopkg.in/tomb.v2"
)

const DefaultThreads = 5

var ErrUnexpectedExit = errors.New("unexpected exit")

type App struct {
	Name    string
	Scheme  string
	Host    string
	Port    int
	Command *exec.Cmd
	Public  bool
	Events  *Events

	lines linebuffer.LineBuffer
	lastLogLine string

	address string
	dir     string

	t tomb.Tomb

	stdout  io.Reader
	pool    *AppPool
	lastUse time.Time

	lock sync.Mutex

	booting bool

	readyChan chan struct{}
}

func (a *App) eventAdd(name string, args ...interface{}) {
	args = append([]interface{}{"app", a.Name}, args...)

	str := a.Events.Add(name, args...)
	a.lines.Append("#event " + str)
}

func (a *App) SetAddress(scheme, host string, port int) {
	a.Scheme = scheme
	a.Host = host
	a.Port = port

	if a.Port == 0 {
		a.address = host
	} else {
		a.address = fmt.Sprintf("%s:%d", a.Host, a.Port)
	}
}

func (a *App) Address() string {
	if a.Port == 0 {
		return a.Host
	}

	return fmt.Sprintf("%s:%d", a.Host, a.Port)
}

func (a *App) Kill(reason string) error {
	a.eventAdd("killing_app",
		"pid", a.Command.Process.Pid,
		"reason", reason,
	)

	fmt.Printf("! Killing '%s' (%d)\n", a.Name, a.Command.Process.Pid)
	err := a.Command.Process.Signal(syscall.SIGTERM)
	if err != nil {
		a.eventAdd("killing_error",
			"pid", a.Command.Process.Pid,
			"error", err.Error(),
		)
		fmt.Printf("! Error trying to kill %s: %s", a.Name, err)
	}
	return err
}

func (a *App) watch() error {
	c := make(chan error)

	go func() {
		r := bufio.NewReader(a.stdout)

		for {
			line, err := r.ReadString('\n')
			if line != "" {
				a.lines.Append(line)
				a.lastLogLine = line
				fmt.Fprintf(os.Stdout, "%s[%d]: %s", a.Name, a.Command.Process.Pid, line)
			}

			if err != nil {
				c <- err
				return
			}
		}
	}()

	var err error

	reason := "detected interval shutdown"

	select {
	case err = <-c:
		reason = "stdout/stderr closed"
		err = fmt.Errorf("%s:\n\t%s", ErrUnexpectedExit, a.lastLogLine)
	case <-a.t.Dying():
		err = nil
	}

	a.Kill(reason)
	a.Command.Wait()
	a.pool.remove(a)

	if a.Scheme == "httpu" {
		os.Remove(a.Address())
	}

	a.eventAdd("shutdown")

	fmt.Printf("* App '%s' shutdown and cleaned up\n", a.Name)

	return err
}

func (a *App) idleMonitor() error {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if a.pool.maybeIdle(a) {
				a.Kill("app is idle")
				return nil
			}
		case <-a.t.Dying():
			return nil
		}
	}

	return nil
}

func (a *App) restartMonitor() error {
	tmpDir := filepath.Join(a.dir, "tmp")
	err := os.MkdirAll(tmpDir, 0755)
	if err != nil {
		return err
	}

	restart := filepath.Join(tmpDir, "restart.txt")

	f, err := os.Create(restart)
	if err != nil {
		return err
	}
	f.Close()

	return watch.Watch(restart, a.t.Dying(), func() {
		a.Kill("restart.txt touched")
	})
}

func (a *App) WaitTilReady() error {
	select {
	case <-a.readyChan:
		// double check we aren't also dying
		select {
		case <-a.t.Dying():
			return a.t.Err()
		default:
			a.lastUse = time.Now()
			return nil
		}
	case <-a.t.Dying():
		return a.t.Err()
	}
}

const (
	Booting = iota
	Running
	Dead
)

func (a *App) Status() int {
	// These are done in order as separate selects because go's
	// select does not execute case's sequentially, it runs bodies
	// after sampling all channels and picking a random body.
	select {
	case <-a.t.Dying():
		return Dead
	default:
		select {
		case <-a.readyChan:
			return Running
		default:
			return Dead
		}
	}
}

func (a *App) Log() string {
	var buf bytes.Buffer
	a.lines.WriteTo(&buf)
	return buf.String()
}

const executionShell = `exec bash -c '
cd %s

if test -e ~/.powconfig; then
	source ~/.powconfig
fi

if test -e .env; then
	source .env
fi

if test -e .powrc; then
	source .powrc
fi

if test -e .powenv; then
	source .powenv
fi

if test -e Gemfile && bundle exec puma -V &>/dev/null; then
	exec bundle exec puma -C $CONFIG --tag puma-dev:%s -w $WORKERS -t 0:$THREADS -b unix:%s
fi

exec puma -C $CONFIG --tag puma-dev:%s -w $WORKERS -t 0:$THREADS -b unix:%s'
`

func (pool *AppPool) LaunchApp(name, dir string) (*App, error) {
	tmpDir := filepath.Join(dir, "tmp")
	err := os.MkdirAll(tmpDir, 0755)
	if err != nil {
		return nil, err
	}

	socket := filepath.Join(tmpDir, fmt.Sprintf("puma-dev-%d.sock", os.Getpid()))

	shell := os.Getenv("SHELL")

	if shell == "" {
		log.Printf("[WARN] SHELL env var not set, using /bin/bash by default")
		shell = "/bin/bash"
	}

	cmd := exec.Command(shell, "-l", "-i", "-c",
		fmt.Sprintf(executionShell, dir, name, socket, name, socket))

	cmd.Dir = dir

	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env,
		fmt.Sprintf("THREADS=%d", DefaultThreads),
		"WORKERS=0",
		"CONFIG=-",
	)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	cmd.Stderr = cmd.Stdout

	err = cmd.Start()
	if err != nil {
		return nil, errors.Context(err, "starting app")
	}

	fmt.Printf("! Booting app '%s' on socket %s\n", name, socket)

	app := &App{
		Name:      name,
		Command:   cmd,
		Events:    pool.Events,
		stdout:    stdout,
		dir:       dir,
		pool:      pool,
		readyChan: make(chan struct{}),
		lastUse:   time.Now(),
	}

	app.eventAdd("booting_app", "socket", socket)

	stat, err := os.Stat(filepath.Join(dir, "public"))
	if err == nil {
		app.Public = stat.IsDir()
	}

	app.SetAddress("httpu", socket, 0)

	app.t.Go(app.watch)
	app.t.Go(app.idleMonitor)
	app.t.Go(app.restartMonitor)

	app.t.Go(func() error {
		// This is a poor substitute for getting an actual readiness signal
		// from puma but it's good enough.

		app.eventAdd("waiting_on_app")

		ticker := time.NewTicker(250 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-app.t.Dying():
				app.eventAdd("dying_on_start")
				fmt.Printf("! Detecting app '%s' dying on start\n", name)
				return fmt.Errorf("app died before booting")
			case <-ticker.C:
				c, err := net.Dial("unix", socket)
				if err == nil {
					c.Close()
					app.eventAdd("app_ready")
					fmt.Printf("! App '%s' booted\n", name)
					close(app.readyChan)
					return nil
				}
			}
		}
	})

	return app, nil
}

func (pool *AppPool) readProxy(name, path string) (*App, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	app := &App{
		Name:      name,
		Events:    pool.Events,
		pool:      pool,
		readyChan: make(chan struct{}),
		lastUse:   time.Now(),
	}

	data = bytes.TrimSpace(data)

	port, err := strconv.Atoi(string(data))
	if err == nil {
		app.SetAddress("http", "127.0.0.1", port)
	} else {
		u, err := url.Parse(string(data))
		if err != nil {
			return nil, err
		}

		var (
			sport, host string
			port        int
		)

		host, sport, err = net.SplitHostPort(u.Host)
		if err == nil {
			port, err = strconv.Atoi(sport)
			if err != nil {
				return nil, err
			}
		} else {
			host = u.Host
		}

		app.SetAddress(u.Scheme, host, port)
	}

	app.eventAdd("proxy_created",
		"destination", fmt.Sprintf("%s://%s"), app.Scheme, app.Address())

	fmt.Printf("* Generated proxy connection for '%s' to %s://%s\n",
		name, app.Scheme, app.Address())

	// to satisfy the tomb
	app.t.Go(func() error {
		<-app.t.Dying()
		return nil
	})

	close(app.readyChan)

	return app, nil
}

type AppPool struct {
	Dir      string
	IdleTime time.Duration
	Debug    bool
	Events   *Events

	AppClosed func(*App)

	lock sync.Mutex
	apps map[string]*App
}

func (a *AppPool) maybeIdle(app *App) bool {
	a.lock.Lock()
	defer a.lock.Unlock()

	diff := time.Since(app.lastUse)
	if diff > a.IdleTime {
		app.eventAdd("idle_app", "last_used", diff.String())
		delete(a.apps, app.Name)
		return true
	}

	return false
}

var ErrUnknownApp = errors.New("unknown app")

func (a *AppPool) App(name string) (*App, error) {
	a.lock.Lock()
	defer a.lock.Unlock()

	if a.apps == nil {
		a.apps = make(map[string]*App)
	}

	app, ok := a.apps[name]
	if ok {
		return app, nil
	}

	path := filepath.Join(a.Dir, name)

	a.Events.Add("app_lookup", "path", path)

	stat, err := os.Stat(path)
	destPath, _ := os.Readlink(path)

	if err != nil {
		if os.IsNotExist(err) {
			// Check there might be a link there but it's not valid
			_, err := os.Lstat(path)
			if err == nil {
				fmt.Printf("! Bad symlink detected '%s'. Destination '%s' doesn't exist\n", path, destPath)
				a.Events.Add("bad_symlink", "path", path, "dest", destPath)
			}

			return nil, ErrUnknownApp
		}

		return nil, err
	}

	canonicalName := name
	aliasName := ""

	// Handle multiple symlinks to the same app
	destStat, err := os.Stat(destPath)
	if err == nil {
		destName := destStat.Name()
		if destName != canonicalName {
			canonicalName = destName
			aliasName = name
		}
	}

	app, ok = a.apps[canonicalName]

	if !ok {
		if stat.IsDir() {
			app, err = a.LaunchApp(canonicalName, path)
		} else {
			app, err = a.readProxy(canonicalName, path)
		}
	}

	if err != nil {
		a.Events.Add("error_starting_app", "app", canonicalName, "error", err.Error())
		return nil, err
	}

	a.apps[canonicalName] = app

	if aliasName != "" {
		a.apps[aliasName] = app
	}

	return app, nil
}

func (a *AppPool) remove(app *App) {
	a.lock.Lock()
	defer a.lock.Unlock()

	// Find all instance references so aliases are removed too
	for name, candidate := range a.apps {
		if candidate == app {
			delete(a.apps, name)
		}
	}

	if a.AppClosed != nil {
		a.AppClosed(app)
	}
}

func (a *AppPool) ForApps(f func(*App)) {
	a.lock.Lock()
	defer a.lock.Unlock()

	for _, app := range a.apps {
		f(app)
	}
}

func (a *AppPool) Purge() {
	a.lock.Lock()

	var apps []*App

	for _, app := range a.apps {
		apps = append(apps, app)
	}

	a.lock.Unlock()

	for _, app := range apps {
		app.eventAdd("purging_app")
		app.t.Kill(nil)
	}

	for _, app := range apps {
		app.t.Wait()
	}

	a.Events.Add("apps_purged")
}
<commit_msg>gofmt<commit_after>package dev

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/puma/puma-dev/linebuffer"
	"github.com/puma/puma-dev/watch"
	"github.com/vektra/errors"
	"gopkg.in/tomb.v2"
)

const DefaultThreads = 5

var ErrUnexpectedExit = errors.New("unexpected exit")

type App struct {
	Name    string
	Scheme  string
	Host    string
	Port    int
	Command *exec.Cmd
	Public  bool
	Events  *Events

	lines       linebuffer.LineBuffer
	lastLogLine string

	address string
	dir     string

	t tomb.Tomb

	stdout  io.Reader
	pool    *AppPool
	lastUse time.Time

	lock sync.Mutex

	booting bool

	readyChan chan struct{}
}

func (a *App) eventAdd(name string, args ...interface{}) {
	args = append([]interface{}{"app", a.Name}, args...)

	str := a.Events.Add(name, args...)
	a.lines.Append("#event " + str)
}

func (a *App) SetAddress(scheme, host string, port int) {
	a.Scheme = scheme
	a.Host = host
	a.Port = port

	if a.Port == 0 {
		a.address = host
	} else {
		a.address = fmt.Sprintf("%s:%d", a.Host, a.Port)
	}
}

func (a *App) Address() string {
	if a.Port == 0 {
		return a.Host
	}

	return fmt.Sprintf("%s:%d", a.Host, a.Port)
}

func (a *App) Kill(reason string) error {
	a.eventAdd("killing_app",
		"pid", a.Command.Process.Pid,
		"reason", reason,
	)

	fmt.Printf("! Killing '%s' (%d)\n", a.Name, a.Command.Process.Pid)
	err := a.Command.Process.Signal(syscall.SIGTERM)
	if err != nil {
		a.eventAdd("killing_error",
			"pid", a.Command.Process.Pid,
			"error", err.Error(),
		)
		fmt.Printf("! Error trying to kill %s: %s", a.Name, err)
	}
	return err
}

func (a *App) watch() error {
	c := make(chan error)

	go func() {
		r := bufio.NewReader(a.stdout)

		for {
			line, err := r.ReadString('\n')
			if line != "" {
				a.lines.Append(line)
				a.lastLogLine = line
				fmt.Fprintf(os.Stdout, "%s[%d]: %s", a.Name, a.Command.Process.Pid, line)
			}

			if err != nil {
				c <- err
				return
			}
		}
	}()

	var err error

	reason := "detected interval shutdown"

	select {
	case err = <-c:
		reason = "stdout/stderr closed"
		err = fmt.Errorf("%s:\n\t%s", ErrUnexpectedExit, a.lastLogLine)
	case <-a.t.Dying():
		err = nil
	}

	a.Kill(reason)
	a.Command.Wait()
	a.pool.remove(a)

	if a.Scheme == "httpu" {
		os.Remove(a.Address())
	}

	a.eventAdd("shutdown")

	fmt.Printf("* App '%s' shutdown and cleaned up\n", a.Name)

	return err
}

func (a *App) idleMonitor() error {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if a.pool.maybeIdle(a) {
				a.Kill("app is idle")
				return nil
			}
		case <-a.t.Dying():
			return nil
		}
	}

	return nil
}

func (a *App) restartMonitor() error {
	tmpDir := filepath.Join(a.dir, "tmp")
	err := os.MkdirAll(tmpDir, 0755)
	if err != nil {
		return err
	}

	restart := filepath.Join(tmpDir, "restart.txt")

	f, err := os.Create(restart)
	if err != nil {
		return err
	}
	f.Close()

	return watch.Watch(restart, a.t.Dying(), func() {
		a.Kill("restart.txt touched")
	})
}

func (a *App) WaitTilReady() error {
	select {
	case <-a.readyChan:
		// double check we aren't also dying
		select {
		case <-a.t.Dying():
			return a.t.Err()
		default:
			a.lastUse = time.Now()
			return nil
		}
	case <-a.t.Dying():
		return a.t.Err()
	}
}

const (
	Booting = iota
	Running
	Dead
)

func (a *App) Status() int {
	// These are done in order as separate selects because go's
	// select does not execute case's sequentially, it runs bodies
	// after sampling all channels and picking a random body.
	select {
	case <-a.t.Dying():
		return Dead
	default:
		select {
		case <-a.readyChan:
			return Running
		default:
			return Dead
		}
	}
}

func (a *App) Log() string {
	var buf bytes.Buffer
	a.lines.WriteTo(&buf)
	return buf.String()
}

const executionShell = `exec bash -c '
cd %s

if test -e ~/.powconfig; then
	source ~/.powconfig
fi

if test -e .env; then
	source .env
fi

if test -e .powrc; then
	source .powrc
fi

if test -e .powenv; then
	source .powenv
fi

if test -e Gemfile && bundle exec puma -V &>/dev/null; then
	exec bundle exec puma -C $CONFIG --tag puma-dev:%s -w $WORKERS -t 0:$THREADS -b unix:%s
fi

exec puma -C $CONFIG --tag puma-dev:%s -w $WORKERS -t 0:$THREADS -b unix:%s'
`

func (pool *AppPool) LaunchApp(name, dir string) (*App, error) {
	tmpDir := filepath.Join(dir, "tmp")
	err := os.MkdirAll(tmpDir, 0755)
	if err != nil {
		return nil, err
	}

	socket := filepath.Join(tmpDir, fmt.Sprintf("puma-dev-%d.sock", os.Getpid()))

	shell := os.Getenv("SHELL")

	if shell == "" {
		log.Printf("[WARN] SHELL env var not set, using /bin/bash by default")
		shell = "/bin/bash"
	}

	cmd := exec.Command(shell, "-l", "-i", "-c",
		fmt.Sprintf(executionShell, dir, name, socket, name, socket))

	cmd.Dir = dir

	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env,
		fmt.Sprintf("THREADS=%d", DefaultThreads),
		"WORKERS=0",
		"CONFIG=-",
	)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	cmd.Stderr = cmd.Stdout

	err = cmd.Start()
	if err != nil {
		return nil, errors.Context(err, "starting app")
	}

	fmt.Printf("! Booting app '%s' on socket %s\n", name, socket)

	app := &App{
		Name:      name,
		Command:   cmd,
		Events:    pool.Events,
		stdout:    stdout,
		dir:       dir,
		pool:      pool,
		readyChan: make(chan struct{}),
		lastUse:   time.Now(),
	}

	app.eventAdd("booting_app", "socket", socket)

	stat, err := os.Stat(filepath.Join(dir, "public"))
	if err == nil {
		app.Public = stat.IsDir()
	}

	app.SetAddress("httpu", socket, 0)

	app.t.Go(app.watch)
	app.t.Go(app.idleMonitor)
	app.t.Go(app.restartMonitor)

	app.t.Go(func() error {
		// This is a poor substitute for getting an actual readiness signal
		// from puma but it's good enough.

		app.eventAdd("waiting_on_app")

		ticker := time.NewTicker(250 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-app.t.Dying():
				app.eventAdd("dying_on_start")
				fmt.Printf("! Detecting app '%s' dying on start\n", name)
				return fmt.Errorf("app died before booting")
			case <-ticker.C:
				c, err := net.Dial("unix", socket)
				if err == nil {
					c.Close()
					app.eventAdd("app_ready")
					fmt.Printf("! App '%s' booted\n", name)
					close(app.readyChan)
					return nil
				}
			}
		}
	})

	return app, nil
}

func (pool *AppPool) readProxy(name, path string) (*App, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	app := &App{
		Name:      name,
		Events:    pool.Events,
		pool:      pool,
		readyChan: make(chan struct{}),
		lastUse:   time.Now(),
	}

	data = bytes.TrimSpace(data)

	port, err := strconv.Atoi(string(data))
	if err == nil {
		app.SetAddress("http", "127.0.0.1", port)
	} else {
		u, err := url.Parse(string(data))
		if err != nil {
			return nil, err
		}

		var (
			sport, host string
			port        int
		)

		host, sport, err = net.SplitHostPort(u.Host)
		if err == nil {
			port, err = strconv.Atoi(sport)
			if err != nil {
				return nil, err
			}
		} else {
			host = u.Host
		}

		app.SetAddress(u.Scheme, host, port)
	}

	app.eventAdd("proxy_created",
		"destination", fmt.Sprintf("%s://%s"), app.Scheme, app.Address())

	fmt.Printf("* Generated proxy connection for '%s' to %s://%s\n",
		name, app.Scheme, app.Address())

	// to satisfy the tomb
	app.t.Go(func() error {
		<-app.t.Dying()
		return nil
	})

	close(app.readyChan)

	return app, nil
}

type AppPool struct {
	Dir      string
	IdleTime time.Duration
	Debug    bool
	Events   *Events

	AppClosed func(*App)

	lock sync.Mutex
	apps map[string]*App
}

func (a *AppPool) maybeIdle(app *App) bool {
	a.lock.Lock()
	defer a.lock.Unlock()

	diff := time.Since(app.lastUse)
	if diff > a.IdleTime {
		app.eventAdd("idle_app", "last_used", diff.String())
		delete(a.apps, app.Name)
		return true
	}

	return false
}

var ErrUnknownApp = errors.New("unknown app")

func (a *AppPool) App(name string) (*App, error) {
	a.lock.Lock()
	defer a.lock.Unlock()

	if a.apps == nil {
		a.apps = make(map[string]*App)
	}

	app, ok := a.apps[name]
	if ok {
		return app, nil
	}

	path := filepath.Join(a.Dir, name)

	a.Events.Add("app_lookup", "path", path)

	stat, err := os.Stat(path)
	destPath, _ := os.Readlink(path)

	if err != nil {
		if os.IsNotExist(err) {
			// Check there might be a link there but it's not valid
			_, err := os.Lstat(path)
			if err == nil {
				fmt.Printf("! Bad symlink detected '%s'. Destination '%s' doesn't exist\n", path, destPath)
				a.Events.Add("bad_symlink", "path", path, "dest", destPath)
			}

			return nil, ErrUnknownApp
		}

		return nil, err
	}

	canonicalName := name
	aliasName := ""

	// Handle multiple symlinks to the same app
	destStat, err := os.Stat(destPath)
	if err == nil {
		destName := destStat.Name()
		if destName != canonicalName {
			canonicalName = destName
			aliasName = name
		}
	}

	app, ok = a.apps[canonicalName]

	if !ok {
		if stat.IsDir() {
			app, err = a.LaunchApp(canonicalName, path)
		} else {
			app, err = a.readProxy(canonicalName, path)
		}
	}

	if err != nil {
		a.Events.Add("error_starting_app", "app", canonicalName, "error", err.Error())
		return nil, err
	}

	a.apps[canonicalName] = app

	if aliasName != "" {
		a.apps[aliasName] = app
	}

	return app, nil
}

func (a *AppPool) remove(app *App) {
	a.lock.Lock()
	defer a.lock.Unlock()

	// Find all instance references so aliases are removed too
	for name, candidate := range a.apps {
		if candidate == app {
			delete(a.apps, name)
		}
	}

	if a.AppClosed != nil {
		a.AppClosed(app)
	}
}

func (a *AppPool) ForApps(f func(*App)) {
	a.lock.Lock()
	defer a.lock.Unlock()

	for _, app := range a.apps {
		f(app)
	}
}

func (a *AppPool) Purge() {
	a.lock.Lock()

	var apps []*App

	for _, app := range a.apps {
		apps = append(apps, app)
	}

	a.lock.Unlock()

	for _, app := range apps {
		app.eventAdd("purging_app")
		app.t.Kill(nil)
	}

	for _, app := range apps {
		app.t.Wait()
	}

	a.Events.Add("apps_purged")
}
<|endoftext|>


<commit_before>package endpoints

import (
	"fmt"
	"strings"

	workloadutil "github.com/rancher/rancher/pkg/controllers/user/workload"
	"github.com/rancher/types/apis/core/v1"
	"github.com/rancher/types/apis/extensions/v1beta1"
	managementv3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/apis/project.cattle.io/v3"
	"github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	extensionsv1beta1 "k8s.io/api/extensions/v1beta1"
	"k8s.io/apimachinery/pkg/labels"
)

// This controller is responsible for monitoring workloads
// and setting public endpoints on them based on HostPort pods
// and NodePort services backing up the workload

type WorkloadEndpointsController struct {
	ingressLister      v1beta1.IngressLister
	serviceLister      v1.ServiceLister
	podLister          v1.PodLister
	WorkloadController workloadutil.CommonController
	machinesLister     managementv3.NodeLister
	clusterName        string
	isRKE              bool
}

func (c *WorkloadEndpointsController) UpdateEndpoints(key string, obj *workloadutil.Workload) error {
	if obj == nil && key != workloadutil.AllWorkloads {
		return nil
	}
	// do not update endpoints for job, cronJob and for workload owned by controller (ReplicaSet)
	if strings.EqualFold(obj.Kind, "job") || strings.EqualFold(obj.Kind, "cronJob") {
		return nil
	}
	for _, o := range obj.OwnerReferences {
		if *o.Controller {
			return nil
		}
	}

	var workloads []*workloadutil.Workload
	var services []*corev1.Service
	var ingresses []*extensionsv1beta1.Ingress
	var err error
	if strings.HasSuffix(key, workloadutil.AllWorkloads) {
		namespace := ""
		if !strings.EqualFold(key, allEndpoints) {
			namespace = strings.TrimSuffix(key, fmt.Sprintf("/%s", allEndpoints))
		}
		workloads, err = c.WorkloadController.GetAllWorkloads(namespace)
		if err != nil {
			return err
		}
		services, err = c.serviceLister.List(namespace, labels.NewSelector())
		if err != nil {
			return err
		}
		ingresses, err = c.ingressLister.List(namespace, labels.NewSelector())

	} else {
		ingresses, err = c.ingressLister.List(obj.Namespace, labels.NewSelector())
		if err != nil {
			return err
		}
		services, err = c.serviceLister.List(obj.Namespace, labels.NewSelector())
		workloads = append(workloads, obj)
	}
	if err != nil {
		return err
	}
	nodeNameToMachine, err := getNodeNameToMachine(c.clusterName, c.machinesLister)
	if err != nil {
		return err
	}
	// get ingress endpoint group by service
	serviceToIngressEndpoints := make(map[string][]v3.PublicEndpoint)
	for _, ingress := range ingresses {
		epsMap, err := convertIngressToServicePublicEndpointsMap(ingress, c.isRKE)
		if err != nil {
			return err
		}
		for k, v := range epsMap {
			eps := serviceToIngressEndpoints[k]
			serviceToIngressEndpoints[k] = append(eps, v...)
		}
	}

	for _, w := range workloads {
		// 1. Get endpoints from services
		var newPublicEps []v3.PublicEndpoint
		for _, svc := range services {
			set := labels.Set{}
			for key, val := range svc.Spec.Selector {
				set[key] = val
			}
			selector := labels.SelectorFromSet(set)
			found := false
			if selector.Matches(labels.Set(w.Labels)) {
				// direct selector match
				found = true
			} else {
				// match based off the workload
				value, ok := svc.Annotations[workloadutil.WorkloadAnnotation]
				if !ok || value == "" {
					continue
				}
				workloadIDs := strings.Split(value, ",")
				for _, workloadID := range workloadIDs {
					splitted := strings.Split(workloadID, ":")
					if len(splitted) != 3 {
						continue
					}
					namespace := splitted[1]
					name := splitted[2]
					if w.Name == name && w.Namespace == namespace {
						found = true
						break
					}
				}
			}
			if found {
				eps, err := convertServiceToPublicEndpoints(svc, "", nil)
				if err != nil {
					return err
				}
				newPublicEps = append(newPublicEps, eps...)
				if ingressEndpoints, ok := serviceToIngressEndpoints[svc.Name]; ok {
					newPublicEps = append(newPublicEps, ingressEndpoints...)
				}
			}
		}

		// 2. Get endpoints from HostPort pods matching the selector
		set := labels.Set{}
		for key, val := range w.SelectorLabels {
			set[key] = val
		}
		pods, err := c.podLister.List(w.Namespace, labels.SelectorFromSet(set))
		if err != nil {
			return err
		}
		for _, pod := range pods {
			eps, err := convertHostPortToEndpoint(pod, c.clusterName, nodeNameToMachine[pod.Spec.NodeName])
			if err != nil {
				return err
			}
			newPublicEps = append(newPublicEps, eps...)
		}

		existingPublicEps := getPublicEndpointsFromAnnotations(w.Annotations)
		if areEqualEndpoints(existingPublicEps, newPublicEps) {
			return nil
		}
		epsToUpdate, err := publicEndpointsToString(newPublicEps)
		if err != nil {
			return err
		}

		logrus.Infof("Updating workload [%s] with public endpoints [%v]", key, epsToUpdate)

		annotations := map[string]string{
			endpointsAnnotation: epsToUpdate,
		}
		if err = c.WorkloadController.UpdateWorkload(w, annotations); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Decode workloadIDs by json format<commit_after>package endpoints

import (
	"encoding/json"
	"fmt"
	"strings"

	workloadutil "github.com/rancher/rancher/pkg/controllers/user/workload"
	"github.com/rancher/types/apis/core/v1"
	"github.com/rancher/types/apis/extensions/v1beta1"
	managementv3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/apis/project.cattle.io/v3"
	"github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	extensionsv1beta1 "k8s.io/api/extensions/v1beta1"
	"k8s.io/apimachinery/pkg/labels"
)

// This controller is responsible for monitoring workloads
// and setting public endpoints on them based on HostPort pods
// and NodePort services backing up the workload

type WorkloadEndpointsController struct {
	ingressLister      v1beta1.IngressLister
	serviceLister      v1.ServiceLister
	podLister          v1.PodLister
	WorkloadController workloadutil.CommonController
	machinesLister     managementv3.NodeLister
	clusterName        string
	isRKE              bool
}

func (c *WorkloadEndpointsController) UpdateEndpoints(key string, obj *workloadutil.Workload) error {
	if obj == nil && key != workloadutil.AllWorkloads {
		return nil
	}
	// do not update endpoints for job, cronJob and for workload owned by controller (ReplicaSet)
	if strings.EqualFold(obj.Kind, "job") || strings.EqualFold(obj.Kind, "cronJob") {
		return nil
	}
	for _, o := range obj.OwnerReferences {
		if *o.Controller {
			return nil
		}
	}

	var workloads []*workloadutil.Workload
	var services []*corev1.Service
	var ingresses []*extensionsv1beta1.Ingress
	var err error
	if strings.HasSuffix(key, workloadutil.AllWorkloads) {
		namespace := ""
		if !strings.EqualFold(key, allEndpoints) {
			namespace = strings.TrimSuffix(key, fmt.Sprintf("/%s", allEndpoints))
		}
		workloads, err = c.WorkloadController.GetAllWorkloads(namespace)
		if err != nil {
			return err
		}
		services, err = c.serviceLister.List(namespace, labels.NewSelector())
		if err != nil {
			return err
		}
		ingresses, err = c.ingressLister.List(namespace, labels.NewSelector())

	} else {
		ingresses, err = c.ingressLister.List(obj.Namespace, labels.NewSelector())
		if err != nil {
			return err
		}
		services, err = c.serviceLister.List(obj.Namespace, labels.NewSelector())
		workloads = append(workloads, obj)
	}
	if err != nil {
		return err
	}
	nodeNameToMachine, err := getNodeNameToMachine(c.clusterName, c.machinesLister)
	if err != nil {
		return err
	}
	// get ingress endpoint group by service
	serviceToIngressEndpoints := make(map[string][]v3.PublicEndpoint)
	for _, ingress := range ingresses {
		epsMap, err := convertIngressToServicePublicEndpointsMap(ingress, c.isRKE)
		if err != nil {
			return err
		}
		for k, v := range epsMap {
			eps := serviceToIngressEndpoints[k]
			serviceToIngressEndpoints[k] = append(eps, v...)
		}
	}

	for _, w := range workloads {
		// 1. Get endpoints from services
		var newPublicEps []v3.PublicEndpoint
		for _, svc := range services {
			set := labels.Set{}
			for key, val := range svc.Spec.Selector {
				set[key] = val
			}
			selector := labels.SelectorFromSet(set)
			found := false
			if selector.Matches(labels.Set(w.Labels)) {
				// direct selector match
				found = true
			} else {
				// match based off the workload
				value, ok := svc.Annotations[workloadutil.WorkloadAnnotation]
				if !ok || value == "" {
					continue
				}
				var workloadIDs []string
				err := json.Unmarshal([]byte(value), &workloadIDs)
				if err != nil {
					logrus.WithError(err).Errorf("Unmarshalling %s workloadIDs of %s Error", value, svc.Name)
					continue
				}
				for _, workloadID := range workloadIDs {
					splitted := strings.Split(workloadID, ":")
					if len(splitted) != 3 {
						continue
					}
					namespace := splitted[1]
					name := splitted[2]
					if w.Name == name && w.Namespace == namespace {
						found = true
						break
					}
				}
			}
			if found {
				eps, err := convertServiceToPublicEndpoints(svc, "", nil)
				if err != nil {
					return err
				}
				newPublicEps = append(newPublicEps, eps...)
				if ingressEndpoints, ok := serviceToIngressEndpoints[svc.Name]; ok {
					newPublicEps = append(newPublicEps, ingressEndpoints...)
				}
			}
		}

		// 2. Get endpoints from HostPort pods matching the selector
		set := labels.Set{}
		for key, val := range w.SelectorLabels {
			set[key] = val
		}
		pods, err := c.podLister.List(w.Namespace, labels.SelectorFromSet(set))
		if err != nil {
			return err
		}
		for _, pod := range pods {
			eps, err := convertHostPortToEndpoint(pod, c.clusterName, nodeNameToMachine[pod.Spec.NodeName])
			if err != nil {
				return err
			}
			newPublicEps = append(newPublicEps, eps...)
		}

		existingPublicEps := getPublicEndpointsFromAnnotations(w.Annotations)
		if areEqualEndpoints(existingPublicEps, newPublicEps) {
			return nil
		}
		epsToUpdate, err := publicEndpointsToString(newPublicEps)
		if err != nil {
			return err
		}

		logrus.Infof("Updating workload [%s] with public endpoints [%v]", key, epsToUpdate)

		annotations := map[string]string{
			endpointsAnnotation: epsToUpdate,
		}
		if err = c.WorkloadController.UpdateWorkload(w, annotations); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package dht

import (
	"context"
	"fmt"
	"sync"
	"time"

	ggio "github.com/gogo/protobuf/io"
	ctxio "github.com/jbenet/go-context/io"
	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
)

var dhtReadMessageTimeout = time.Minute
var ErrReadTimeout = fmt.Errorf("timed out reading response")

// handleNewStream implements the inet.StreamHandler
func (dht *IpfsDHT) handleNewStream(s inet.Stream) {
	go dht.handleNewMessage(s)
}

func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
	defer s.Close()

	ctx := dht.Context()
	cr := ctxio.NewReader(ctx, s) // ok to use. we defer close stream in this func
	cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func
	r := ggio.NewDelimitedReader(cr, inet.MessageSizeMax)
	w := ggio.NewDelimitedWriter(cw)
	mPeer := s.Conn().RemotePeer()

	for {
		// receive msg
		pmes := new(pb.Message)
		if err := r.ReadMsg(pmes); err != nil {
			s.Reset()
			log.Debugf("Error unmarshaling data: %s", err)
			return
		}

		// update the peer (on valid msgs only)
		dht.updateFromMessage(ctx, mPeer, pmes)

		// get handler for this msg type.
		handler := dht.handlerForMsgType(pmes.GetType())
		if handler == nil {
			s.Reset()
			log.Debug("got back nil handler from handlerForMsgType")
			return
		}

		// dispatch handler.
		rpmes, err := handler(ctx, mPeer, pmes)
		if err != nil {
			s.Reset()
			log.Debugf("handle message error: %s", err)
			return
		}

		// if nil response, return it before serializing
		if rpmes == nil {
			log.Debug("got back nil response from request")
			continue
		}

		// send out response msg
		if err := w.WriteMsg(rpmes); err != nil {
			s.Reset()
			log.Debugf("send response error: %s", err)
			return
		}
	}
}

// sendRequest sends out a request, but also makes sure to
// measure the RTT for latency measurements.
func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {

	ms, err := dht.messageSenderForPeer(p)
	if err != nil {
		return nil, err
	}

	start := time.Now()

	rpmes, err := ms.SendRequest(ctx, pmes)
	if err != nil {
		return nil, err
	}

	// update the peer (on valid msgs only)
	dht.updateFromMessage(ctx, p, rpmes)

	dht.peerstore.RecordLatency(p, time.Since(start))
	log.Event(ctx, "dhtReceivedMessage", dht.self, p, rpmes)
	return rpmes, nil
}

// sendMessage sends out a message
func (dht *IpfsDHT) sendMessage(ctx context.Context, p peer.ID, pmes *pb.Message) error {
	ms, err := dht.messageSenderForPeer(p)
	if err != nil {
		return err
	}

	if err := ms.SendMessage(ctx, pmes); err != nil {
		return err
	}
	log.Event(ctx, "dhtSentMessage", dht.self, p, pmes)
	return nil
}

func (dht *IpfsDHT) updateFromMessage(ctx context.Context, p peer.ID, mes *pb.Message) error {
	dht.Update(ctx, p)
	return nil
}

func (dht *IpfsDHT) messageSenderForPeer(p peer.ID) (*messageSender, error) {
	dht.smlk.Lock()
	ms, ok := dht.strmap[p]
	if ok {
		dht.smlk.Unlock()
		return ms, nil
	}
	ms = &messageSender{p: p, dht: dht}
	dht.strmap[p] = ms
	dht.smlk.Unlock()

	if err := ms.prepOrInvalidate(); err != nil {
		dht.smlk.Lock()
		defer dht.smlk.Unlock()

		if msCur, ok := dht.strmap[p]; ok {
			// Changed. Use the new one, old one is invalid and
			// not in the map so we can just throw it away.
			if ms != msCur {
				return msCur, nil
			}
			// Not changed, remove the now invalid stream from the
			// map.
			delete(dht.strmap, p)
		}
		// Invalid but not in map. Must have been removed by a disconnect.
		return nil, err
	}
	// All ready to go.
	return ms, nil
}

type messageSender struct {
	s   inet.Stream
	r   ggio.ReadCloser
	w   ggio.WriteCloser
	lk  sync.Mutex
	p   peer.ID
	dht *IpfsDHT

	invalid   bool
	singleMes int
}

// invalidate is called after this messageSender is removed form the strmap.
func (ms *messageSender) invalidate() {
	ms.invalid = true
	if ms.s != nil {
		ms.s.Reset()
		ms.s = nil
	}
}

func (ms *messageSender) prepOrInvalidate() error {
	ms.lk.Lock()
	defer ms.lk.Unlock()
	if err := ms.prep(); err != nil {
		ms.invalidate()
		return err
	}
	return nil
}

func (ms *messageSender) prep() error {
	if ms.invalid {
		return fmt.Errorf("message sender has been invalidated")
	}
	if ms.s != nil {
		return nil
	}

	nstr, err := ms.dht.host.NewStream(ms.dht.ctx, ms.p, ProtocolDHT, ProtocolDHTOld)
	if err != nil {
		return err
	}

	ms.r = ggio.NewDelimitedReader(nstr, inet.MessageSizeMax)
	ms.w = ggio.NewDelimitedWriter(nstr)
	ms.s = nstr

	return nil
}

// streamReuseTries is the number of times we will try to reuse a stream to a
// given peer before giving up and reverting to the old one-message-per-stream
// behaviour.
const streamReuseTries = 3

func (ms *messageSender) SendMessage(ctx context.Context, pmes *pb.Message) error {
	ms.lk.Lock()
	defer ms.lk.Unlock()
	retry := false
	for {
		if err := ms.prep(); err != nil {
			return err
		}

		if err := ms.w.WriteMsg(pmes); err != nil {
			ms.s.Reset()
			ms.s = nil

			if retry {
				log.Info("error writing message, bailing: ", err)
				return err
			} else {
				log.Info("error writing message, trying again: ", err)
				retry = true
				continue
			}
		}

		log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)

		if ms.singleMes > streamReuseTries {
			ms.s.Close()
			ms.s = nil
		} else if retry {
			ms.singleMes++
		}

		return nil
	}
}

func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb.Message, error) {
	ms.lk.Lock()
	defer ms.lk.Unlock()
	retry := false
	for {
		if err := ms.prep(); err != nil {
			return nil, err
		}

		if err := ms.w.WriteMsg(pmes); err != nil {
			ms.s.Reset()
			ms.s = nil

			if retry {
				log.Info("error writing message, bailing: ", err)
				return nil, err
			} else {
				log.Info("error writing message, trying again: ", err)
				retry = true
				continue
			}
		}

		mes := new(pb.Message)
		if err := ms.ctxReadMsg(ctx, mes); err != nil {
			ms.s.Reset()
			ms.s = nil

			if retry {
				log.Info("error reading message, bailing: ", err)
				return nil, err
			} else {
				log.Info("error reading message, trying again: ", err)
				retry = true
				continue
			}
		}

		log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)

		if ms.singleMes > streamReuseTries {
			ms.s.Close()
			ms.s = nil
		} else if retry {
			ms.singleMes++
		}

		return mes, nil
	}
}

func (ms *messageSender) ctxReadMsg(ctx context.Context, mes *pb.Message) error {
	errc := make(chan error, 1)
	go func(r ggio.ReadCloser) {
		errc <- r.ReadMsg(mes)
	}(ms.r)

	t := time.NewTimer(dhtReadMessageTimeout)
	defer t.Stop()

	select {
	case err := <-errc:
		return err
	case <-ctx.Done():
		return ctx.Err()
	case <-t.C:
		return ErrReadTimeout
	}
}
<commit_msg>fix comment on invalidate.<commit_after>package dht

import (
	"context"
	"fmt"
	"sync"
	"time"

	ggio "github.com/gogo/protobuf/io"
	ctxio "github.com/jbenet/go-context/io"
	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
)

var dhtReadMessageTimeout = time.Minute
var ErrReadTimeout = fmt.Errorf("timed out reading response")

// handleNewStream implements the inet.StreamHandler
func (dht *IpfsDHT) handleNewStream(s inet.Stream) {
	go dht.handleNewMessage(s)
}

func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
	defer s.Close()

	ctx := dht.Context()
	cr := ctxio.NewReader(ctx, s) // ok to use. we defer close stream in this func
	cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func
	r := ggio.NewDelimitedReader(cr, inet.MessageSizeMax)
	w := ggio.NewDelimitedWriter(cw)
	mPeer := s.Conn().RemotePeer()

	for {
		// receive msg
		pmes := new(pb.Message)
		if err := r.ReadMsg(pmes); err != nil {
			s.Reset()
			log.Debugf("Error unmarshaling data: %s", err)
			return
		}

		// update the peer (on valid msgs only)
		dht.updateFromMessage(ctx, mPeer, pmes)

		// get handler for this msg type.
		handler := dht.handlerForMsgType(pmes.GetType())
		if handler == nil {
			s.Reset()
			log.Debug("got back nil handler from handlerForMsgType")
			return
		}

		// dispatch handler.
		rpmes, err := handler(ctx, mPeer, pmes)
		if err != nil {
			s.Reset()
			log.Debugf("handle message error: %s", err)
			return
		}

		// if nil response, return it before serializing
		if rpmes == nil {
			log.Debug("got back nil response from request")
			continue
		}

		// send out response msg
		if err := w.WriteMsg(rpmes); err != nil {
			s.Reset()
			log.Debugf("send response error: %s", err)
			return
		}
	}
}

// sendRequest sends out a request, but also makes sure to
// measure the RTT for latency measurements.
func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {

	ms, err := dht.messageSenderForPeer(p)
	if err != nil {
		return nil, err
	}

	start := time.Now()

	rpmes, err := ms.SendRequest(ctx, pmes)
	if err != nil {
		return nil, err
	}

	// update the peer (on valid msgs only)
	dht.updateFromMessage(ctx, p, rpmes)

	dht.peerstore.RecordLatency(p, time.Since(start))
	log.Event(ctx, "dhtReceivedMessage", dht.self, p, rpmes)
	return rpmes, nil
}

// sendMessage sends out a message
func (dht *IpfsDHT) sendMessage(ctx context.Context, p peer.ID, pmes *pb.Message) error {
	ms, err := dht.messageSenderForPeer(p)
	if err != nil {
		return err
	}

	if err := ms.SendMessage(ctx, pmes); err != nil {
		return err
	}
	log.Event(ctx, "dhtSentMessage", dht.self, p, pmes)
	return nil
}

func (dht *IpfsDHT) updateFromMessage(ctx context.Context, p peer.ID, mes *pb.Message) error {
	dht.Update(ctx, p)
	return nil
}

func (dht *IpfsDHT) messageSenderForPeer(p peer.ID) (*messageSender, error) {
	dht.smlk.Lock()
	ms, ok := dht.strmap[p]
	if ok {
		dht.smlk.Unlock()
		return ms, nil
	}
	ms = &messageSender{p: p, dht: dht}
	dht.strmap[p] = ms
	dht.smlk.Unlock()

	if err := ms.prepOrInvalidate(); err != nil {
		dht.smlk.Lock()
		defer dht.smlk.Unlock()

		if msCur, ok := dht.strmap[p]; ok {
			// Changed. Use the new one, old one is invalid and
			// not in the map so we can just throw it away.
			if ms != msCur {
				return msCur, nil
			}
			// Not changed, remove the now invalid stream from the
			// map.
			delete(dht.strmap, p)
		}
		// Invalid but not in map. Must have been removed by a disconnect.
		return nil, err
	}
	// All ready to go.
	return ms, nil
}

type messageSender struct {
	s   inet.Stream
	r   ggio.ReadCloser
	w   ggio.WriteCloser
	lk  sync.Mutex
	p   peer.ID
	dht *IpfsDHT

	invalid   bool
	singleMes int
}

// invalidate is called before this messageSender is removed from the strmap.
// It prevents the messageSender from being reused/reinitialized and then
// forgotten (leaving the stream open).
func (ms *messageSender) invalidate() {
	ms.invalid = true
	if ms.s != nil {
		ms.s.Reset()
		ms.s = nil
	}
}

func (ms *messageSender) prepOrInvalidate() error {
	ms.lk.Lock()
	defer ms.lk.Unlock()
	if err := ms.prep(); err != nil {
		ms.invalidate()
		return err
	}
	return nil
}

func (ms *messageSender) prep() error {
	if ms.invalid {
		return fmt.Errorf("message sender has been invalidated")
	}
	if ms.s != nil {
		return nil
	}

	nstr, err := ms.dht.host.NewStream(ms.dht.ctx, ms.p, ProtocolDHT, ProtocolDHTOld)
	if err != nil {
		return err
	}

	ms.r = ggio.NewDelimitedReader(nstr, inet.MessageSizeMax)
	ms.w = ggio.NewDelimitedWriter(nstr)
	ms.s = nstr

	return nil
}

// streamReuseTries is the number of times we will try to reuse a stream to a
// given peer before giving up and reverting to the old one-message-per-stream
// behaviour.
const streamReuseTries = 3

func (ms *messageSender) SendMessage(ctx context.Context, pmes *pb.Message) error {
	ms.lk.Lock()
	defer ms.lk.Unlock()
	retry := false
	for {
		if err := ms.prep(); err != nil {
			return err
		}

		if err := ms.w.WriteMsg(pmes); err != nil {
			ms.s.Reset()
			ms.s = nil

			if retry {
				log.Info("error writing message, bailing: ", err)
				return err
			} else {
				log.Info("error writing message, trying again: ", err)
				retry = true
				continue
			}
		}

		log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)

		if ms.singleMes > streamReuseTries {
			ms.s.Close()
			ms.s = nil
		} else if retry {
			ms.singleMes++
		}

		return nil
	}
}

func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb.Message, error) {
	ms.lk.Lock()
	defer ms.lk.Unlock()
	retry := false
	for {
		if err := ms.prep(); err != nil {
			return nil, err
		}

		if err := ms.w.WriteMsg(pmes); err != nil {
			ms.s.Reset()
			ms.s = nil

			if retry {
				log.Info("error writing message, bailing: ", err)
				return nil, err
			} else {
				log.Info("error writing message, trying again: ", err)
				retry = true
				continue
			}
		}

		mes := new(pb.Message)
		if err := ms.ctxReadMsg(ctx, mes); err != nil {
			ms.s.Reset()
			ms.s = nil

			if retry {
				log.Info("error reading message, bailing: ", err)
				return nil, err
			} else {
				log.Info("error reading message, trying again: ", err)
				retry = true
				continue
			}
		}

		log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)

		if ms.singleMes > streamReuseTries {
			ms.s.Close()
			ms.s = nil
		} else if retry {
			ms.singleMes++
		}

		return mes, nil
	}
}

func (ms *messageSender) ctxReadMsg(ctx context.Context, mes *pb.Message) error {
	errc := make(chan error, 1)
	go func(r ggio.ReadCloser) {
		errc <- r.ReadMsg(mes)
	}(ms.r)

	t := time.NewTimer(dhtReadMessageTimeout)
	defer t.Stop()

	select {
	case err := <-errc:
		return err
	case <-ctx.Done():
		return ctx.Err()
	case <-t.C:
		return ErrReadTimeout
	}
}
<|endoftext|>


<commit_before>package repository

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"regexp"
	"strings"
	"time"

	"github.com/ksonnet/ksonnet/pkg/app"
	log "github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"

	"github.com/argoproj/argo-cd/common"
	"github.com/argoproj/argo-cd/pkg/apis/application/v1alpha1"
	"github.com/argoproj/argo-cd/util"
	"github.com/argoproj/argo-cd/util/cache"
	"github.com/argoproj/argo-cd/util/git"
	"github.com/argoproj/argo-cd/util/helm"
	"github.com/argoproj/argo-cd/util/ksonnet"
	"github.com/argoproj/argo-cd/util/kube"
	"github.com/argoproj/argo-cd/util/kustomize"
)

const (
	// DefaultRepoCacheExpiration is the duration for items to live in the repo cache
	DefaultRepoCacheExpiration = 24 * time.Hour
)

type AppSourceType string

const (
	AppSourceKsonnet   AppSourceType = "ksonnet"
	AppSourceHelm      AppSourceType = "helm"
	AppSourceKustomize AppSourceType = "kustomize"
	AppSourceDirectory AppSourceType = "directory"
)

// Service implements ManifestService interface
type Service struct {
	repoLock   *util.KeyLock
	gitFactory git.ClientFactory
	cache      cache.Cache
}

// NewService returns a new instance of the Manifest service
func NewService(gitFactory git.ClientFactory, cache cache.Cache) *Service {
	return &Service{
		repoLock:   util.NewKeyLock(),
		gitFactory: gitFactory,
		cache:      cache,
	}
}

// ListDir lists the contents of a GitHub repo
func (s *Service) ListDir(ctx context.Context, q *ListDirRequest) (*FileList, error) {
	appRepoPath := tempRepoPath(q.Repo.Repo)
	s.repoLock.Lock(appRepoPath)
	defer s.repoLock.Unlock(appRepoPath)

	gitClient := s.gitFactory.NewClient(q.Repo.Repo, appRepoPath, q.Repo.Username, q.Repo.Password, q.Repo.SSHPrivateKey)
	err := gitClient.Init()
	if err != nil {
		return nil, err
	}

	commitSHA, err := gitClient.LsRemote(q.Revision)
	if err != nil {
		return nil, err
	}
	cacheKey := listDirCacheKey(commitSHA, q)
	var res FileList
	err = s.cache.Get(cacheKey, &res)
	if err == nil {
		log.Infof("listdir cache hit: %s", cacheKey)
		return &res, nil
	}

	err = checkoutRevision(gitClient, q.Revision)
	if err != nil {
		return nil, err
	}

	lsFiles, err := gitClient.LsFiles(q.Path)
	if err != nil {
		return nil, err
	}

	res = FileList{
		Items: lsFiles,
	}
	err = s.cache.Set(&cache.Item{
		Key:        cacheKey,
		Object:     &res,
		Expiration: DefaultRepoCacheExpiration,
	})
	if err != nil {
		log.Warnf("listdir cache set error %s: %v", cacheKey, err)
	}
	return &res, nil
}

func (s *Service) GetFile(ctx context.Context, q *GetFileRequest) (*GetFileResponse, error) {
	appRepoPath := tempRepoPath(q.Repo.Repo)
	s.repoLock.Lock(appRepoPath)
	defer s.repoLock.Unlock(appRepoPath)

	gitClient := s.gitFactory.NewClient(q.Repo.Repo, appRepoPath, q.Repo.Username, q.Repo.Password, q.Repo.SSHPrivateKey)
	err := gitClient.Init()
	if err != nil {
		return nil, err
	}
	err = checkoutRevision(gitClient, q.Revision)
	if err != nil {
		return nil, err
	}
	data, err := ioutil.ReadFile(path.Join(gitClient.Root(), q.Path))
	if err != nil {
		return nil, err
	}
	res := GetFileResponse{
		Data: data,
	}
	return &res, nil
}

func (s *Service) GenerateManifest(c context.Context, q *ManifestRequest) (*ManifestResponse, error) {
	var res ManifestResponse
	if git.IsCommitSHA(q.Revision) {
		cacheKey := manifestCacheKey(q.Revision, q)
		err := s.cache.Get(cacheKey, res)
		if err == nil {
			log.Infof("manifest cache hit: %s", cacheKey)
			return &res, nil
		}
	}
	appRepoPath := tempRepoPath(q.Repo.Repo)
	s.repoLock.Lock(appRepoPath)
	defer s.repoLock.Unlock(appRepoPath)

	gitClient := s.gitFactory.NewClient(q.Repo.Repo, appRepoPath, q.Repo.Username, q.Repo.Password, q.Repo.SSHPrivateKey)
	err := gitClient.Init()
	if err != nil {
		return nil, err
	}
	commitSHA, err := gitClient.LsRemote(q.Revision)
	if err != nil {
		return nil, err
	}
	cacheKey := manifestCacheKey(commitSHA, q)
	err = s.cache.Get(cacheKey, &res)
	if err == nil {
		log.Infof("manifest cache hit: %s", cacheKey)
		return &res, nil
	}
	if err != cache.ErrCacheMiss {
		log.Warnf("manifest cache error %s: %v", cacheKey, err)
	} else {
		log.Infof("manifest cache miss: %s", cacheKey)
	}

	err = checkoutRevision(gitClient, q.Revision)
	if err != nil {
		return nil, err
	}
	appPath := path.Join(appRepoPath, q.Path)

	genRes, err := generateManifests(appPath, q)
	if err != nil {
		return nil, err
	}
	res = *genRes
	res.Revision = commitSHA
	err = s.cache.Set(&cache.Item{
		Key:        cacheKey,
		Object:     res,
		Expiration: DefaultRepoCacheExpiration,
	})
	if err != nil {
		log.Warnf("manifest cache set error %s: %v", cacheKey, err)
	}
	return &res, nil
}

// generateManifests generates manifests from a path
func generateManifests(appPath string, q *ManifestRequest) (*ManifestResponse, error) {
	var targetObjs []*unstructured.Unstructured
	var params []*v1alpha1.ComponentParameter
	var env *app.EnvironmentConfig
	var err error

	appSourceType := IdentifyAppSourceTypeByAppDir(appPath)
	switch appSourceType {
	case AppSourceKsonnet:
		targetObjs, params, env, err = ksShow(appPath, q.Environment, q.ComponentParameterOverrides)
	case AppSourceHelm:
		h := helm.NewHelmApp(appPath)
		targetObjs, err = h.Template(q.AppLabel, q.ValueFiles, q.ComponentParameterOverrides)
		if err != nil {
			return nil, err
		}
		params, err = h.GetParameters(q.ValueFiles)
		if err != nil {
			return nil, err
		}
	case AppSourceKustomize:
		k := kustomize.NewKustomizeApp(appPath)
		targetObjs, err = k.Build()
	case AppSourceDirectory:
		targetObjs, err = findManifests(appPath)
	}
	if err != nil {
		return nil, err
	}
	// TODO(jessesuen): we need to sort objects based on their dependency order of creation

	manifests := make([]string, len(targetObjs))
	for i, target := range targetObjs {
		if q.AppLabel != "" {
			err = kube.SetLabel(target, common.LabelApplicationName, q.AppLabel)
			if err != nil {
				return nil, err
			}
		}
		manifestStr, err := json.Marshal(target.Object)
		if err != nil {
			return nil, err
		}
		manifests[i] = string(manifestStr)
	}

	res := ManifestResponse{
		Manifests: manifests,
		Params:    params,
	}
	if env != nil {
		res.Namespace = env.Destination.Namespace
		res.Server = env.Destination.Server
	}
	return &res, nil
}

// tempRepoPath returns a formulated temporary directory location to clone a repository
func tempRepoPath(repo string) string {
	return path.Join(os.TempDir(), strings.Replace(repo, "/", "_", -1))
}

// IdentifyAppSourceTypeByAppDir examines a directory and determines its application source type
func IdentifyAppSourceTypeByAppDir(appDirPath string) AppSourceType {
	if pathExists(path.Join(appDirPath, "app.yaml")) {
		return AppSourceKsonnet
	}
	if pathExists(path.Join(appDirPath, "Chart.yaml")) {
		return AppSourceHelm
	}
	if pathExists(path.Join(appDirPath, "kustomization.yaml")) {
		return AppSourceKustomize
	}
	return AppSourceDirectory
}

// IdentifyAppSourceTypeByAppPath determines application source type by app file path
func IdentifyAppSourceTypeByAppPath(appFilePath string) AppSourceType {
	if strings.HasSuffix(appFilePath, "app.yaml") {
		return AppSourceKsonnet
	}
	if strings.HasSuffix(appFilePath, "Chart.yaml") {
		return AppSourceHelm
	}
	if strings.HasSuffix(appFilePath, "kustomization.yaml") {
		return AppSourceKustomize
	}
	return AppSourceDirectory
}

// checkoutRevision is a convenience function to initialize a repo, fetch, and checkout a revision
func checkoutRevision(gitClient git.Client, revision string) error {
	err := gitClient.Fetch()
	if err != nil {
		return err
	}
	err = gitClient.Reset()
	if err != nil {
		log.Warn(err)
	}
	err = gitClient.Checkout(revision)
	if err != nil {
		return err
	}
	return nil
}

func manifestCacheKey(commitSHA string, q *ManifestRequest) string {
	pStr, _ := json.Marshal(q.ComponentParameterOverrides)
	valuesFiles := strings.Join(q.ValueFiles, ",")
	return fmt.Sprintf("mfst|%s|%s|%s|%s|%s|%s", q.AppLabel, q.Path, q.Environment, commitSHA, string(pStr), valuesFiles)
}

func listDirCacheKey(commitSHA string, q *ListDirRequest) string {
	return fmt.Sprintf("ldir|%s|%s", q.Path, commitSHA)
}

// ksShow runs `ks show` in an app directory after setting any component parameter overrides
func ksShow(appPath, envName string, overrides []*v1alpha1.ComponentParameter) ([]*unstructured.Unstructured, []*v1alpha1.ComponentParameter, *app.EnvironmentConfig, error) {
	ksApp, err := ksonnet.NewKsonnetApp(appPath)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("unable to load application from %s: %v", appPath, err)
	}
	params, err := ksApp.ListEnvParams(envName)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("Failed to list ksonnet app params: %v", err)
	}
	if overrides != nil {
		for _, override := range overrides {
			err = ksApp.SetComponentParams(envName, override.Component, override.Name, override.Value)
			if err != nil {
				return nil, nil, nil, err
			}
		}
	}
	appSpec := ksApp.App()
	env, err := appSpec.Environment(envName)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("environment '%s' does not exist in ksonnet app", envName)
	}
	targetObjs, err := ksApp.Show(envName)
	if err != nil {
		return nil, nil, nil, err
	}
	return targetObjs, params, env, nil
}

var manifestFile = regexp.MustCompile(`^.*\.(yaml|yml|json)$`)

// findManifests looks at all yaml files in a directory and unmarshals them into a list of unstructured objects
func findManifests(appPath string) ([]*unstructured.Unstructured, error) {
	files, err := ioutil.ReadDir(appPath)
	if err != nil {
		return nil, fmt.Errorf("Failed to read dir %s: %v", appPath, err)
	}
	var objs []*unstructured.Unstructured
	for _, f := range files {
		if f.IsDir() || !manifestFile.MatchString(f.Name()) {
			continue
		}
		out, err := ioutil.ReadFile(path.Join(appPath, f.Name()))
		if err != nil {
			return nil, err
		}
		if strings.HasSuffix(f.Name(), ".json") {
			var obj unstructured.Unstructured
			err = json.Unmarshal(out, &obj)
			if err != nil {
				return nil, fmt.Errorf("Failed to unmarshal '%s': %v", f.Name(), err)
			}
			objs = append(objs, &obj)
		} else {
			yamlObjs, err := kube.SplitYAML(string(out))
			if err != nil {
				if len(yamlObjs) > 0 {
					// If we get here, we had a multiple objects in a single YAML file which had some
					// valid k8s objects, but errors parsing others (within the same file). It's very
					// likely the user messed up a portion of the YAML, so report on that.
					return nil, fmt.Errorf("Failed to unmarshal '%s': %v", f.Name(), err)
				}
				// Otherwise, it might be a unrelated YAML file which we will ignore
				continue
			}
			objs = append(objs, yamlObjs...)
		}
	}
	return objs, nil
}

// pathExists reports whether the named file or directory exists.
func pathExists(name string) bool {
	if _, err := os.Stat(name); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}
<commit_msg>Use gRPC error codes instead of fmt.Errorf (#532)<commit_after>package repository

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"regexp"
	"strings"
	"time"

	"github.com/ksonnet/ksonnet/pkg/app"
	log "github.com/sirupsen/logrus"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"

	"github.com/argoproj/argo-cd/common"
	"github.com/argoproj/argo-cd/pkg/apis/application/v1alpha1"
	"github.com/argoproj/argo-cd/util"
	"github.com/argoproj/argo-cd/util/cache"
	"github.com/argoproj/argo-cd/util/git"
	"github.com/argoproj/argo-cd/util/helm"
	"github.com/argoproj/argo-cd/util/ksonnet"
	"github.com/argoproj/argo-cd/util/kube"
	"github.com/argoproj/argo-cd/util/kustomize"
)

const (
	// DefaultRepoCacheExpiration is the duration for items to live in the repo cache
	DefaultRepoCacheExpiration = 24 * time.Hour
)

type AppSourceType string

const (
	AppSourceKsonnet   AppSourceType = "ksonnet"
	AppSourceHelm      AppSourceType = "helm"
	AppSourceKustomize AppSourceType = "kustomize"
	AppSourceDirectory AppSourceType = "directory"
)

// Service implements ManifestService interface
type Service struct {
	repoLock   *util.KeyLock
	gitFactory git.ClientFactory
	cache      cache.Cache
}

// NewService returns a new instance of the Manifest service
func NewService(gitFactory git.ClientFactory, cache cache.Cache) *Service {
	return &Service{
		repoLock:   util.NewKeyLock(),
		gitFactory: gitFactory,
		cache:      cache,
	}
}

// ListDir lists the contents of a GitHub repo
func (s *Service) ListDir(ctx context.Context, q *ListDirRequest) (*FileList, error) {
	appRepoPath := tempRepoPath(q.Repo.Repo)
	s.repoLock.Lock(appRepoPath)
	defer s.repoLock.Unlock(appRepoPath)

	gitClient := s.gitFactory.NewClient(q.Repo.Repo, appRepoPath, q.Repo.Username, q.Repo.Password, q.Repo.SSHPrivateKey)
	err := gitClient.Init()
	if err != nil {
		return nil, err
	}

	commitSHA, err := gitClient.LsRemote(q.Revision)
	if err != nil {
		return nil, err
	}
	cacheKey := listDirCacheKey(commitSHA, q)
	var res FileList
	err = s.cache.Get(cacheKey, &res)
	if err == nil {
		log.Infof("listdir cache hit: %s", cacheKey)
		return &res, nil
	}

	err = checkoutRevision(gitClient, q.Revision)
	if err != nil {
		return nil, err
	}

	lsFiles, err := gitClient.LsFiles(q.Path)
	if err != nil {
		return nil, err
	}

	res = FileList{
		Items: lsFiles,
	}
	err = s.cache.Set(&cache.Item{
		Key:        cacheKey,
		Object:     &res,
		Expiration: DefaultRepoCacheExpiration,
	})
	if err != nil {
		log.Warnf("listdir cache set error %s: %v", cacheKey, err)
	}
	return &res, nil
}

func (s *Service) GetFile(ctx context.Context, q *GetFileRequest) (*GetFileResponse, error) {
	appRepoPath := tempRepoPath(q.Repo.Repo)
	s.repoLock.Lock(appRepoPath)
	defer s.repoLock.Unlock(appRepoPath)

	gitClient := s.gitFactory.NewClient(q.Repo.Repo, appRepoPath, q.Repo.Username, q.Repo.Password, q.Repo.SSHPrivateKey)
	err := gitClient.Init()
	if err != nil {
		return nil, err
	}
	err = checkoutRevision(gitClient, q.Revision)
	if err != nil {
		return nil, err
	}
	data, err := ioutil.ReadFile(path.Join(gitClient.Root(), q.Path))
	if err != nil {
		return nil, err
	}
	res := GetFileResponse{
		Data: data,
	}
	return &res, nil
}

func (s *Service) GenerateManifest(c context.Context, q *ManifestRequest) (*ManifestResponse, error) {
	var res ManifestResponse
	if git.IsCommitSHA(q.Revision) {
		cacheKey := manifestCacheKey(q.Revision, q)
		err := s.cache.Get(cacheKey, res)
		if err == nil {
			log.Infof("manifest cache hit: %s", cacheKey)
			return &res, nil
		}
	}
	appRepoPath := tempRepoPath(q.Repo.Repo)
	s.repoLock.Lock(appRepoPath)
	defer s.repoLock.Unlock(appRepoPath)

	gitClient := s.gitFactory.NewClient(q.Repo.Repo, appRepoPath, q.Repo.Username, q.Repo.Password, q.Repo.SSHPrivateKey)
	err := gitClient.Init()
	if err != nil {
		return nil, err
	}
	commitSHA, err := gitClient.LsRemote(q.Revision)
	if err != nil {
		return nil, err
	}
	cacheKey := manifestCacheKey(commitSHA, q)
	err = s.cache.Get(cacheKey, &res)
	if err == nil {
		log.Infof("manifest cache hit: %s", cacheKey)
		return &res, nil
	}
	if err != cache.ErrCacheMiss {
		log.Warnf("manifest cache error %s: %v", cacheKey, err)
	} else {
		log.Infof("manifest cache miss: %s", cacheKey)
	}

	err = checkoutRevision(gitClient, q.Revision)
	if err != nil {
		return nil, err
	}
	appPath := path.Join(appRepoPath, q.Path)

	genRes, err := generateManifests(appPath, q)
	if err != nil {
		return nil, err
	}
	res = *genRes
	res.Revision = commitSHA
	err = s.cache.Set(&cache.Item{
		Key:        cacheKey,
		Object:     res,
		Expiration: DefaultRepoCacheExpiration,
	})
	if err != nil {
		log.Warnf("manifest cache set error %s: %v", cacheKey, err)
	}
	return &res, nil
}

// generateManifests generates manifests from a path
func generateManifests(appPath string, q *ManifestRequest) (*ManifestResponse, error) {
	var targetObjs []*unstructured.Unstructured
	var params []*v1alpha1.ComponentParameter
	var env *app.EnvironmentConfig
	var err error

	appSourceType := IdentifyAppSourceTypeByAppDir(appPath)
	switch appSourceType {
	case AppSourceKsonnet:
		targetObjs, params, env, err = ksShow(appPath, q.Environment, q.ComponentParameterOverrides)
	case AppSourceHelm:
		h := helm.NewHelmApp(appPath)
		targetObjs, err = h.Template(q.AppLabel, q.ValueFiles, q.ComponentParameterOverrides)
		if err != nil {
			return nil, err
		}
		params, err = h.GetParameters(q.ValueFiles)
		if err != nil {
			return nil, err
		}
	case AppSourceKustomize:
		k := kustomize.NewKustomizeApp(appPath)
		targetObjs, err = k.Build()
	case AppSourceDirectory:
		targetObjs, err = findManifests(appPath)
	}
	if err != nil {
		return nil, err
	}
	// TODO(jessesuen): we need to sort objects based on their dependency order of creation

	manifests := make([]string, len(targetObjs))
	for i, target := range targetObjs {
		if q.AppLabel != "" {
			err = kube.SetLabel(target, common.LabelApplicationName, q.AppLabel)
			if err != nil {
				return nil, err
			}
		}
		manifestStr, err := json.Marshal(target.Object)
		if err != nil {
			return nil, err
		}
		manifests[i] = string(manifestStr)
	}

	res := ManifestResponse{
		Manifests: manifests,
		Params:    params,
	}
	if env != nil {
		res.Namespace = env.Destination.Namespace
		res.Server = env.Destination.Server
	}
	return &res, nil
}

// tempRepoPath returns a formulated temporary directory location to clone a repository
func tempRepoPath(repo string) string {
	return path.Join(os.TempDir(), strings.Replace(repo, "/", "_", -1))
}

// IdentifyAppSourceTypeByAppDir examines a directory and determines its application source type
func IdentifyAppSourceTypeByAppDir(appDirPath string) AppSourceType {
	if pathExists(path.Join(appDirPath, "app.yaml")) {
		return AppSourceKsonnet
	}
	if pathExists(path.Join(appDirPath, "Chart.yaml")) {
		return AppSourceHelm
	}
	if pathExists(path.Join(appDirPath, "kustomization.yaml")) {
		return AppSourceKustomize
	}
	return AppSourceDirectory
}

// IdentifyAppSourceTypeByAppPath determines application source type by app file path
func IdentifyAppSourceTypeByAppPath(appFilePath string) AppSourceType {
	if strings.HasSuffix(appFilePath, "app.yaml") {
		return AppSourceKsonnet
	}
	if strings.HasSuffix(appFilePath, "Chart.yaml") {
		return AppSourceHelm
	}
	if strings.HasSuffix(appFilePath, "kustomization.yaml") {
		return AppSourceKustomize
	}
	return AppSourceDirectory
}

// checkoutRevision is a convenience function to initialize a repo, fetch, and checkout a revision
func checkoutRevision(gitClient git.Client, revision string) error {
	err := gitClient.Fetch()
	if err != nil {
		return err
	}
	err = gitClient.Reset()
	if err != nil {
		log.Warn(err)
	}
	err = gitClient.Checkout(revision)
	if err != nil {
		return err
	}
	return nil
}

func manifestCacheKey(commitSHA string, q *ManifestRequest) string {
	pStr, _ := json.Marshal(q.ComponentParameterOverrides)
	valuesFiles := strings.Join(q.ValueFiles, ",")
	return fmt.Sprintf("mfst|%s|%s|%s|%s|%s|%s", q.AppLabel, q.Path, q.Environment, commitSHA, string(pStr), valuesFiles)
}

func listDirCacheKey(commitSHA string, q *ListDirRequest) string {
	return fmt.Sprintf("ldir|%s|%s", q.Path, commitSHA)
}

// ksShow runs `ks show` in an app directory after setting any component parameter overrides
func ksShow(appPath, envName string, overrides []*v1alpha1.ComponentParameter) ([]*unstructured.Unstructured, []*v1alpha1.ComponentParameter, *app.EnvironmentConfig, error) {
	ksApp, err := ksonnet.NewKsonnetApp(appPath)
	if err != nil {
		return nil, nil, nil, status.Errorf(codes.FailedPrecondition, "unable to load application from %s: %v", appPath, err)
	}
	params, err := ksApp.ListEnvParams(envName)
	if err != nil {
		return nil, nil, nil, status.Errorf(codes.InvalidArgument, "Failed to list ksonnet app params: %v", err)
	}
	if overrides != nil {
		for _, override := range overrides {
			err = ksApp.SetComponentParams(envName, override.Component, override.Name, override.Value)
			if err != nil {
				return nil, nil, nil, err
			}
		}
	}
	appSpec := ksApp.App()
	env, err := appSpec.Environment(envName)
	if err != nil {
		return nil, nil, nil, status.Errorf(codes.NotFound, "environment %q does not exist in ksonnet app", envName)
	}
	targetObjs, err := ksApp.Show(envName)
	if err != nil {
		return nil, nil, nil, err
	}
	return targetObjs, params, env, nil
}

var manifestFile = regexp.MustCompile(`^.*\.(yaml|yml|json)$`)

// findManifests looks at all yaml files in a directory and unmarshals them into a list of unstructured objects
func findManifests(appPath string) ([]*unstructured.Unstructured, error) {
	files, err := ioutil.ReadDir(appPath)
	if err != nil {
		return nil, status.Errorf(codes.FailedPrecondition, "Failed to read dir %s: %v", appPath, err)
	}
	var objs []*unstructured.Unstructured
	for _, f := range files {
		if f.IsDir() || !manifestFile.MatchString(f.Name()) {
			continue
		}
		out, err := ioutil.ReadFile(path.Join(appPath, f.Name()))
		if err != nil {
			return nil, err
		}
		if strings.HasSuffix(f.Name(), ".json") {
			var obj unstructured.Unstructured
			err = json.Unmarshal(out, &obj)
			if err != nil {
				return nil, status.Errorf(codes.FailedPrecondition, "Failed to unmarshal %q: %v", f.Name(), err)
			}
			objs = append(objs, &obj)
		} else {
			yamlObjs, err := kube.SplitYAML(string(out))
			if err != nil {
				if len(yamlObjs) > 0 {
					// If we get here, we had a multiple objects in a single YAML file which had some
					// valid k8s objects, but errors parsing others (within the same file). It's very
					// likely the user messed up a portion of the YAML, so report on that.
					return nil, status.Errorf(codes.FailedPrecondition, "Failed to unmarshal %q: %v", f.Name(), err)
				}
				// Otherwise, it might be a unrelated YAML file which we will ignore
				continue
			}
			objs = append(objs, yamlObjs...)
		}
	}
	return objs, nil
}

// pathExists reports whether the named file or directory exists.
func pathExists(name string) bool {
	if _, err := os.Stat(name); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>// +build linux

package fscommon

import (
	"io/ioutil"
	"os"
	"syscall"

	securejoin "github.com/cyphar/filepath-securejoin"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

func WriteFile(dir, file, data string) error {
	if dir == "" {
		return errors.Errorf("no directory specified for %s", file)
	}
	path, err := securejoin.SecureJoin(dir, file)
	if err != nil {
		return err
	}
	if err := retryingWriteFile(path, []byte(data), 0700); err != nil {
		return errors.Wrapf(err, "failed to write %q to %q", data, path)
	}
	return nil
}

func ReadFile(dir, file string) (string, error) {
	if dir == "" {
		return "", errors.Errorf("no directory specified for %s", file)
	}
	path, err := securejoin.SecureJoin(dir, file)
	if err != nil {
		return "", err
	}
	data, err := ioutil.ReadFile(path)
	return string(data), err
}

func retryingWriteFile(filename string, data []byte, perm os.FileMode) error {
	for {
		err := ioutil.WriteFile(filename, data, perm)
		if isInterruptedWriteFile(err) {
			logrus.Infof("interrupted while writing %s to %s", string(data), filename)
			continue
		}
		return err
	}
}

func isInterruptedWriteFile(err error) bool {
	if patherr, ok := err.(*os.PathError); ok {
		errno, ok2 := patherr.Err.(syscall.Errno)
		if ok2 && errno == syscall.EINTR {
			return true
		}
	}
	return false
}
<commit_msg>libct/cgroups/fscommon.WriteFile: use errors.Unwrap<commit_after>// +build linux

package fscommon

import (
	"io/ioutil"
	"os"
	"syscall"

	securejoin "github.com/cyphar/filepath-securejoin"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

func WriteFile(dir, file, data string) error {
	if dir == "" {
		return errors.Errorf("no directory specified for %s", file)
	}
	path, err := securejoin.SecureJoin(dir, file)
	if err != nil {
		return err
	}
	if err := retryingWriteFile(path, []byte(data), 0700); err != nil {
		return errors.Wrapf(err, "failed to write %q to %q", data, path)
	}
	return nil
}

func ReadFile(dir, file string) (string, error) {
	if dir == "" {
		return "", errors.Errorf("no directory specified for %s", file)
	}
	path, err := securejoin.SecureJoin(dir, file)
	if err != nil {
		return "", err
	}
	data, err := ioutil.ReadFile(path)
	return string(data), err
}

func retryingWriteFile(filename string, data []byte, perm os.FileMode) error {
	for {
		err := ioutil.WriteFile(filename, data, perm)
		if errors.Unwrap(err) == syscall.EINTR {
			logrus.Infof("interrupted while writing %s to %s", string(data), filename)
			continue
		}
		return err
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 The LUCI 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 main

import (
	"context"
	"time"

	"cloud.google.com/go/spanner"
	tspb "github.com/golang/protobuf/ptypes/timestamp"
	"google.golang.org/grpc/metadata"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/grpc/grpcutil"

	"go.chromium.org/luci/resultdb/cmd/recorder/chromium"
	"go.chromium.org/luci/resultdb/internal/span"
	"go.chromium.org/luci/resultdb/pbutil"
	pb "go.chromium.org/luci/resultdb/proto/rpc/v1"
)

// updateTokenMetadataKey is the metadata.MD key for the secret update token
// required to mutate an invocation.
// It is returned by CreateInvocation RPC in response header metadata,
// and is required by all RPCs mutating an invocation.
const updateTokenMetadataKey = "update-token"

// mutateInvocation checks if the invocation can be mutated and also
// finalizes the invocation if it's deadline is exceeded.
// If the invocation is active, continue with the other mutation(s) in f.
func mutateInvocation(ctx context.Context, invID string, f func(context.Context, *spanner.ReadWriteTransaction) error) error {
	var retErr error

	_, err := span.ReadWriteTransaction(ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
		userToken, err := extractUserUpdateToken(ctx)
		if err != nil {
			return err
		}

		now := clock.Now(ctx)

		var updateToken spanner.NullString
		var state pb.Invocation_State
		var deadline time.Time
		err = span.ReadInvocation(ctx, txn, invID, map[string]interface{}{
			"UpdateToken": &updateToken,
			"State":       &state,
			"Deadline":    &deadline,
		})

		switch {
		case err != nil:
			return err

		case state != pb.Invocation_ACTIVE:
			return errors.Reason("%q is not active", pbutil.InvocationName(invID)).Tag(grpcutil.FailedPreconditionTag).Err()

		case deadline.Before(now):
			retErr = errors.Reason("%q is not active", pbutil.InvocationName(invID)).Tag(grpcutil.FailedPreconditionTag).Err()

			// The invocation has exceeded deadline, finalize it now.
			return finalizeInvocation(txn, invID, true, pbutil.MustTimestampProto(deadline))
		}

		if err = validateUserUpdateToken(updateToken, userToken); err != nil {
			return err
		}

		return f(ctx, txn)
	})

	if err != nil {
		retErr = err
	}
	return retErr
}

func extractUserUpdateToken(ctx context.Context) (string, error) {
	md, _ := metadata.FromIncomingContext(ctx)
	userToken := md.Get(updateTokenMetadataKey)
	switch {
	case len(userToken) == 0:
		return "", errors.
			Reason("missing %q metadata value in the request", updateTokenMetadataKey).
			Tag(grpcutil.UnauthenticatedTag).
			Err()

	case len(userToken) > 1:
		return "", errors.
			Reason("expected exactly one %q metadata value, got %d", updateTokenMetadataKey, len(userToken)).
			Tag(grpcutil.InvalidArgumentTag).
			Err()

	default:
		return userToken[0], nil
	}
}

func finalizeInvocation(txn *spanner.ReadWriteTransaction, invID string, interrupted bool, finalizeTime *tspb.Timestamp) error {
	state := pb.Invocation_COMPLETED
	if interrupted {
		state = pb.Invocation_INTERRUPTED
	}

	// TODO(chanli): Also update all inclusions that include this invocation.
	return txn.BufferWrite([]*spanner.Mutation{
		spanner.UpdateMap("Invocations", span.ToSpannerMap(map[string]interface{}{
			"InvocationId": invID,
			"State":        state,
			"FinalizeTime": finalizeTime,
		})),
	})
}

func validateUserUpdateToken(updateToken spanner.NullString, userToken string) error {
	if !updateToken.Valid {
		return errors.Reason("no update token in active invocation").Tag(grpcutil.InternalTag).Err()
	}

	if userToken != updateToken.StringVal {
		return errors.Reason("invalid update token").Tag(grpcutil.PermissionDeniedTag).Err()
	}

	return nil
}

func readInvocationState(ctx context.Context, txn span.Txn, invID string) (pb.Invocation_State, error) {
	var state pb.Invocation_State
	err := span.ReadInvocation(ctx, txn, invID, map[string]interface{}{"State": &state})
	return state, err
}

func insertInvocationsByTag(invID string, inv *pb.Invocation) []*spanner.Mutation {
	muts := make([]*spanner.Mutation, len(inv.Tags))
	for i, tag := range inv.Tags {
		muts[i] = spanner.InsertMap("InvocationsByTag", map[string]interface{}{
			"TagId":        span.TagID(tag),
			"InvocationId": invID,
		})
	}
	return muts
}

// insertInvocation returns an spanner mutation that inserts an Invocation row.
// Uses the value of clock.Now(ctx) to compute expiration times.
// Assumes inv is complete and valid; may panic otherwise.
func insertInvocation(ctx context.Context, inv *pb.Invocation, updateToken, createRequestID string) *spanner.Mutation {
	row := map[string]interface{}{
		"InvocationId": pbutil.MustParseInvocationName(inv.Name),
		"State":        inv.State,
		"Realm":        chromium.Realm, // TODO(crbug.com/1013316): accept realm in the proto

		"UpdateToken": updateToken,

		"CreateTime": inv.CreateTime,
		"Deadline":   inv.Deadline,

		"BaseTestVariantDef": inv.BaseTestVariantDef,
		"Tags":               inv.Tags,
	}

	if inv.FinalizeTime != nil {
		row["FinalizeTime"] = inv.FinalizeTime
	}

	if createRequestID != "" {
		row["CreateRequestId"] = createRequestID
	}

	populateExpirations(row, clock.Now(ctx))

	return spanner.InsertMap("Invocations", span.ToSpannerMap(row))
}

// populateExpirations populates the invocation row's expiration fields using the given current time.
func populateExpirations(invRow map[string]interface{}, now time.Time) {
	invExp := now.Add(invocationExpirationDuration)
	invRow["InvocationExpirationTime"] = invExp
	invRow["InvocationExpirationWeek"] = invExp.Truncate(week)

	resultsExp := now.Add(expectedTestResultsExpirationDuration)
	invRow["ExpectedTestResultsExpirationTime"] = resultsExp
	invRow["ExpectedTestResultsExpirationWeek"] = resultsExp.Truncate(week)
}
<commit_msg>[resultdb] remove TODO from recorder/invocation.go<commit_after>// Copyright 2019 The LUCI 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 main

import (
	"context"
	"time"

	"cloud.google.com/go/spanner"
	tspb "github.com/golang/protobuf/ptypes/timestamp"
	"google.golang.org/grpc/metadata"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/grpc/grpcutil"

	"go.chromium.org/luci/resultdb/cmd/recorder/chromium"
	"go.chromium.org/luci/resultdb/internal/span"
	"go.chromium.org/luci/resultdb/pbutil"
	pb "go.chromium.org/luci/resultdb/proto/rpc/v1"
)

// updateTokenMetadataKey is the metadata.MD key for the secret update token
// required to mutate an invocation.
// It is returned by CreateInvocation RPC in response header metadata,
// and is required by all RPCs mutating an invocation.
const updateTokenMetadataKey = "update-token"

// mutateInvocation checks if the invocation can be mutated and also
// finalizes the invocation if it's deadline is exceeded.
// If the invocation is active, continue with the other mutation(s) in f.
func mutateInvocation(ctx context.Context, invID string, f func(context.Context, *spanner.ReadWriteTransaction) error) error {
	var retErr error

	_, err := span.ReadWriteTransaction(ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
		userToken, err := extractUserUpdateToken(ctx)
		if err != nil {
			return err
		}

		now := clock.Now(ctx)

		var updateToken spanner.NullString
		var state pb.Invocation_State
		var deadline time.Time
		err = span.ReadInvocation(ctx, txn, invID, map[string]interface{}{
			"UpdateToken": &updateToken,
			"State":       &state,
			"Deadline":    &deadline,
		})

		switch {
		case err != nil:
			return err

		case state != pb.Invocation_ACTIVE:
			return errors.Reason("%q is not active", pbutil.InvocationName(invID)).Tag(grpcutil.FailedPreconditionTag).Err()

		case deadline.Before(now):
			retErr = errors.Reason("%q is not active", pbutil.InvocationName(invID)).Tag(grpcutil.FailedPreconditionTag).Err()

			// The invocation has exceeded deadline, finalize it now.
			return finalizeInvocation(txn, invID, true, pbutil.MustTimestampProto(deadline))
		}

		if err = validateUserUpdateToken(updateToken, userToken); err != nil {
			return err
		}

		return f(ctx, txn)
	})

	if err != nil {
		retErr = err
	}
	return retErr
}

func extractUserUpdateToken(ctx context.Context) (string, error) {
	md, _ := metadata.FromIncomingContext(ctx)
	userToken := md.Get(updateTokenMetadataKey)
	switch {
	case len(userToken) == 0:
		return "", errors.
			Reason("missing %q metadata value in the request", updateTokenMetadataKey).
			Tag(grpcutil.UnauthenticatedTag).
			Err()

	case len(userToken) > 1:
		return "", errors.
			Reason("expected exactly one %q metadata value, got %d", updateTokenMetadataKey, len(userToken)).
			Tag(grpcutil.InvalidArgumentTag).
			Err()

	default:
		return userToken[0], nil
	}
}

func finalizeInvocation(txn *spanner.ReadWriteTransaction, invID string, interrupted bool, finalizeTime *tspb.Timestamp) error {
	state := pb.Invocation_COMPLETED
	if interrupted {
		state = pb.Invocation_INTERRUPTED
	}

	return txn.BufferWrite([]*spanner.Mutation{
		spanner.UpdateMap("Invocations", span.ToSpannerMap(map[string]interface{}{
			"InvocationId": invID,
			"State":        state,
			"FinalizeTime": finalizeTime,
		})),
	})
}

func validateUserUpdateToken(updateToken spanner.NullString, userToken string) error {
	if !updateToken.Valid {
		return errors.Reason("no update token in active invocation").Tag(grpcutil.InternalTag).Err()
	}

	if userToken != updateToken.StringVal {
		return errors.Reason("invalid update token").Tag(grpcutil.PermissionDeniedTag).Err()
	}

	return nil
}

func readInvocationState(ctx context.Context, txn span.Txn, invID string) (pb.Invocation_State, error) {
	var state pb.Invocation_State
	err := span.ReadInvocation(ctx, txn, invID, map[string]interface{}{"State": &state})
	return state, err
}

func insertInvocationsByTag(invID string, inv *pb.Invocation) []*spanner.Mutation {
	muts := make([]*spanner.Mutation, len(inv.Tags))
	for i, tag := range inv.Tags {
		muts[i] = spanner.InsertMap("InvocationsByTag", map[string]interface{}{
			"TagId":        span.TagID(tag),
			"InvocationId": invID,
		})
	}
	return muts
}

// insertInvocation returns an spanner mutation that inserts an Invocation row.
// Uses the value of clock.Now(ctx) to compute expiration times.
// Assumes inv is complete and valid; may panic otherwise.
func insertInvocation(ctx context.Context, inv *pb.Invocation, updateToken, createRequestID string) *spanner.Mutation {
	row := map[string]interface{}{
		"InvocationId": pbutil.MustParseInvocationName(inv.Name),
		"State":        inv.State,
		"Realm":        chromium.Realm, // TODO(crbug.com/1013316): accept realm in the proto

		"UpdateToken": updateToken,

		"CreateTime": inv.CreateTime,
		"Deadline":   inv.Deadline,

		"BaseTestVariantDef": inv.BaseTestVariantDef,
		"Tags":               inv.Tags,
	}

	if inv.FinalizeTime != nil {
		row["FinalizeTime"] = inv.FinalizeTime
	}

	if createRequestID != "" {
		row["CreateRequestId"] = createRequestID
	}

	populateExpirations(row, clock.Now(ctx))

	return spanner.InsertMap("Invocations", span.ToSpannerMap(row))
}

// populateExpirations populates the invocation row's expiration fields using the given current time.
func populateExpirations(invRow map[string]interface{}, now time.Time) {
	invExp := now.Add(invocationExpirationDuration)
	invRow["InvocationExpirationTime"] = invExp
	invRow["InvocationExpirationWeek"] = invExp.Truncate(week)

	resultsExp := now.Add(expectedTestResultsExpirationDuration)
	invRow["ExpectedTestResultsExpirationTime"] = resultsExp
	invRow["ExpectedTestResultsExpirationWeek"] = resultsExp.Truncate(week)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/xml"
	"fmt"
	id3 "github.com/mikkyang/id3-go"
	kingpin "gopkg.in/alecthomas/kingpin.v1"
	"net/url"
	"os"
	"path/filepath"
)

const (
	Header = `<?xml version="1.0" encoding="UTF-8"?>` + "\n"
	iTunesNs = "http://www.itunes.com/dtds/podcast-1.0.dtd"
)

type Rss struct {
	XMLName xml.Name `xml:"rss"`
	Channel Channel  `xml:"channel"`
	Version string   `xml:"version,attr"`
	NS      string   `xml:"xmlns:itunes,attr"`
}

type Channel struct {
	Title       string `xml:"title"`
	Link        string `xml:"link"`
	Description string `xml:"description"`
	Items       []Item `xml:"item"`
}

type Item struct {
	Title       string    `xml:"title"`
	Description string    `xml:"description"`
	Enclosure   Enclosure `xml:"enclosure"`
	Guid        string    `xml:"guid"`
	Subtitle    string    `xml:"itunes:subtitle"`
}

type Enclosure struct {
	Url    string `xml:"url,attr"`
	Length int64  `xml:"length,attr"`
	Type   string `xml:"type,attr"`
}

func fileUrl(relativePath string, baseUrl string) (string, error) {
	Url, err := url.Parse(baseUrl)
	if err != nil {
		return "", err
	}
	Url.Path += relativePath
	return Url.String(), nil
}

func addMeta(path string, f os.FileInfo, item *Item) string {
	fd, err := id3.Open(path)
	name := ""
	if err != nil {
		item.Title = f.Name()
	} else {
		title := fd.Title()
		author := fd.Artist()
		if len(title) > 0 {
			item.Title = title
		} else {
			item.Title = author
			if len(author) > 0 {
				item.Title += " - "
			}
			item.Title += f.Name()
		}
		item.Subtitle = author
	}
	return name
}

func visitFiles(workDir string, channel *Channel, publicUrl string) filepath.WalkFunc {
	return func(path string, f os.FileInfo, err error) error {

		if err != nil {
			return err
		}

		if !!f.IsDir() {
			return nil
		}

		matched, err := filepath.Match("*.mp3", f.Name())
		if err != nil {
			fmt.Println(err) // malformed pattern
			return err
		}

		if matched {
			relativePath := path[len(workDir)-1:]
			url, err := fileUrl(relativePath, publicUrl)
			if err != nil {
				return err
			}

			enclosure := Enclosure{Length: f.Size(), Type: "audio/mpeg",
				Url: url}
			item := Item{Enclosure: enclosure, Guid: enclosure.Url}
			addMeta(path, f, &item)
			channel.Items = append(channel.Items, item)
		}

		return nil

	}
}

var (
	baseUrl     = kingpin.Flag("server", "hostname (and path) to the root e.g. http://myserver.com/rss").Short('s').Default("http://localhost").String()
	title       = kingpin.Flag("title", "RSS channel title").Short('t').Default("RSS FEED").String()
	description = kingpin.Flag("description", "RSS channel description").Short('d').String()
	path        = kingpin.Arg("directory", "directory to read files relative from").Required().String()
)

func main() {

	kingpin.Version("0.0.1")
	kingpin.Parse()

	file, err := os.Open(*path)
	if err != nil {
		fmt.Printf("%s: %v\n", os.Args[0], err)
		return
	}
	stat, err := file.Stat()
	if err != nil || !stat.IsDir() {
		fmt.Printf("%s: %v: Is a file\n", os.Args[0], *path)
		return
	}

	channel := &Channel{Title: *title, Link: *baseUrl, Description: *description}
	err = filepath.Walk(*path, visitFiles(*path, channel, *baseUrl))

	if err != nil {
		fmt.Printf("%s: %v\n", os.Args[0], err)
	} else {
		output, err := xml.MarshalIndent(
			&Rss{Channel: *channel, Version: "2.0", NS: iTunesNs}, "", "  ")
		if err != nil {
			fmt.Printf("error: %v\n", err)
		} else {
			os.Stdout.WriteString(Header)
			os.Stdout.Write(output)
		}
	}

}
<commit_msg>don't return anything from addMeta<commit_after>package main

import (
	"encoding/xml"
	"fmt"
	id3 "github.com/mikkyang/id3-go"
	kingpin "gopkg.in/alecthomas/kingpin.v1"
	"net/url"
	"os"
	"path/filepath"
)

const (
	Header = `<?xml version="1.0" encoding="UTF-8"?>` + "\n"
	iTunesNs = "http://www.itunes.com/dtds/podcast-1.0.dtd"
)

type Rss struct {
	XMLName xml.Name `xml:"rss"`
	Channel Channel  `xml:"channel"`
	Version string   `xml:"version,attr"`
	NS      string   `xml:"xmlns:itunes,attr"`
}

type Channel struct {
	Title       string `xml:"title"`
	Link        string `xml:"link"`
	Description string `xml:"description"`
	Items       []Item `xml:"item"`
}

type Item struct {
	Title       string    `xml:"title"`
	Description string    `xml:"description"`
	Enclosure   Enclosure `xml:"enclosure"`
	Guid        string    `xml:"guid"`
	Subtitle    string    `xml:"itunes:subtitle"`
}

type Enclosure struct {
	Url    string `xml:"url,attr"`
	Length int64  `xml:"length,attr"`
	Type   string `xml:"type,attr"`
}

func fileUrl(relativePath string, baseUrl string) (string, error) {
	Url, err := url.Parse(baseUrl)
	if err != nil {
		return "", err
	}
	Url.Path += relativePath
	return Url.String(), nil
}

func addMeta(path string, f os.FileInfo, item *Item) {
	fd, err := id3.Open(path)
	if err != nil {
		item.Title = f.Name()
	} else {
		title := fd.Title()
		author := fd.Artist()
		if len(title) > 0 {
			item.Title = title
		} else {
			item.Title = author
			if len(author) > 0 {
				item.Title += " - "
			}
			item.Title += f.Name()
		}
		item.Subtitle = author
	}
}

func visitFiles(workDir string, channel *Channel, publicUrl string) filepath.WalkFunc {
	return func(path string, f os.FileInfo, err error) error {

		if err != nil {
			return err
		}

		if !!f.IsDir() {
			return nil
		}

		matched, err := filepath.Match("*.mp3", f.Name())
		if err != nil {
			fmt.Println(err) // malformed pattern
			return err
		}

		if matched {
			relativePath := path[len(workDir)-1:]
			url, err := fileUrl(relativePath, publicUrl)
			if err != nil {
				return err
			}

			enclosure := Enclosure{Length: f.Size(), Type: "audio/mpeg",
				Url: url}
			item := Item{Enclosure: enclosure, Guid: enclosure.Url}
			addMeta(path, f, &item)
			channel.Items = append(channel.Items, item)
		}

		return nil

	}
}

var (
	baseUrl     = kingpin.Flag("server", "hostname (and path) to the root e.g. http://myserver.com/rss").Short('s').Default("http://localhost").String()
	title       = kingpin.Flag("title", "RSS channel title").Short('t').Default("RSS FEED").String()
	description = kingpin.Flag("description", "RSS channel description").Short('d').String()
	path        = kingpin.Arg("directory", "directory to read files relative from").Required().String()
)

func main() {

	kingpin.Version("0.0.1")
	kingpin.Parse()

	file, err := os.Open(*path)
	if err != nil {
		fmt.Printf("%s: %v\n", os.Args[0], err)
		return
	}
	stat, err := file.Stat()
	if err != nil || !stat.IsDir() {
		fmt.Printf("%s: %v: Is a file\n", os.Args[0], *path)
		return
	}

	channel := &Channel{Title: *title, Link: *baseUrl, Description: *description}
	err = filepath.Walk(*path, visitFiles(*path, channel, *baseUrl))

	if err != nil {
		fmt.Printf("%s: %v\n", os.Args[0], err)
	} else {
		output, err := xml.MarshalIndent(
			&Rss{Channel: *channel, Version: "2.0", NS: iTunesNs}, "", "  ")
		if err != nil {
			fmt.Printf("error: %v\n", err)
		} else {
			os.Stdout.WriteString(Header)
			os.Stdout.Write(output)
		}
	}

}
<|endoftext|>