<commit_before>package runner

import (
	"io"
	"log"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/kylelemons/godebug/pretty"
	"github.com/pkg/errors"
	"github.com/yext/edward/home"
	"github.com/yext/edward/instance"
	"github.com/yext/edward/instance/processes"
	"github.com/yext/edward/services"
)

// Runner provides state and functions for running a given service
type Runner struct {
	Service       *services.ServiceConfig
	backendRunner services.Runner
	DirConfig     *home.EdwardConfiguration

	logFile *os.File

	commandWait sync.WaitGroup
	NoWatch     bool
	WorkingDir  string

	status instance.Status

	instanceId string

	standardLog *Log
	errorLog    *Log

	shutdownChan chan struct{}
}

func NewRunner(
	cfg services.OperationConfig,
	service *services.ServiceConfig,
	dirConfig *home.EdwardConfiguration,
	noWatch bool,
	workingDir string,
) (*Runner, error) {
	r := &Runner{
		Service:    service,
		DirConfig:  dirConfig,
		NoWatch:    noWatch,
		WorkingDir: workingDir,
	}
	var err error
	r.backendRunner, err = services.GetRunner(cfg, service)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	return r, nil
}

func (r *Runner) Run(args []string) error {
	r.updateServiceState(instance.StateStarting)

	// Allow shutdown through signals
	r.configureSignals()

	log.Printf("Signals configured")

	r.shutdownChan = make(chan struct{})

	r.status = instance.Status{
		StartTime: time.Now(),
	}

	if r.WorkingDir != "" {
		err := os.Chdir(r.WorkingDir)
		if err != nil {
			r.updateServiceState(instance.StateDied)
			return errors.WithStack(err)
		}
	}

	log.Printf("Service config: %s", pretty.Sprint(r.Service))

	// Set the instance id
	command, err := instance.Load(r.DirConfig, &processes.Processes{}, r.Service, services.ContextOverride{})
	if err != nil {
		log.Printf("Could not get service command: %v\n", err)
	}
	r.instanceId = command.InstanceId

	err = r.configureLogs()
	if err != nil {
		r.updateServiceState(instance.StateDied)
		return errors.WithStack(err)
	}

	statusTick := time.NewTicker(10 * time.Second)
	defer func() {
		if statusTick != nil {
			statusTick.Stop()
		}
	}()
	go func() {
		for _ = range statusTick.C {
			r.updateStatusDetail()
		}
	}()

	r.commandWait.Add(1)

	err = r.startService()
	if err != nil {
		return errors.WithStack(err)
	}

	r.updateStatusDetail()
	r.updateServiceState(instance.StateRunning)

	closeWatchers := r.configureWatch()
	if closeWatchers != nil {
		defer closeWatchers()
	}

	r.commandWait.Wait()

	// Wait for shutdown.
	// If the service stopped and an interrupt was not sent, do not set the "DIED" state.
	select {
	case <-r.shutdownChan:
		r.updateServiceState(instance.StateStopped)
		log.Printf("Service stopped\n")
		return nil
	default:
		r.updateServiceState(instance.StateDied)
		statusTick.Stop()
		statusTick = nil
	}

	return nil
}

func (r *Runner) updateServiceState(newState instance.State) {
	r.status.State = newState
	err := instance.SaveStatusForService(r.Service, r.instanceId, r.status, r.DirConfig.StateDir)
	if err != nil {
		log.Printf("could not save state: %v", err)
	}
}

func (r *Runner) updateStatusDetail() {
	r.status.StdoutLines = r.standardLog.Len()
	r.status.StderrLines = r.errorLog.Len()

	backendStatus, err := r.backendRunner.Status()
	if err != nil {
		log.Printf("could not save state: %v", err)
		return
	}
	r.status.Ports = backendStatus.Ports
	r.status.MemoryInfo = backendStatus.MemoryInfo

	dir := r.DirConfig.StateDir
	err = instance.SaveStatusForService(r.Service, r.instanceId, r.status, dir)
	if err != nil {
		log.Printf("could not save state: %v", err)
	}
}

func (r *Runner) configureLogs() error {
	logLocation := r.Service.GetRunLog(r.DirConfig.LogDir)
	os.Remove(logLocation)

	var err error
	r.logFile, err = os.Create(logLocation)
	if err != nil {
		return errors.WithStack(err)
	}

	// Tee the logs to stdout and the service log file
	log.SetOutput(io.MultiWriter(
		os.Stdout,
		&Log{
			file:   r.logFile,
			name:   r.Service.Name,
			stream: "messages",
		},
	))
	log.SetPrefix("Edward> ")
	return nil
}

func (r *Runner) configureSignals() {
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	go func() {
		for range signalChan {
			log.Printf("Received interrupt\n")
			err := r.stopService()
			if err != nil {
				log.Printf("Could not stop service: %v", err)
			}
			close(r.shutdownChan)
		}
	}()
}

func (r *Runner) configureWatch() func() {
	if !r.NoWatch {
		closeWatchers, err := BeginWatch(r.DirConfig, r.Service, r.restartService)
		if err != nil {
			log.Printf("Could not enable auto-restart: %v\n", err)
			return nil
		}
		if closeWatchers != nil {
			log.Printf("Auto-restart enabled. This service will restart when files in its watch directories are edited.\nThis can be disabled using the --no-watch flag.\n")
		}
		return closeWatchers
	}
	return nil
}

func (r *Runner) restartService() error {
	log.Printf("Restarting service\n")

	// Increment the counter to prevent exiting unexpectedly
	r.commandWait.Add(1)

	err := r.stopService()
	if err != nil {
		return errors.WithStack(err)
	}
	err = r.startService()
	if err != nil {
		return errors.WithStack(err)
	}
	return nil
}

func (r *Runner) stopService() error {
	wd, err := os.Getwd()
	if err != nil {
		return errors.WithStack(err)
	}

	var scriptErr error
	var scriptOutput []byte

	// TODO: Get the right env function from the client
	scriptOutput, scriptErr = r.backendRunner.Stop(wd, nil)
	if scriptErr != nil {
		log.Printf("Stop failed:\n%v\n", string(scriptOutput))
		return errors.WithStack(err)
	}
	return nil
}

func (r *Runner) startService() error {
	log.Printf("Service starting\n")

	r.standardLog = &Log{
		file:   r.logFile,
		name:   r.Service.Name,
		stream: "stdout",
	}
	r.errorLog = &Log{
		file:   r.logFile,
		name:   r.Service.Name,
		stream: "stderr",
	}

	err := r.backendRunner.Start(r.standardLog, r.errorLog)
	if err != nil {
		return errors.WithStack(err)
	}
	go func() {
		r.backendRunner.Wait()
		r.commandWait.Done()
	}()
	return nil
}
<commit_msg>Support env var replacement in stop commands<commit_after>package runner

import (
	"io"
	"log"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/kylelemons/godebug/pretty"
	"github.com/pkg/errors"
	"github.com/yext/edward/home"
	"github.com/yext/edward/instance"
	"github.com/yext/edward/instance/processes"
	"github.com/yext/edward/services"
)

// Runner provides state and functions for running a given service
type Runner struct {
	Service       *services.ServiceConfig
	backendRunner services.Runner
	DirConfig     *home.EdwardConfiguration

	logFile *os.File

	commandWait sync.WaitGroup
	NoWatch     bool
	WorkingDir  string

	status instance.Status

	instanceId string

	standardLog *Log
	errorLog    *Log

	shutdownChan chan struct{}
}

func NewRunner(
	cfg services.OperationConfig,
	service *services.ServiceConfig,
	dirConfig *home.EdwardConfiguration,
	noWatch bool,
	workingDir string,
) (*Runner, error) {
	r := &Runner{
		Service:    service,
		DirConfig:  dirConfig,
		NoWatch:    noWatch,
		WorkingDir: workingDir,
	}
	var err error
	r.backendRunner, err = services.GetRunner(cfg, service)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	return r, nil
}

func (r *Runner) Run(args []string) error {
	r.updateServiceState(instance.StateStarting)

	// Allow shutdown through signals
	r.configureSignals()

	log.Printf("Signals configured")

	r.shutdownChan = make(chan struct{})

	r.status = instance.Status{
		StartTime: time.Now(),
	}

	if r.WorkingDir != "" {
		err := os.Chdir(r.WorkingDir)
		if err != nil {
			r.updateServiceState(instance.StateDied)
			return errors.WithStack(err)
		}
	}

	log.Printf("Service config: %s", pretty.Sprint(r.Service))

	// Set the instance id
	command, err := instance.Load(r.DirConfig, &processes.Processes{}, r.Service, services.ContextOverride{})
	if err != nil {
		log.Printf("Could not get service command: %v\n", err)
	}
	r.instanceId = command.InstanceId

	err = r.configureLogs()
	if err != nil {
		r.updateServiceState(instance.StateDied)
		return errors.WithStack(err)
	}

	statusTick := time.NewTicker(10 * time.Second)
	defer func() {
		if statusTick != nil {
			statusTick.Stop()
		}
	}()
	go func() {
		for _ = range statusTick.C {
			r.updateStatusDetail()
		}
	}()

	r.commandWait.Add(1)

	err = r.startService()
	if err != nil {
		return errors.WithStack(err)
	}

	r.updateStatusDetail()
	r.updateServiceState(instance.StateRunning)

	closeWatchers := r.configureWatch()
	if closeWatchers != nil {
		defer closeWatchers()
	}

	r.commandWait.Wait()

	// Wait for shutdown.
	// If the service stopped and an interrupt was not sent, do not set the "DIED" state.
	select {
	case <-r.shutdownChan:
		r.updateServiceState(instance.StateStopped)
		log.Printf("Service stopped\n")
		return nil
	default:
		r.updateServiceState(instance.StateDied)
		statusTick.Stop()
		statusTick = nil
	}

	return nil
}

func (r *Runner) updateServiceState(newState instance.State) {
	r.status.State = newState
	err := instance.SaveStatusForService(r.Service, r.instanceId, r.status, r.DirConfig.StateDir)
	if err != nil {
		log.Printf("could not save state: %v", err)
	}
}

func (r *Runner) updateStatusDetail() {
	r.status.StdoutLines = r.standardLog.Len()
	r.status.StderrLines = r.errorLog.Len()

	backendStatus, err := r.backendRunner.Status()
	if err != nil {
		log.Printf("could not save state: %v", err)
		return
	}
	r.status.Ports = backendStatus.Ports
	r.status.MemoryInfo = backendStatus.MemoryInfo

	dir := r.DirConfig.StateDir
	err = instance.SaveStatusForService(r.Service, r.instanceId, r.status, dir)
	if err != nil {
		log.Printf("could not save state: %v", err)
	}
}

func (r *Runner) configureLogs() error {
	logLocation := r.Service.GetRunLog(r.DirConfig.LogDir)
	os.Remove(logLocation)

	var err error
	r.logFile, err = os.Create(logLocation)
	if err != nil {
		return errors.WithStack(err)
	}

	// Tee the logs to stdout and the service log file
	log.SetOutput(io.MultiWriter(
		os.Stdout,
		&Log{
			file:   r.logFile,
			name:   r.Service.Name,
			stream: "messages",
		},
	))
	log.SetPrefix("Edward> ")
	return nil
}

func (r *Runner) configureSignals() {
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	go func() {
		for range signalChan {
			log.Printf("Received interrupt\n")
			err := r.stopService()
			if err != nil {
				log.Printf("Could not stop service: %v", err)
			}
			close(r.shutdownChan)
		}
	}()
}

func (r *Runner) configureWatch() func() {
	if !r.NoWatch {
		closeWatchers, err := BeginWatch(r.DirConfig, r.Service, r.restartService)
		if err != nil {
			log.Printf("Could not enable auto-restart: %v\n", err)
			return nil
		}
		if closeWatchers != nil {
			log.Printf("Auto-restart enabled. This service will restart when files in its watch directories are edited.\nThis can be disabled using the --no-watch flag.\n")
		}
		return closeWatchers
	}
	return nil
}

func (r *Runner) restartService() error {
	log.Printf("Restarting service\n")

	// Increment the counter to prevent exiting unexpectedly
	r.commandWait.Add(1)

	err := r.stopService()
	if err != nil {
		return errors.WithStack(err)
	}
	err = r.startService()
	if err != nil {
		return errors.WithStack(err)
	}
	return nil
}

func (r *Runner) stopService() error {
	wd, err := os.Getwd()
	if err != nil {
		return errors.WithStack(err)
	}

	var scriptErr error
	var scriptOutput []byte

	c, err := instance.Load(r.DirConfig, &processes.Processes{}, r.Service, services.ContextOverride{})
	if err != nil {
		return errors.WithStack(err)
	}

	scriptOutput, scriptErr = r.backendRunner.Stop(wd, c.Getenv)
	if scriptErr != nil {
		log.Printf("Stop failed:%v\n%v\n", scriptErr, string(scriptOutput))
		return errors.WithStack(err)
	}
	return nil
}

func (r *Runner) startService() error {
	log.Printf("Service starting\n")

	r.standardLog = &Log{
		file:   r.logFile,
		name:   r.Service.Name,
		stream: "stdout",
	}
	r.errorLog = &Log{
		file:   r.logFile,
		name:   r.Service.Name,
		stream: "stderr",
	}

	err := r.backendRunner.Start(r.standardLog, r.errorLog)
	if err != nil {
		return errors.WithStack(err)
	}
	go func() {
		r.backendRunner.Wait()
		r.commandWait.Done()
	}()
	return nil
}
<|endoftext|>


<commit_before><commit_msg>Added begging of the solution for Mars Lander ep 3<commit_after><|endoftext|>


<commit_before>// Package samlsp provides helpers that can be used to protect web
// services using SAML.
package samlsp

import (
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/crewjam/saml"
)

// Options represents the parameters for creating a new middleware
type Options struct {
	URL               string
	Key               string
	Certificate       string
	AllowIDPInitiated bool
	IDPMetadata       *saml.Metadata
	IDPMetadataURL    string
}

// New creates a new Middleware
func New(opts Options) (*Middleware, error) {
	m := &Middleware{
		ServiceProvider: saml.ServiceProvider{
			Key:         opts.Key,
			Certificate: opts.Certificate,
			MetadataURL: opts.URL + "/saml/metadata",
			AcsURL:      opts.URL + "/saml/acs",
			IDPMetadata: opts.IDPMetadata,
		},
		AllowIDPInitiated: opts.AllowIDPInitiated,
	}

	// fetch the IDP metadata if needed.
	if opts.IDPMetadataURL == "" {
		return m, nil
	}

	for i := 0; true; i++ {
		resp, err := http.Get(opts.IDPMetadataURL)
		if err == nil && resp.StatusCode != http.StatusOK {
			err = fmt.Errorf("%d %s", resp.StatusCode, resp.Status)
		}
		var data []byte
		if err == nil {
			data, err = ioutil.ReadAll(resp.Body)
			resp.Body.Close()
		}
		if err != nil {
			if i > 10 {
				return nil, err
			}
			log.Printf("ERROR: %s: %s (will retry)", opts.IDPMetadataURL, err)
			time.Sleep(5 * time.Second)
			continue
		}

		entity := &saml.Metadata{}
		err = xml.Unmarshal(data, entity)
		if err == nil {
			m.ServiceProvider.IDPMetadata = entity
			return m, nil
		}
		entities := &saml.EntitiesDescriptor{}
		err = xml.Unmarshal(data, entities)
		if err != nil {
			return nil, err
		}
		for _, entity := range entities.EntityDescriptor {
			if entity.IDPSSODescriptor != nil {
				m.ServiceProvider.IDPMetadata = entity
				return m, nil
			}
		}
		return nil, fmt.Errorf("no entity returned with IDPSSODescriptor: %#v",
			entities)
	}

	panic("unreachable")
}
<commit_msg>when samlsp.New() fails to parse metadata, report the first error<commit_after>// Package samlsp provides helpers that can be used to protect web
// services using SAML.
package samlsp

import (
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/crewjam/saml"
)

// Options represents the parameters for creating a new middleware
type Options struct {
	URL               string
	Key               string
	Certificate       string
	AllowIDPInitiated bool
	IDPMetadata       *saml.Metadata
	IDPMetadataURL    string
}

// New creates a new Middleware
func New(opts Options) (*Middleware, error) {
	m := &Middleware{
		ServiceProvider: saml.ServiceProvider{
			Key:         opts.Key,
			Certificate: opts.Certificate,
			MetadataURL: opts.URL + "/saml/metadata",
			AcsURL:      opts.URL + "/saml/acs",
			IDPMetadata: opts.IDPMetadata,
		},
		AllowIDPInitiated: opts.AllowIDPInitiated,
	}

	// fetch the IDP metadata if needed.
	if opts.IDPMetadataURL == "" {
		return m, nil
	}

	for i := 0; true; i++ {
		resp, err := http.Get(opts.IDPMetadataURL)
		if err == nil && resp.StatusCode != http.StatusOK {
			err = fmt.Errorf("%d %s", resp.StatusCode, resp.Status)
		}
		var data []byte
		if err == nil {
			data, err = ioutil.ReadAll(resp.Body)
			resp.Body.Close()
		}
		if err != nil {
			if i > 10 {
				return nil, err
			}
			log.Printf("ERROR: %s: %s (will retry)", opts.IDPMetadataURL, err)
			time.Sleep(5 * time.Second)
			continue
		}

		entity := &saml.Metadata{}
		firstError := xml.Unmarshal(data, entity)
		if firstError == nil {
			m.ServiceProvider.IDPMetadata = entity
			return m, nil
		}
		entities := &saml.EntitiesDescriptor{}
		secondError := xml.Unmarshal(data, entities)
		if secondError != nil {
			if firstError != nil {
				return nil, firstError
			}
			return nil, secondError
		}
		for _, entity := range entities.EntityDescriptor {
			if entity.IDPSSODescriptor != nil {
				m.ServiceProvider.IDPMetadata = entity
				return m, nil
			}
		}
		return nil, fmt.Errorf("no entity returned with IDPSSODescriptor: %#v",
			entities)
	}

	panic("unreachable")
}
<|endoftext|>


<commit_before>package main

import (
	"time"
	"github.com/galaktor/gogre3d"
)

func main() {
	root := gogre3d.NewRoot("", "", "ogre.log")
	// setup OpenGL
        root.LoadPlugin("RenderSystem_GL")
        rs := root.GetRenderSystemByName("OpenGL Rendering Subsystem")
        rs.SetConfigOption("Full Screen", "No")
        rs.SetConfigOption("VSync", "No")
        rs.SetConfigOption("Video Mode", "800 x 600 @ 32-bit")
        root.SetRenderSystem(rs)

	window := root.Initialise(true, "gogre3d sample")
	window.Reposition(0, 0)
	sm := root.CreateSceneManager("DefaultSceneManager", "The SceneManager")
	cam := sm.CreateCamera("MyCamera")
	cam.SetPosition(0, 0, 80)
	cam.LookAt(0, 0, -300)
	cam.SetNearClipDistance(5)

	vp := window.AddViewport(cam)
	vp.SetBackgroundColour(0.01, 0, 0, 0)

	cam.SetAspectRatio(vp.GetActualWidth(), vp.GetActualHeight())

	gogre3d.SetDefaultNumMipmaps(5)

	rgm := gogre3d.GetResourceGroupManager()
	rgm.AddResourceLocation("./media/fonts", "FileSystem", "Default")
	rgm.AddResourceLocation("./media/models", "FileSystem", "Default")
	rgm.AddResourceLocation("./media/materials/scripts", "FileSystem", "Default")
	rgm.AddResourceLocation("./media/materials/programs", "FileSystem", "Default")
	rgm.AddResourceLocation("./media/materials/textures", "FileSystem", "Default")
	rgm.InitialiseAllResourceGroups()

	head := sm.CreateEntity("Head", "ogrehead.mesh", "Default")
	headnode := sm.CreateChildSceneNode("Head")
	headnode.Attach(head)

	sm.SetAmbientLight(0.5, 0.5, 0.5)

	light := sm.CreateLight("MyLight")
	light.SetPosition(20, 80, 50)

	// OIS tests
	i := gogre3d.NewInput(window)
	kb := i.NewKeyboard(false)

	ticker := time.Tick(40 * time.Millisecond)

	running := true
	for running {
		gogre3d.MessagePump()
		if window.IsClosed() {
			running = false
			break
		}

		kb.Capture()
		switch {
		case kb.KeyDown(gogre3d.KC_LEFT):
			headnode.Yaw(-0.1, gogre3d.TS_LOCAL)
		case kb.KeyDown(gogre3d.KC_RIGHT):
			headnode.Yaw(0.1, gogre3d.TS_LOCAL)
		case kb.KeyDown(gogre3d.KC_UP):
			headnode.Pitch(-0.1, gogre3d.TS_LOCAL)
		case kb.KeyDown(gogre3d.KC_DOWN):
			headnode.Pitch(0.1, gogre3d.TS_LOCAL)
		case kb.KeyDown(gogre3d.KC_ESCAPE):
			running = false
			break
		}

		if error := root.RenderOneFrame(); error == false {
			running = false
			break
		}

		<-ticker
	}

	root.Destroy()
	return
}<commit_msg>allow usage of several keys at once in sample<commit_after>package main

import (
	"time"
	"github.com/galaktor/gogre3d"
)

func main() {
	root := gogre3d.NewRoot("", "", "ogre.log")
	// setup OpenGL
        root.LoadPlugin("RenderSystem_GL")
        rs := root.GetRenderSystemByName("OpenGL Rendering Subsystem")
        rs.SetConfigOption("Full Screen", "No")
        rs.SetConfigOption("VSync", "No")
        rs.SetConfigOption("Video Mode", "800 x 600 @ 32-bit")
        root.SetRenderSystem(rs)

	window := root.Initialise(true, "gogre3d sample")
	window.Reposition(0, 0)
	sm := root.CreateSceneManager("DefaultSceneManager", "The SceneManager")
	cam := sm.CreateCamera("MyCamera")
	cam.SetPosition(0, 0, 80)
	cam.LookAt(0, 0, -300)
	cam.SetNearClipDistance(5)

	vp := window.AddViewport(cam)
	vp.SetBackgroundColour(0.01, 0, 0, 0)

	cam.SetAspectRatio(vp.GetActualWidth(), vp.GetActualHeight())

	gogre3d.SetDefaultNumMipmaps(5)

	rgm := gogre3d.GetResourceGroupManager()
	rgm.AddResourceLocation("./media/fonts", "FileSystem", "Default")
	rgm.AddResourceLocation("./media/models", "FileSystem", "Default")
	rgm.AddResourceLocation("./media/materials/scripts", "FileSystem", "Default")
	rgm.AddResourceLocation("./media/materials/programs", "FileSystem", "Default")
	rgm.AddResourceLocation("./media/materials/textures", "FileSystem", "Default")
	rgm.InitialiseAllResourceGroups()

	head := sm.CreateEntity("Head", "ogrehead.mesh", "Default")
	headnode := sm.CreateChildSceneNode("Head")
	headnode.Attach(head)

	sm.SetAmbientLight(0.5, 0.5, 0.5)

	light := sm.CreateLight("MyLight")
	light.SetPosition(20, 80, 50)

	// OIS tests
	i := gogre3d.NewInput(window)
	kb := i.NewKeyboard(false)

	ticker := time.Tick(40 * time.Millisecond)

	running := true
	for running {
		gogre3d.MessagePump()
		if window.IsClosed() {
			running = false
			break
		}

		kb.Capture()
		if kb.KeyDown(gogre3d.KC_LEFT){
			headnode.Yaw(-0.1, gogre3d.TS_LOCAL)
		}
		if kb.KeyDown(gogre3d.KC_RIGHT) {
			headnode.Yaw(0.1, gogre3d.TS_LOCAL)
		}
		if kb.KeyDown(gogre3d.KC_UP) {
			headnode.Pitch(-0.1, gogre3d.TS_LOCAL)
		}
		if kb.KeyDown(gogre3d.KC_DOWN) {
			headnode.Pitch(0.1, gogre3d.TS_LOCAL)
		}
		if kb.KeyDown(gogre3d.KC_ESCAPE) {
			running = false
			break
		}

		if error := root.RenderOneFrame(); error == false {
			running = false
			break
		}

		<-ticker
	}

	root.Destroy()
	return
}<|endoftext|>


<commit_before>/*
Copyright 2014 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.
*/

// A small utility program to lookup hostnames of endpoints in a service.
package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"sort"
	"strings"
	"time"

	"k8s.io/apimachinery/pkg/util/sets"
)

const (
	pollPeriod = 1 * time.Second
)

var (
	onChange  = flag.String("on-change", "", "Script to run on change, must accept a new line separated list of peers via stdin.")
	onStart   = flag.String("on-start", "", "Script to run on start, must accept a new line separated list of peers via stdin.")
	svc       = flag.String("service", "", "Governing service responsible for the DNS records of the domain this pod is in.")
	namespace = flag.String("ns", "", "The namespace this pod is running in. If unspecified, the POD_NAMESPACE env var is used.")
	domain    = flag.String("domain", "cluster.local", "The Cluster Domain which is used by the Cluster.")
)

func lookup(svcName string) (sets.String, error) {
	endpoints := sets.NewString()
	_, srvRecords, err := net.LookupSRV("", "", svcName)
	if err != nil {
		return endpoints, err
	}
	for _, srvRecord := range srvRecords {
		// The SRV records ends in a "." for the root domain
		ep := fmt.Sprintf("%v", srvRecord.Target[:len(srvRecord.Target)-1])
		endpoints.Insert(ep)
	}
	return endpoints, nil
}

func shellOut(sendStdin, script string) {
	log.Printf("execing: %v with stdin: %v", script, sendStdin)
	// TODO: Switch to sending stdin from go
	out, err := exec.Command("bash", "-c", fmt.Sprintf("echo -e '%v' | %v", sendStdin, script)).CombinedOutput()
	if err != nil {
		log.Fatalf("Failed to execute %v: %v, err: %v", script, string(out), err)
	}
	log.Print(string(out))
}

func main() {
	flag.Parse()

	ns := *namespace
	if ns == "" {
		ns = os.Getenv("POD_NAMESPACE")
	}
	if *svc == "" || ns == "" || (*onChange == "" && *onStart == "") {
		log.Fatalf("Incomplete args, require -on-change and/or -on-start, -service and -ns or an env var for POD_NAMESPACE.")
	}

	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalf("Failed to get hostname: %s", err)
	}

	svcLocalSuffix := strings.Join([]string{"svc", *domain}, ".")
	myName := strings.Join([]string{hostname, *svc, ns, svcLocalSuffix}, ".")
	script := *onStart
	if script == "" {
		script = *onChange
		log.Printf("No on-start supplied, on-change %v will be applied on start.", script)
	}
	for newPeers, peers := sets.NewString(), sets.NewString(); script != ""; time.Sleep(pollPeriod) {
		newPeers, err = lookup(*svc)
		if err != nil {
			log.Printf("%v", err)
			continue
		}
		if newPeers.Equal(peers) || !newPeers.Has(myName) {
			continue
		}
		peerList := newPeers.List()
		sort.Strings(peerList)
		log.Printf("Peer list updated\nwas %v\nnow %v", peers.List(), newPeers.List())
		shellOut(strings.Join(peerList, "\n"), script)
		peers = newPeers
		script = *onChange
	}
	// TODO: Exit if there's no on-change?
	log.Printf("Peer finder exiting")
}
<commit_msg>Improve domain default, add messaging when in loop to help anyone debug a seemingly infinite loop<commit_after>/*
Copyright 2014 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.
*/

// A small utility program to lookup hostnames of endpoints in a service.
package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"sort"
	"strings"
	"time"

	"k8s.io/apimachinery/pkg/util/sets"
	"io/ioutil"
)

const (
	pollPeriod = 1 * time.Second
)

var (
	onChange  = flag.String("on-change", "", "Script to run on change, must accept a new line separated list of peers via stdin.")
	onStart   = flag.String("on-start", "", "Script to run on start, must accept a new line separated list of peers via stdin.")
	svc       = flag.String("service", "", "Governing service responsible for the DNS records of the domain this pod is in.")
	namespace = flag.String("ns", "", "The namespace this pod is running in. If unspecified, the POD_NAMESPACE env var is used.")
	domain    = flag.String("domain", "", "The Cluster Domain which is used by the Cluster.")
)

func lookup(svcName string) (sets.String, error) {
	endpoints := sets.NewString()
	_, srvRecords, err := net.LookupSRV("", "", svcName)
	if err != nil {
		return endpoints, err
	}
	for _, srvRecord := range srvRecords {
		// The SRV records ends in a "." for the root domain
		ep := fmt.Sprintf("%v", srvRecord.Target[:len(srvRecord.Target)-1])
		endpoints.Insert(ep)
	}
	return endpoints, nil
}

func shellOut(sendStdin, script string) {
	log.Printf("execing: %v with stdin: %v", script, sendStdin)
	// TODO: Switch to sending stdin from go
	out, err := exec.Command("bash", "-c", fmt.Sprintf("echo -e '%v' | %v", sendStdin, script)).CombinedOutput()
	if err != nil {
		log.Fatalf("Failed to execute %v: %v, err: %v", script, string(out), err)
	}
	log.Print(string(out))
}

func main() {
	flag.Parse()

	ns := *namespace
	if ns == "" {
		ns = os.Getenv("POD_NAMESPACE")
	}
	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalf("Failed to get hostname: %s", err)
	}
	var domainName string

	// If domain is not provided, try to get it from resolv.conf
	if (*domain == "") || (ns == "") {
		resolvConfBytes, err := ioutil.ReadFile("/etc/resolv.conf")
		resolvConf := string(resolvConfBytes)
		if err != nil {
			log.Fatal("Unable to read /etc/resolv.conf")
		}
		domainName = strings.Split(strings.Split(resolvConf, ("search " + ns + "."))[1], " ")[0]
		if ns != "" {
			domainName = strings.Join([]string{ns, domainName}, ".")
		}
	} else {
		domainName = strings.Join([]string{ns, "svc", *domain}, ".")
	}

	if *svc == "" || domainName == "" || (*onChange == "" && *onStart == "") {
		log.Fatalf("Incomplete args, require -on-change and/or -on-start, -service and -ns or an env var for POD_NAMESPACE.")
	}

	myName := strings.Join([]string{hostname, *svc, domainName}, ".")
	script := *onStart
	if script == "" {
		script = *onChange
		log.Printf("No on-start supplied, on-change %v will be applied on start.", script)
	}
	for newPeers, peers := sets.NewString(), sets.NewString(); script != ""; time.Sleep(pollPeriod) {
		newPeers, err = lookup(*svc)
		if err != nil {
			log.Printf("%v", err)
			continue
		}
		if newPeers.Equal(peers) || !newPeers.Has(myName) {
			log.Printf( "Have not found myself in list yet.\nMy Hostname: %s\nHosts in list: %s", myName, strings.Join(newPeers.List(), ", "))
			continue
		}
		peerList := newPeers.List()
		sort.Strings(peerList)
		log.Printf("Peer list updated\nwas %v\nnow %v", peers.List(), newPeers.List())
		shellOut(strings.Join(peerList, "\n"), script)
		peers = newPeers
		script = *onChange
	}
	// TODO: Exit if there's no on-change?
	log.Printf("Peer finder exiting")
}
<|endoftext|>


<commit_before>// Copyright 2017 Istio 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 model

import (
	"fmt"
	"strings"
	"testing"

	networking "istio.io/api/networking/v1alpha3"
)

var (
	port8000 = []*Port{
		{
			Name: "port1",
			Port: 8000,
		},
	}

	port9000 = []*Port{
		{
			Name: "port1",
			Port: 9000,
		},
	}

	configs1 = &Config{
		ConfigMeta: ConfigMeta{
			Name:      "foo",
			Namespace: "not-default",
		},
		Spec: &networking.Sidecar{
			Egress: []*networking.IstioEgressListener{
				{
					Port: &networking.Port{
						Number:   9000,
						Protocol: "HTTP",
						Name:     "uds",
					},
					Bind:  "1.1.1.1",
					Hosts: []string{"*/*"},
				},
				{
					Hosts: []string{"*/*"},
				},
			},
		},
	}

	configs2 = &Config{
		ConfigMeta: ConfigMeta{
			Name:      "foo",
			Namespace: "not-default",
		},
		Spec: &networking.Sidecar{},
	}

	configs3 = &Config{
		ConfigMeta: ConfigMeta{
			Name:      "foo",
			Namespace: "not-default",
		},
		Spec: &networking.Sidecar{
			Egress: []*networking.IstioEgressListener{
				{
					Hosts: []string{"foo/bar", "*/*"}, // MBMBMB not valid; one no-port egress
				},
			},
		},
	}

	services1 = []*Service{
		{Hostname: "bar"},
	}

	services2 = []*Service{
		{
			Hostname: "bar",
			Ports:    port8000,
		},
		{Hostname: "barprime"},
	}

	services3 = []*Service{
		{
			Hostname: "bar",
			Ports:    port9000,
		},
		{Hostname: "barprime"},
	}

	services4 = []*Service{
		{Hostname: "bar"},
		{Hostname: "barprime"},
	}
)

func TestCreateSidecarScope(t *testing.T) {
	tests := []struct {
		name          string
		sidecarConfig *Config
		// list of available service for a given proxy
		services []*Service
		// list of services expected to be in the listener
		excpectedServices []string
	}{
		{
			"no-sidecar-config",
			nil,
			nil,
			nil,
		},
		{
			"no-sidecar-config-with-service",
			nil,
			services1,
			[]string{"bar"},
		},
		{
			"sidecar-with-multiple-egress",
			configs1,
			nil,
			nil,
		},
		{
			"sidecar-with-multiple-egress-with-service",
			configs1,
			services1,
			[]string{"bar"},
		},
		{
			"sidecar-with-multiple-egress-with-service-on-same-port",
			configs1,
			services3,
			[]string{"bar", "barprime"},
		},
		{
			"sidecar-with-multiple-egress-with-multiple-service",
			configs1,
			services4,
			[]string{"bar", "barprime"},
		},
		{
			"sidecar-with-zero-egress",
			configs2,
			nil,
			nil,
		},
		{
			"sidecar-with-zero-egress-multiple-service",
			configs2,
			services4,
			nil,
		},
		{
			"sidecar-with-multiple-egress-noport",
			configs3,
			nil,
			nil,
		},
		{
			"sidecar-with-multiple-egress-noport-with-specific-service",
			configs3,
			services2,
			[]string{"bar", "barprime"},
		},
		{
			"sidecar-with-multiple-egress-noport-with-services",
			configs3,
			services4,
			[]string{"bar", "barprime"},
		},
	}

	for idx, tt := range tests {
		t.Run(fmt.Sprintf("[%d] %s", idx, tt.name), func(t *testing.T) {
			var found bool
			ps := NewPushContext()
			meshConfig := DefaultMeshConfig()
			ps.Env = &Environment{
				Mesh: &meshConfig,
			}
			if tt.services != nil {
				ps.publicServices = append(ps.publicServices, tt.services...)
			}

			sidecarConfig := tt.sidecarConfig
			sidecarScope := ConvertToSidecarScope(ps, sidecarConfig, "mynamespace")
			configuredListeneres := 1
			if sidecarConfig != nil {
				r := sidecarConfig.Spec.(*networking.Sidecar)
				configuredListeneres = len(r.Egress)
			}

			numberListeners := len(sidecarScope.EgressListeners)
			if numberListeners != configuredListeneres {
				t.Errorf("Expected %d listeners, Got: %d", configuredListeneres, numberListeners)
			}

			if sidecarConfig != nil {
				a := sidecarConfig.Spec.(*networking.Sidecar)
				for _, egress := range a.Egress {
					for _, host := range egress.Hosts {
						parts := strings.SplitN(host, "/", 2)
						found = false
						for _, listeners := range sidecarScope.EgressListeners {
							if sidecarScopeHosts, ok := listeners.listenerHosts[parts[0]]; ok {
								for _, sidecarScopeHost := range sidecarScopeHosts {
									if sidecarScopeHost == Hostname(parts[1]) &&
										listeners.IstioListener.Port == egress.Port {
										found = true
										break
									}
								}
							}
							if found {
								break
							}
						}
						if !found {
							t.Errorf("Did not find %v entry in any listener", host)
						}
					}
				}
			}

			for _, s1 := range sidecarScope.services {
				found = false
				for _, s2 := range tt.excpectedServices {
					if string(s1.Hostname) == s2 {
						found = true
						break
					}
				}
				if !found {
					t.Errorf("Unexpected service %v found in SidecarScope", s1.Hostname)
				}
			}

			for _, s1 := range tt.excpectedServices {
				found = false
				for _, s2 := range sidecarScope.services {
					if s1 == string(s2.Hostname) {
						found = true
						break
					}
				}
				if !found {
					t.Errorf("Expected service %v in SidecarScope, but did not find it", s1)
				}
			}
			// TODO destination rule
		})
	}
}
<commit_msg>cleans up unnecessary left over comment (#13137)<commit_after>// Copyright 2017 Istio 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 model

import (
	"fmt"
	"strings"
	"testing"

	networking "istio.io/api/networking/v1alpha3"
)

var (
	port8000 = []*Port{
		{
			Name: "port1",
			Port: 8000,
		},
	}

	port9000 = []*Port{
		{
			Name: "port1",
			Port: 9000,
		},
	}

	configs1 = &Config{
		ConfigMeta: ConfigMeta{
			Name:      "foo",
			Namespace: "not-default",
		},
		Spec: &networking.Sidecar{
			Egress: []*networking.IstioEgressListener{
				{
					Port: &networking.Port{
						Number:   9000,
						Protocol: "HTTP",
						Name:     "uds",
					},
					Bind:  "1.1.1.1",
					Hosts: []string{"*/*"},
				},
				{
					Hosts: []string{"*/*"},
				},
			},
		},
	}

	configs2 = &Config{
		ConfigMeta: ConfigMeta{
			Name:      "foo",
			Namespace: "not-default",
		},
		Spec: &networking.Sidecar{},
	}

	configs3 = &Config{
		ConfigMeta: ConfigMeta{
			Name:      "foo",
			Namespace: "not-default",
		},
		Spec: &networking.Sidecar{
			Egress: []*networking.IstioEgressListener{
				{
					Hosts: []string{"foo/bar", "*/*"},
				},
			},
		},
	}

	services1 = []*Service{
		{Hostname: "bar"},
	}

	services2 = []*Service{
		{
			Hostname: "bar",
			Ports:    port8000,
		},
		{Hostname: "barprime"},
	}

	services3 = []*Service{
		{
			Hostname: "bar",
			Ports:    port9000,
		},
		{Hostname: "barprime"},
	}

	services4 = []*Service{
		{Hostname: "bar"},
		{Hostname: "barprime"},
	}
)

func TestCreateSidecarScope(t *testing.T) {
	tests := []struct {
		name          string
		sidecarConfig *Config
		// list of available service for a given proxy
		services []*Service
		// list of services expected to be in the listener
		excpectedServices []string
	}{
		{
			"no-sidecar-config",
			nil,
			nil,
			nil,
		},
		{
			"no-sidecar-config-with-service",
			nil,
			services1,
			[]string{"bar"},
		},
		{
			"sidecar-with-multiple-egress",
			configs1,
			nil,
			nil,
		},
		{
			"sidecar-with-multiple-egress-with-service",
			configs1,
			services1,
			[]string{"bar"},
		},
		{
			"sidecar-with-multiple-egress-with-service-on-same-port",
			configs1,
			services3,
			[]string{"bar", "barprime"},
		},
		{
			"sidecar-with-multiple-egress-with-multiple-service",
			configs1,
			services4,
			[]string{"bar", "barprime"},
		},
		{
			"sidecar-with-zero-egress",
			configs2,
			nil,
			nil,
		},
		{
			"sidecar-with-zero-egress-multiple-service",
			configs2,
			services4,
			nil,
		},
		{
			"sidecar-with-multiple-egress-noport",
			configs3,
			nil,
			nil,
		},
		{
			"sidecar-with-multiple-egress-noport-with-specific-service",
			configs3,
			services2,
			[]string{"bar", "barprime"},
		},
		{
			"sidecar-with-multiple-egress-noport-with-services",
			configs3,
			services4,
			[]string{"bar", "barprime"},
		},
	}

	for idx, tt := range tests {
		t.Run(fmt.Sprintf("[%d] %s", idx, tt.name), func(t *testing.T) {
			var found bool
			ps := NewPushContext()
			meshConfig := DefaultMeshConfig()
			ps.Env = &Environment{
				Mesh: &meshConfig,
			}
			if tt.services != nil {
				ps.publicServices = append(ps.publicServices, tt.services...)
			}

			sidecarConfig := tt.sidecarConfig
			sidecarScope := ConvertToSidecarScope(ps, sidecarConfig, "mynamespace")
			configuredListeneres := 1
			if sidecarConfig != nil {
				r := sidecarConfig.Spec.(*networking.Sidecar)
				configuredListeneres = len(r.Egress)
			}

			numberListeners := len(sidecarScope.EgressListeners)
			if numberListeners != configuredListeneres {
				t.Errorf("Expected %d listeners, Got: %d", configuredListeneres, numberListeners)
			}

			if sidecarConfig != nil {
				a := sidecarConfig.Spec.(*networking.Sidecar)
				for _, egress := range a.Egress {
					for _, host := range egress.Hosts {
						parts := strings.SplitN(host, "/", 2)
						found = false
						for _, listeners := range sidecarScope.EgressListeners {
							if sidecarScopeHosts, ok := listeners.listenerHosts[parts[0]]; ok {
								for _, sidecarScopeHost := range sidecarScopeHosts {
									if sidecarScopeHost == Hostname(parts[1]) &&
										listeners.IstioListener.Port == egress.Port {
										found = true
										break
									}
								}
							}
							if found {
								break
							}
						}
						if !found {
							t.Errorf("Did not find %v entry in any listener", host)
						}
					}
				}
			}

			for _, s1 := range sidecarScope.services {
				found = false
				for _, s2 := range tt.excpectedServices {
					if string(s1.Hostname) == s2 {
						found = true
						break
					}
				}
				if !found {
					t.Errorf("Unexpected service %v found in SidecarScope", s1.Hostname)
				}
			}

			for _, s1 := range tt.excpectedServices {
				found = false
				for _, s2 := range sidecarScope.services {
					if s1 == string(s2.Hostname) {
						found = true
						break
					}
				}
				if !found {
					t.Errorf("Expected service %v in SidecarScope, but did not find it", s1)
				}
			}
			// TODO destination rule
		})
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 sshutil

import (
	"bytes"
	"fmt"
	"io"
	"path/filepath"

	"github.com/docker/machine/libmachine/drivers"
	machinessh "github.com/docker/machine/libmachine/ssh"
	"golang.org/x/crypto/ssh"
)

// SSHSession provides methods for running commands on a host.
type SSHSession interface {
	Close() error
	StdinPipe() (io.WriteCloser, error)
	Run(cmd string) error
	Wait() error
}

// NewSSHClient returns an SSH client object for running commands.
func NewSSHClient(d drivers.Driver) (*ssh.Client, error) {
	h, err := newSSHHost(d)
	if err != nil {
		return nil, err

	}
	auth := &machinessh.Auth{}
	if h.SSHKeyPath != "" {
		auth.Keys = []string{h.SSHKeyPath}
	}
	config, err := machinessh.NewNativeConfig(h.Username, auth)
	if err != nil {
		return nil, err
	}

	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", h.IP, h.Port), &config)
	if err != nil {
		return nil, err
	}
	return client, nil
}

// Transfer uses an SSH session to copy a file to the remote machine.
func Transfer(data []byte, remotedir, filename string, perm string, c *ssh.Client) error {
	// Delete the old file first. This makes sure permissions get reset.
	deleteCmd := fmt.Sprintf("sudo rm -f %s", filepath.Join(remotedir, filename))
	mkdirCmd := fmt.Sprintf("sudo mkdir -p %s", remotedir)
	for _, cmd := range []string{deleteCmd, mkdirCmd} {
		if err := RunCommand(c, cmd); err != nil {
			return err
		}
	}

	s, err := c.NewSession()
	if err != nil {
		return err
	}

	go func() {
		w, err := s.StdinPipe()
		defer w.Close()
		if err != nil {
			return
		}
		header := fmt.Sprintf("C%s %d %s\n", perm, len(data), filename)
		fmt.Fprint(w, header)
		reader := bytes.NewReader(data)
		io.Copy(w, reader)
		fmt.Fprint(w, "\x00")
	}()

	scpcmd := fmt.Sprintf("sudo /usr/local/bin/scp -t %s", remotedir)
	if err := s.Run(scpcmd); err != nil {
		return err
	}

	return nil
}

func RunCommand(c *ssh.Client, cmd string) error {
	s, err := c.NewSession()
	defer s.Close()
	if err != nil {
		return err
	}

	return s.Run(cmd)
}

type sshHost struct {
	IP         string
	Port       int
	SSHKeyPath string
	Username   string
}

func newSSHHost(d drivers.Driver) (*sshHost, error) {

	ip, err := d.GetSSHHostname()
	if err != nil {
		return nil, err
	}
	port, err := d.GetSSHPort()
	if err != nil {
		return nil, err
	}
	return &sshHost{
		IP:         ip,
		Port:       port,
		SSHKeyPath: d.GetSSHKeyPath(),
		Username:   d.GetSSHUsername(),
	}, nil
}
<commit_msg>Fix a flake in sshutil.<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 sshutil

import (
	"bytes"
	"fmt"
	"io"
	"path/filepath"
	"sync"

	"github.com/docker/machine/libmachine/drivers"
	machinessh "github.com/docker/machine/libmachine/ssh"
	"golang.org/x/crypto/ssh"
)

// SSHSession provides methods for running commands on a host.
type SSHSession interface {
	Close() error
	StdinPipe() (io.WriteCloser, error)
	Run(cmd string) error
	Wait() error
}

// NewSSHClient returns an SSH client object for running commands.
func NewSSHClient(d drivers.Driver) (*ssh.Client, error) {
	h, err := newSSHHost(d)
	if err != nil {
		return nil, err

	}
	auth := &machinessh.Auth{}
	if h.SSHKeyPath != "" {
		auth.Keys = []string{h.SSHKeyPath}
	}
	config, err := machinessh.NewNativeConfig(h.Username, auth)
	if err != nil {
		return nil, err
	}

	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", h.IP, h.Port), &config)
	if err != nil {
		return nil, err
	}
	return client, nil
}

// Transfer uses an SSH session to copy a file to the remote machine.
func Transfer(data []byte, remotedir, filename string, perm string, c *ssh.Client) error {
	// Delete the old file first. This makes sure permissions get reset.
	deleteCmd := fmt.Sprintf("sudo rm -f %s", filepath.Join(remotedir, filename))
	mkdirCmd := fmt.Sprintf("sudo mkdir -p %s", remotedir)
	for _, cmd := range []string{deleteCmd, mkdirCmd} {
		if err := RunCommand(c, cmd); err != nil {
			return err
		}
	}

	s, err := c.NewSession()
	if err != nil {
		return err
	}

	w, err := s.StdinPipe()
	if err != nil {
		return err
	}
	// The scpcmd below *should not* return until all data is copied and the
	// StdinPipe is closed. But let's use a WaitGroup to make it expicit.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		defer w.Close()
		header := fmt.Sprintf("C%s %d %s\n", perm, len(data), filename)
		fmt.Fprint(w, header)
		reader := bytes.NewReader(data)
		io.Copy(w, reader)
		fmt.Fprint(w, "\x00")
	}()
	scpcmd := fmt.Sprintf("sudo /usr/local/bin/scp -t %s", remotedir)
	if err := s.Run(scpcmd); err != nil {
		return err
	}
	wg.Wait()

	return nil
}

func RunCommand(c *ssh.Client, cmd string) error {
	s, err := c.NewSession()
	defer s.Close()
	if err != nil {
		return err
	}

	return s.Run(cmd)
}

type sshHost struct {
	IP         string
	Port       int
	SSHKeyPath string
	Username   string
}

func newSSHHost(d drivers.Driver) (*sshHost, error) {

	ip, err := d.GetSSHHostname()
	if err != nil {
		return nil, err
	}
	port, err := d.GetSSHPort()
	if err != nil {
		return nil, err
	}
	return &sshHost{
		IP:         ip,
		Port:       port,
		SSHKeyPath: d.GetSSHKeyPath(),
		Username:   d.GetSSHUsername(),
	}, nil
}
<|endoftext|>


<commit_before>package rancher

import (
	"context"
	"fmt"
	"text/template"
	"time"

	"github.com/cenkalti/backoff/v3"
	"github.com/containous/traefik/v2/pkg/config/dynamic"
	"github.com/containous/traefik/v2/pkg/job"
	"github.com/containous/traefik/v2/pkg/log"
	"github.com/containous/traefik/v2/pkg/provider"
	"github.com/containous/traefik/v2/pkg/safe"
	rancher "github.com/rancher/go-rancher-metadata/metadata"
)

const (
	// DefaultTemplateRule The default template for the default rule.
	DefaultTemplateRule = "Host(`{{ normalize .Name }}`)"
)

// Health
const (
	healthy         = "healthy"
	updatingHealthy = "updating-healthy"
)

// State
const (
	active          = "active"
	running         = "running"
	upgraded        = "upgraded"
	upgrading       = "upgrading"
	updatingActive  = "updating-active"
	updatingRunning = "updating-running"
)

var _ provider.Provider = (*Provider)(nil)

// Provider holds configurations of the provider.
type Provider struct {
	Constraints               string `description:"Constraints is an expression that Traefik matches against the container's labels to determine whether to create any route for that container." json:"constraints,omitempty" toml:"constraints,omitempty" yaml:"constraints,omitempty" export:"true"`
	Watch                     bool   `description:"Watch provider." json:"watch,omitempty" toml:"watch,omitempty" yaml:"watch,omitempty" export:"true"`
	DefaultRule               string `description:"Default rule." json:"defaultRule,omitempty" toml:"defaultRule,omitempty" yaml:"defaultRule,omitempty"`
	ExposedByDefault          bool   `description:"Expose containers by default." json:"exposedByDefault,omitempty" toml:"exposedByDefault,omitempty" yaml:"exposedByDefault,omitempty" export:"true"`
	EnableServiceHealthFilter bool   `description:"Filter services with unhealthy states and inactive states." json:"enableServiceHealthFilter,omitempty" toml:"enableServiceHealthFilter,omitempty" yaml:"enableServiceHealthFilter,omitempty" export:"true"`
	RefreshSeconds            int    `description:"Defines the polling interval in seconds." json:"refreshSeconds,omitempty" toml:"refreshSeconds,omitempty" yaml:"refreshSeconds,omitempty" export:"true"`
	IntervalPoll              bool   `description:"Poll the Rancher metadata service every 'rancher.refreshseconds' (less accurate)." json:"intervalPoll,omitempty" toml:"intervalPoll,omitempty" yaml:"intervalPoll,omitempty"`
	Prefix                    string `description:"Prefix used for accessing the Rancher metadata service." json:"prefix,omitempty" toml:"prefix,omitempty" yaml:"prefix,omitempty"`
	defaultRuleTpl            *template.Template
}

// SetDefaults sets the default values.
func (p *Provider) SetDefaults() {
	p.Watch = true
	p.ExposedByDefault = true
	p.EnableServiceHealthFilter = true
	p.RefreshSeconds = 15
	p.DefaultRule = DefaultTemplateRule
	p.Prefix = "latest"
}

type rancherData struct {
	Name       string
	Labels     map[string]string
	Containers []string
	Health     string
	State      string
	Port       string
	ExtraConf  configuration
}

// Init the provider.
func (p *Provider) Init() error {
	defaultRuleTpl, err := provider.MakeDefaultRuleTemplate(p.DefaultRule, nil)
	if err != nil {
		return fmt.Errorf("error while parsing default rule: %v", err)
	}

	p.defaultRuleTpl = defaultRuleTpl
	return nil
}

func (p *Provider) createClient(ctx context.Context) (rancher.Client, error) {
	metadataServiceURL := fmt.Sprintf("http://rancher-metadata.rancher.internal/%s", p.Prefix)
	client, err := rancher.NewClientAndWait(metadataServiceURL)
	if err != nil {
		log.FromContext(ctx).Errorf("Failed to create Rancher metadata service client: %v", err)
		return nil, err
	}

	return client, nil
}

// Provide allows the rancher provider to provide configurations to traefik using the given configuration channel.
func (p *Provider) Provide(configurationChan chan<- dynamic.Message, pool *safe.Pool) error {
	pool.GoCtx(func(routineCtx context.Context) {
		ctxLog := log.With(routineCtx, log.Str(log.ProviderName, "rancher"))
		logger := log.FromContext(ctxLog)

		operation := func() error {
			client, err := p.createClient(ctxLog)
			if err != nil {
				logger.Errorf("Failed to create the metadata client metadata service: %v", err)
				return err
			}

			updateConfiguration := func(_ string) {
				stacks, err := client.GetStacks()
				if err != nil {
					logger.Errorf("Failed to query Rancher metadata service: %v", err)
					return
				}

				rancherData := p.parseMetadataSourcedRancherData(ctxLog, stacks)

				logger.Printf("Received Rancher data %+v", rancherData)

				configuration := p.buildConfiguration(ctxLog, rancherData)
				configurationChan <- dynamic.Message{
					ProviderName:  "rancher",
					Configuration: configuration,
				}
			}
			updateConfiguration("init")

			if p.Watch {
				if p.IntervalPoll {
					p.intervalPoll(ctxLog, client, updateConfiguration)
				} else {
					// Long polling should be favored for the most accurate configuration updates.
					// Holds the connection until there is either a change in the metadata repository or `p.RefreshSeconds` has elapsed.
					client.OnChangeCtx(ctxLog, p.RefreshSeconds, updateConfiguration)
				}
			}

			return nil
		}

		notify := func(err error, time time.Duration) {
			logger.Errorf("Provider connection error %+v, retrying in %s", err, time)
		}
		err := backoff.RetryNotify(safe.OperationWithRecover(operation), backoff.WithContext(job.NewBackOff(backoff.NewExponentialBackOff()), ctxLog), notify)
		if err != nil {
			logger.Errorf("Cannot connect to Provider server: %+v", err)
		}
	})

	return nil
}

func (p *Provider) intervalPoll(ctx context.Context, client rancher.Client, updateConfiguration func(string)) {
	ticker := time.NewTicker(time.Second * time.Duration(p.RefreshSeconds))
	defer ticker.Stop()

	var version string
	for {
		select {
		case <-ticker.C:
			newVersion, err := client.GetVersion()
			if err != nil {
				log.FromContext(ctx).Errorf("Failed to create Rancher metadata service client: %v", err)
			} else if version != newVersion {
				version = newVersion
				updateConfiguration(version)
			}
		case <-ctx.Done():
			return
		}
	}
}

func (p *Provider) parseMetadataSourcedRancherData(ctx context.Context, stacks []rancher.Stack) (rancherDataList []rancherData) {
	for _, stack := range stacks {
		for _, service := range stack.Services {
			ctxSvc := log.With(ctx, log.Str("stack", stack.Name), log.Str("service", service.Name))
			logger := log.FromContext(ctxSvc)

			servicePort := ""
			if len(service.Ports) > 0 {
				servicePort = service.Ports[0]
			}
			for _, port := range service.Ports {
				logger.Debugf("Set Port %s", port)
			}

			var containerIPAddresses []string
			for _, container := range service.Containers {
				if containerFilter(ctxSvc, container.Name, container.HealthState, container.State) {
					containerIPAddresses = append(containerIPAddresses, container.PrimaryIp)
				}
			}

			service := rancherData{
				Name:       service.Name + "/" + stack.Name,
				State:      service.State,
				Labels:     service.Labels,
				Port:       servicePort,
				Containers: containerIPAddresses,
			}

			extraConf, err := p.getConfiguration(service)
			if err != nil {
				logger.Errorf("Skip container %s: %v", service.Name, err)
				continue
			}

			service.ExtraConf = extraConf

			rancherDataList = append(rancherDataList, service)
		}
	}
	return rancherDataList
}

func containerFilter(ctx context.Context, name, healthState, state string) bool {
	logger := log.FromContext(ctx)

	if healthState != "" && healthState != healthy && healthState != updatingHealthy {
		logger.Debugf("Filtering container %s with healthState of %s", name, healthState)
		return false
	}

	if state != "" && state != running && state != updatingRunning && state != upgraded {
		logger.Debugf("Filtering container %s with state of %s", name, state)
		return false
	}

	return true
}
<commit_msg>Change service name in rancher provider to make webui service details view work<commit_after>package rancher

import (
	"context"
	"fmt"
	"text/template"
	"time"

	"github.com/cenkalti/backoff/v3"
	"github.com/containous/traefik/v2/pkg/config/dynamic"
	"github.com/containous/traefik/v2/pkg/job"
	"github.com/containous/traefik/v2/pkg/log"
	"github.com/containous/traefik/v2/pkg/provider"
	"github.com/containous/traefik/v2/pkg/safe"
	rancher "github.com/rancher/go-rancher-metadata/metadata"
)

const (
	// DefaultTemplateRule The default template for the default rule.
	DefaultTemplateRule = "Host(`{{ normalize .Name }}`)"
)

// Health
const (
	healthy         = "healthy"
	updatingHealthy = "updating-healthy"
)

// State
const (
	active          = "active"
	running         = "running"
	upgraded        = "upgraded"
	upgrading       = "upgrading"
	updatingActive  = "updating-active"
	updatingRunning = "updating-running"
)

var _ provider.Provider = (*Provider)(nil)

// Provider holds configurations of the provider.
type Provider struct {
	Constraints               string `description:"Constraints is an expression that Traefik matches against the container's labels to determine whether to create any route for that container." json:"constraints,omitempty" toml:"constraints,omitempty" yaml:"constraints,omitempty" export:"true"`
	Watch                     bool   `description:"Watch provider." json:"watch,omitempty" toml:"watch,omitempty" yaml:"watch,omitempty" export:"true"`
	DefaultRule               string `description:"Default rule." json:"defaultRule,omitempty" toml:"defaultRule,omitempty" yaml:"defaultRule,omitempty"`
	ExposedByDefault          bool   `description:"Expose containers by default." json:"exposedByDefault,omitempty" toml:"exposedByDefault,omitempty" yaml:"exposedByDefault,omitempty" export:"true"`
	EnableServiceHealthFilter bool   `description:"Filter services with unhealthy states and inactive states." json:"enableServiceHealthFilter,omitempty" toml:"enableServiceHealthFilter,omitempty" yaml:"enableServiceHealthFilter,omitempty" export:"true"`
	RefreshSeconds            int    `description:"Defines the polling interval in seconds." json:"refreshSeconds,omitempty" toml:"refreshSeconds,omitempty" yaml:"refreshSeconds,omitempty" export:"true"`
	IntervalPoll              bool   `description:"Poll the Rancher metadata service every 'rancher.refreshseconds' (less accurate)." json:"intervalPoll,omitempty" toml:"intervalPoll,omitempty" yaml:"intervalPoll,omitempty"`
	Prefix                    string `description:"Prefix used for accessing the Rancher metadata service." json:"prefix,omitempty" toml:"prefix,omitempty" yaml:"prefix,omitempty"`
	defaultRuleTpl            *template.Template
}

// SetDefaults sets the default values.
func (p *Provider) SetDefaults() {
	p.Watch = true
	p.ExposedByDefault = true
	p.EnableServiceHealthFilter = true
	p.RefreshSeconds = 15
	p.DefaultRule = DefaultTemplateRule
	p.Prefix = "latest"
}

type rancherData struct {
	Name       string
	Labels     map[string]string
	Containers []string
	Health     string
	State      string
	Port       string
	ExtraConf  configuration
}

// Init the provider.
func (p *Provider) Init() error {
	defaultRuleTpl, err := provider.MakeDefaultRuleTemplate(p.DefaultRule, nil)
	if err != nil {
		return fmt.Errorf("error while parsing default rule: %v", err)
	}

	p.defaultRuleTpl = defaultRuleTpl
	return nil
}

func (p *Provider) createClient(ctx context.Context) (rancher.Client, error) {
	metadataServiceURL := fmt.Sprintf("http://rancher-metadata.rancher.internal/%s", p.Prefix)
	client, err := rancher.NewClientAndWait(metadataServiceURL)
	if err != nil {
		log.FromContext(ctx).Errorf("Failed to create Rancher metadata service client: %v", err)
		return nil, err
	}

	return client, nil
}

// Provide allows the rancher provider to provide configurations to traefik using the given configuration channel.
func (p *Provider) Provide(configurationChan chan<- dynamic.Message, pool *safe.Pool) error {
	pool.GoCtx(func(routineCtx context.Context) {
		ctxLog := log.With(routineCtx, log.Str(log.ProviderName, "rancher"))
		logger := log.FromContext(ctxLog)

		operation := func() error {
			client, err := p.createClient(ctxLog)
			if err != nil {
				logger.Errorf("Failed to create the metadata client metadata service: %v", err)
				return err
			}

			updateConfiguration := func(_ string) {
				stacks, err := client.GetStacks()
				if err != nil {
					logger.Errorf("Failed to query Rancher metadata service: %v", err)
					return
				}

				rancherData := p.parseMetadataSourcedRancherData(ctxLog, stacks)

				logger.Printf("Received Rancher data %+v", rancherData)

				configuration := p.buildConfiguration(ctxLog, rancherData)
				configurationChan <- dynamic.Message{
					ProviderName:  "rancher",
					Configuration: configuration,
				}
			}
			updateConfiguration("init")

			if p.Watch {
				if p.IntervalPoll {
					p.intervalPoll(ctxLog, client, updateConfiguration)
				} else {
					// Long polling should be favored for the most accurate configuration updates.
					// Holds the connection until there is either a change in the metadata repository or `p.RefreshSeconds` has elapsed.
					client.OnChangeCtx(ctxLog, p.RefreshSeconds, updateConfiguration)
				}
			}

			return nil
		}

		notify := func(err error, time time.Duration) {
			logger.Errorf("Provider connection error %+v, retrying in %s", err, time)
		}
		err := backoff.RetryNotify(safe.OperationWithRecover(operation), backoff.WithContext(job.NewBackOff(backoff.NewExponentialBackOff()), ctxLog), notify)
		if err != nil {
			logger.Errorf("Cannot connect to Provider server: %+v", err)
		}
	})

	return nil
}

func (p *Provider) intervalPoll(ctx context.Context, client rancher.Client, updateConfiguration func(string)) {
	ticker := time.NewTicker(time.Second * time.Duration(p.RefreshSeconds))
	defer ticker.Stop()

	var version string
	for {
		select {
		case <-ticker.C:
			newVersion, err := client.GetVersion()
			if err != nil {
				log.FromContext(ctx).Errorf("Failed to create Rancher metadata service client: %v", err)
			} else if version != newVersion {
				version = newVersion
				updateConfiguration(version)
			}
		case <-ctx.Done():
			return
		}
	}
}

func (p *Provider) parseMetadataSourcedRancherData(ctx context.Context, stacks []rancher.Stack) (rancherDataList []rancherData) {
	for _, stack := range stacks {
		for _, service := range stack.Services {
			ctxSvc := log.With(ctx, log.Str("stack", stack.Name), log.Str("service", service.Name))
			logger := log.FromContext(ctxSvc)

			servicePort := ""
			if len(service.Ports) > 0 {
				servicePort = service.Ports[0]
			}
			for _, port := range service.Ports {
				logger.Debugf("Set Port %s", port)
			}

			var containerIPAddresses []string
			for _, container := range service.Containers {
				if containerFilter(ctxSvc, container.Name, container.HealthState, container.State) {
					containerIPAddresses = append(containerIPAddresses, container.PrimaryIp)
				}
			}

			service := rancherData{
				Name:       service.Name + "_" + stack.Name,
				State:      service.State,
				Labels:     service.Labels,
				Port:       servicePort,
				Containers: containerIPAddresses,
			}

			extraConf, err := p.getConfiguration(service)
			if err != nil {
				logger.Errorf("Skip container %s: %v", service.Name, err)
				continue
			}

			service.ExtraConf = extraConf

			rancherDataList = append(rancherDataList, service)
		}
	}
	return rancherDataList
}

func containerFilter(ctx context.Context, name, healthState, state string) bool {
	logger := log.FromContext(ctx)

	if healthState != "" && healthState != healthy && healthState != updatingHealthy {
		logger.Debugf("Filtering container %s with healthState of %s", name, healthState)
		return false
	}

	if state != "" && state != running && state != updatingRunning && state != upgraded {
		logger.Debugf("Filtering container %s with state of %s", name, state)
		return false
	}

	return true
}
<|endoftext|>


<commit_before>package attacks

import (
  "crypto/rsa"
  "crypto/x509"
  "encoding/pem"
  "errors"
  "fmt"
  "math/big"
  "github.com/ncw/gmp"
  "github.com/sourcekris/goRsaTool/libnum"
  "github.com/sourcekris/goRsaTool/utils"
)

type GMPPublicKey struct {
  N *gmp.Int
  E int
}

type GMPPrivateKey struct {
  PublicKey GMPPublicKey
  D *gmp.Int
  Primes []*gmp.Int
  N *gmp.Int
  E int
}

/*
 * wrap rsa.PrivateKey and add a field for cipher and plaintexts
 */
type RSAStuff struct {
  Key GMPPrivateKey
  CipherText []byte
  PlainText []byte
  PastPrimesFile string
}

/*
 * constructor for RSAStuff struct
 */
func NewRSAStuff(key *rsa.PrivateKey, c []byte, m []byte, pf string) (*RSAStuff, error) {
	if key.N == nil {
		return nil, errors.New("Key had no modulus or exponent")
	}

  var pastPrimesFile string
    if len(pf) > 0 {
      pastPrimesFile = pf
  }

  var cipherText []byte
    if len(c) > 0 {
      cipherText = c
  }

  // copy a rsa.PrivateKey to a GMPPrivateKey that uses gmp.Int types
  gmpPrivateKey := RSAtoGMPPrivateKey(key)

	// pack the RSAStuff struct
   return &RSAStuff{
      Key: *gmpPrivateKey,
      PastPrimesFile: pastPrimesFile,
      CipherText: cipherText,
    }, nil
}

/*
 * Given one prime p, pack the Key member of the RSAStuff struct with the private key values, p, q & d
 */
func (targetRSA *RSAStuff) PackGivenP(p *gmp.Int) {
  q := new(gmp.Int).Div(targetRSA.Key.N, p)
  targetRSA.Key.Primes = []*gmp.Int{p, q}
  targetRSA.Key.D      = utils.SolveforD(p, q, targetRSA.Key.E)
}

func (targetRSA *RSAStuff) DumpKey() {
  fmt.Printf("[*] n = %d\n", targetRSA.Key.N)
  fmt.Printf("[*] e = %d\n", targetRSA.Key.E)

  // XXX: Support RSA multiprime [where len(key.Primes) > 2]
  if targetRSA.Key.D!= nil {
    fmt.Printf("[*] d = %d\n", targetRSA.Key.D)
    fmt.Printf("[*] p = %d\n", targetRSA.Key.Primes[0])
    fmt.Printf("[*] q = %d\n", targetRSA.Key.Primes[1])
  }

  if len(targetRSA.CipherText) > 0 {
    fmt.Printf("[*] c = %d\n", libnum.BytesToNumber(targetRSA.CipherText))
  }
}

/*
 * Takes a rsa.PrivateKey and returns a GMPPrivateKey that uses gmp.Int types
 */
func RSAtoGMPPrivateKey(key *rsa.PrivateKey) *GMPPrivateKey {
  gmpPubKey := &GMPPublicKey{
    N: new(gmp.Int).SetBytes(key.N.Bytes()),
    E: key.E,
  }

  gmpPrivateKey := &GMPPrivateKey{
    PublicKey: *gmpPubKey,
    D: new(gmp.Int).SetBytes(key.D.Bytes()),
    Primes: []*gmp.Int{
      new(gmp.Int).SetBytes(key.Primes[0].Bytes()),
      new(gmp.Int).SetBytes(key.Primes[1].Bytes()),
      },
  }

  return gmpPrivateKey
}

func GMPtoRSAPrivateKey(key *GMPPrivateKey) *rsa.PrivateKey {
  pubKey := &rsa.PublicKey{
    N: new(big.Int).SetBytes(key.N.Bytes()),
    E: key.E,
  }

  privateKey := &rsa.PrivateKey{
    PublicKey: *pubKey,
    D: new(big.Int).SetBytes(key.D.Bytes()),
    Primes: []*big.Int{
      new(big.Int).SetBytes(key.Primes[0].Bytes()),
      new(big.Int).SetBytes(key.Primes[1].Bytes()),
      },
  }

  return privateKey
}

func encodeDerToPem(der []byte, t string) string {
  p := pem.EncodeToMemory(
    &pem.Block{
      Type: t,
      Bytes: der,
      },
      )

  return string(p)
}

func EncodePublicKey(pub *rsa.PublicKey) (string, error) {
  pubder, err := x509.MarshalPKIXPublicKey(pub)
  if err != nil {
    return "", err
  }

  return encodeDerToPem(pubder, "RSA PUBLIC KEY"), nil
}

func EncodePrivateKey(priv *rsa.PrivateKey) string {
  privder := x509.MarshalPKCS1PrivateKey(priv)
  return encodeDerToPem(privder, "RSA PRIVATE KEY")
}

func EncodeGMPPrivateKey(priv *GMPPrivateKey) string {
  privder := x509.MarshalPKCS1PrivateKey(GMPtoRSAPrivateKey(priv))
  return encodeDerToPem(privder, "RSA PRIVATE KEY")
}

<commit_msg>Added checks in RSA to GMP key conversion to avoid co-ercing nils to null pointers. Problems is now in GMP code within attacks now...<commit_after>package attacks

import (
  "crypto/rsa"
  "crypto/x509"
  "encoding/pem"
  "errors"
  "fmt"
  "math/big"
  "github.com/ncw/gmp"
  "github.com/sourcekris/goRsaTool/libnum"
  "github.com/sourcekris/goRsaTool/utils"
)

type GMPPublicKey struct {
  N *gmp.Int
  E int
}

type GMPPrivateKey struct {
  PublicKey GMPPublicKey
  D *gmp.Int
  Primes []*gmp.Int
  N *gmp.Int
  E int
}

/*
 * wrap rsa.PrivateKey and add a field for cipher and plaintexts
 */
type RSAStuff struct {
  Key GMPPrivateKey
  CipherText []byte
  PlainText []byte
  PastPrimesFile string
}

/*
 * constructor for RSAStuff struct
 */
func NewRSAStuff(key *rsa.PrivateKey, c []byte, m []byte, pf string) (*RSAStuff, error) {
	if key.N == nil {
		return nil, errors.New("Key had no modulus or exponent")
	}

  var pastPrimesFile string
    if len(pf) > 0 {
      pastPrimesFile = pf
  }

  var cipherText []byte
    if len(c) > 0 {
      cipherText = c
  }

  // copy a rsa.PrivateKey to a GMPPrivateKey that uses gmp.Int types
  gmpPrivateKey := RSAtoGMPPrivateKey(key)

	// pack the RSAStuff struct
   return &RSAStuff{
      Key: *gmpPrivateKey,
      PastPrimesFile: pastPrimesFile,
      CipherText: cipherText,
    }, nil
}

/*
 * Given one prime p, pack the Key member of the RSAStuff struct with the private key values, p, q & d
 */
func (targetRSA *RSAStuff) PackGivenP(p *gmp.Int) {
  q := new(gmp.Int).Div(targetRSA.Key.N, p)
  targetRSA.Key.Primes = []*gmp.Int{p, q}
  targetRSA.Key.D      = utils.SolveforD(p, q, targetRSA.Key.E)
}

func (targetRSA *RSAStuff) DumpKey() {
  fmt.Printf("[*] n = %d\n", targetRSA.Key.N)
  fmt.Printf("[*] e = %d\n", targetRSA.Key.E)

  // XXX: Support RSA multiprime [where len(key.Primes) > 2]
  if targetRSA.Key.D!= nil {
    fmt.Printf("[*] d = %d\n", targetRSA.Key.D)
    fmt.Printf("[*] p = %d\n", targetRSA.Key.Primes[0])
    fmt.Printf("[*] q = %d\n", targetRSA.Key.Primes[1])
  }

  if len(targetRSA.CipherText) > 0 {
    fmt.Printf("[*] c = %d\n", libnum.BytesToNumber(targetRSA.CipherText))
  }
}

/*
 * Takes a rsa.PrivateKey and returns a GMPPrivateKey that uses gmp.Int types
 */
func RSAtoGMPPrivateKey(key *rsa.PrivateKey) *GMPPrivateKey {
  gmpPubKey := &GMPPublicKey{
    N: new(gmp.Int).SetBytes(key.N.Bytes()),
    E: key.E,
  }

  var gmpPrivateKey *GMPPrivateKey
  if key.D != nil {
    gmpPrivateKey = &GMPPrivateKey{
      PublicKey: *gmpPubKey,
      D: new(gmp.Int).SetBytes(key.D.Bytes()),
      Primes: []*gmp.Int{
        new(gmp.Int).SetBytes(key.Primes[0].Bytes()),
        new(gmp.Int).SetBytes(key.Primes[1].Bytes()),
        },
    }
  } else {
    gmpPrivateKey = &GMPPrivateKey{
      PublicKey: *gmpPubKey,
    }
  }

  return gmpPrivateKey
}

func GMPtoRSAPrivateKey(key *GMPPrivateKey) *rsa.PrivateKey {
  pubKey := &rsa.PublicKey{
    N: new(big.Int).SetBytes(key.N.Bytes()),
    E: key.E,
  }

  var privateKey *rsa.PrivateKey
  if key.D != nil {
    privateKey = &rsa.PrivateKey{
      PublicKey: *pubKey,
      D: new(big.Int).SetBytes(key.D.Bytes()),
      Primes: []*big.Int{
        new(big.Int).SetBytes(key.Primes[0].Bytes()),
        new(big.Int).SetBytes(key.Primes[1].Bytes()),
      },
    }
  } else {
    privateKey = &rsa.PrivateKey{
      PublicKey: *pubKey,
    }
  }

  return privateKey
}

func encodeDerToPem(der []byte, t string) string {
  p := pem.EncodeToMemory(
    &pem.Block{
      Type: t,
      Bytes: der,
      },
      )

  return string(p)
}

func EncodePublicKey(pub *rsa.PublicKey) (string, error) {
  pubder, err := x509.MarshalPKIXPublicKey(pub)
  if err != nil {
    return "", err
  }

  return encodeDerToPem(pubder, "RSA PUBLIC KEY"), nil
}

func EncodePrivateKey(priv *rsa.PrivateKey) string {
  privder := x509.MarshalPKCS1PrivateKey(priv)
  return encodeDerToPem(privder, "RSA PRIVATE KEY")
}

func EncodeGMPPrivateKey(priv *GMPPrivateKey) string {
  privder := x509.MarshalPKCS1PrivateKey(GMPtoRSAPrivateKey(priv))
  return encodeDerToPem(privder, "RSA PRIVATE KEY")
}<|endoftext|>


<commit_before>/*
Package secret provides tools for encrypting and decrypting authenticated messages.
See docs/secret.md for more details.
*/
package secret

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"code.google.com/p/go.crypto/nacl/secretbox"
	"github.com/mailgun/lemma/random"
	"github.com/mailgun/metrics"
)

// SecretSevice is an interface for encrypting/decrypting and authenticating messages.
type SecretService interface {
	// Seal takes a plaintext message and returns an encrypted and authenticated ciphertext.
	Seal([]byte) (SealedData, error)

	// Open authenticates the ciphertext and, if it is valid, decrypts and returns plaintext.
	Open(SealedData) ([]byte, error)
}

// SealedData respresents an encrypted and authenticated message.
type SealedData interface {
	CiphertextBytes() []byte
	CiphertextHex() string

	NonceBytes() []byte
	NonceHex() string
}

// Config is used to configure a secret service. It contains either the key path
// or key bytes to use.
type Config struct {
	KeyPath  string
	KeyBytes *[SecretKeyLength]byte

	EmitStats    bool   // toggle emitting metrics or not
	StatsdHost   string // hostname of statsd server
	StatsdPort   int    // port of statsd server
	StatsdPrefix string // prefix to prepend to metrics
}

// SealedBytes contains the ciphertext and nonce for a sealed message.
type SealedBytes struct {
	Ciphertext []byte
	Nonce      []byte
}

func (s *SealedBytes) CiphertextBytes() []byte {
	return s.Ciphertext
}

func (s *SealedBytes) CiphertextHex() string {
	return base64.URLEncoding.EncodeToString(s.Ciphertext)
}

func (s *SealedBytes) NonceBytes() []byte {
	return s.Nonce
}

func (s *SealedBytes) NonceHex() string {
	return base64.URLEncoding.EncodeToString(s.Nonce)
}

// A Service can be used to seal/open (encrypt/decrypt and authenticate) messages.
type Service struct {
	secretKey     *[SecretKeyLength]byte
	metricsClient metrics.Client
}

// New returns a new Service. Config can not be nil.
func New(config *Config) (SecretService, error) {

	var err error
	var keyBytes *[SecretKeyLength]byte
	var metricsClient metrics.Client

	// read in key from keypath or if not given, try getting them from key bytes.
	if config.KeyPath != "" {
		keyBytes, err = ReadKeyFromDisk(config.KeyPath)
		if err != nil {
			return nil, err
		}
	} else {
		if config.KeyBytes == nil {
			return nil, fmt.Errorf("No key bytes provided.")
		}
		keyBytes = config.KeyBytes
	}

	// setup metrics service
	if config.EmitStats {
		// get hostname of box
		hostname, err := os.Hostname()
		if err != nil {
			return nil, fmt.Errorf("failed to obtain hostname: %v", err)
		}

		// build lemma prefix
		prefix := "lemma." + strings.Replace(hostname, ".", "_", -1)
		if config.StatsdPrefix != "" {
			prefix += "." + config.StatsdPrefix
		}

		// build metrics client
		hostport := fmt.Sprintf("%v:%v", config.StatsdHost, config.StatsdPort)
		metricsClient, err = metrics.NewWithOptions(hostport, prefix, metrics.Options{UseBuffering: true})
		if err != nil {
			return nil, err
		}
	} else {
		// if you don't want to emit stats, use the nop client
		metricsClient = metrics.NewNop()
	}

	return &Service{
		secretKey:     keyBytes,
		metricsClient: metricsClient,
	}, nil
}

// Seal takes plaintext and a key and returns encrypted and authenticated ciphertext.
// Allows passing in a key and useful for one off sealing purposes, otherwise
// create a secret.Service to seal multiple times.
func Seal(value []byte, secretKey *[SecretKeyLength]byte) (SealedData, error) {
	if secretKey == nil {
		return nil, fmt.Errorf("secret key is nil")
	}

	secretService, err := New(&Config{KeyBytes: secretKey})
	if err != nil {
		return nil, err
	}

	return secretService.Seal(value)
}

// Open authenticates the ciphertext and if valid, decrypts and returns plaintext.
// Allows passing in a key and useful for one off opening purposes, otherwise
// create a secret.Service to open multiple times.
func Open(e SealedData, secretKey *[SecretKeyLength]byte) ([]byte, error) {
	if secretKey == nil {
		return nil, fmt.Errorf("secret key is nil")
	}

	secretService, err := New(&Config{KeyBytes: secretKey})
	if err != nil {
		return nil, err
	}

	return secretService.Open(e)
}

// Seal takes plaintext and returns encrypted and authenticated ciphertext.
func (s *Service) Seal(value []byte) (SealedData, error) {
	// generate nonce
	nonce, err := generateNonce()
	if err != nil {
		return nil, fmt.Errorf("unable to generate nonce: %v", err)
	}

	// use nacl secret box to encrypt plaintext
	var encrypted []byte
	encrypted = secretbox.Seal(encrypted, value, nonce, s.secretKey)

	// return sealed ciphertext
	return &SealedBytes{
		Ciphertext: encrypted,
		Nonce:      nonce[:],
	}, nil
}

// Open authenticates the ciphertext and if valid, decrypts and returns plaintext.
func (s *Service) Open(e SealedData) (byt []byte, err error) {
	// once function is complete, check if we are returning err or not.
	// if we are, return emit a failure metric, if not a success metric.
	defer func() {
		if err == nil {
			s.metricsClient.Inc("success", 1, 1)
		} else {
			s.metricsClient.Inc("failure", 1, 1)
		}
	}()

	// convert nonce to an array
	nonce, err := nonceSliceToArray(e.NonceBytes())
	if err != nil {
		return nil, err
	}

	// decrypt
	var decrypted []byte
	decrypted, ok := secretbox.Open(decrypted, e.CiphertextBytes(), nonce, s.secretKey)
	if !ok {
		return nil, fmt.Errorf("unable to decrypt message")
	}

	return decrypted, nil
}

func ReadKeyFromDisk(keypath string) (*[SecretKeyLength]byte, error) {
	// load key from disk
	keyBytes, err := ioutil.ReadFile(keypath)
	if err != nil {
		return nil, err
	}

	// strip newline (\n or 0x0a) if it's at the end
	keyBytes = bytes.TrimSuffix(keyBytes, []byte("\n"))

	// decode string and convert to array and return it
	return EncodedStringToKey(string(keyBytes))
}

func KeySliceToArray(bytes []byte) (*[SecretKeyLength]byte, error) {
	// check that the lengths match
	if len(bytes) != SecretKeyLength {
		return nil, fmt.Errorf("wrong key length: %v", len(bytes))
	}

	// copy bytes into array
	var keyBytes [SecretKeyLength]byte
	copy(keyBytes[:], bytes)

	return &keyBytes, nil
}

func nonceSliceToArray(bytes []byte) (*[NonceLength]byte, error) {
	// check that the lengths match
	if len(bytes) != NonceLength {
		return nil, fmt.Errorf("wrong nonce length: %v", len(bytes))
	}

	// copy bytes into array
	var nonceBytes [NonceLength]byte
	copy(nonceBytes[:], bytes)

	return &nonceBytes, nil
}

func generateNonce() (*[NonceLength]byte, error) {
	// get b-bytes of random from /dev/urandom
	bytes, err := randomProvider.Bytes(NonceLength)
	if err != nil {
		return nil, err
	}

	return nonceSliceToArray(bytes)
}

var randomProvider random.RandomProvider

// init sets the package level randomProvider to be a real csprng. this is done
// so during tests, we can use a fake random number generator.
func init() {
	randomProvider = &random.CSPRNG{}
}
<commit_msg>Account for golang crypto relocation<commit_after>/*
Package secret provides tools for encrypting and decrypting authenticated messages.
See docs/secret.md for more details.
*/
package secret

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"golang.org/x/crypto/nacl/secretbox"
	"github.com/mailgun/lemma/random"
	"github.com/mailgun/metrics"
)

// SecretSevice is an interface for encrypting/decrypting and authenticating messages.
type SecretService interface {
	// Seal takes a plaintext message and returns an encrypted and authenticated ciphertext.
	Seal([]byte) (SealedData, error)

	// Open authenticates the ciphertext and, if it is valid, decrypts and returns plaintext.
	Open(SealedData) ([]byte, error)
}

// SealedData respresents an encrypted and authenticated message.
type SealedData interface {
	CiphertextBytes() []byte
	CiphertextHex() string

	NonceBytes() []byte
	NonceHex() string
}

// Config is used to configure a secret service. It contains either the key path
// or key bytes to use.
type Config struct {
	KeyPath  string
	KeyBytes *[SecretKeyLength]byte

	EmitStats    bool   // toggle emitting metrics or not
	StatsdHost   string // hostname of statsd server
	StatsdPort   int    // port of statsd server
	StatsdPrefix string // prefix to prepend to metrics
}

// SealedBytes contains the ciphertext and nonce for a sealed message.
type SealedBytes struct {
	Ciphertext []byte
	Nonce      []byte
}

func (s *SealedBytes) CiphertextBytes() []byte {
	return s.Ciphertext
}

func (s *SealedBytes) CiphertextHex() string {
	return base64.URLEncoding.EncodeToString(s.Ciphertext)
}

func (s *SealedBytes) NonceBytes() []byte {
	return s.Nonce
}

func (s *SealedBytes) NonceHex() string {
	return base64.URLEncoding.EncodeToString(s.Nonce)
}

// A Service can be used to seal/open (encrypt/decrypt and authenticate) messages.
type Service struct {
	secretKey     *[SecretKeyLength]byte
	metricsClient metrics.Client
}

// New returns a new Service. Config can not be nil.
func New(config *Config) (SecretService, error) {

	var err error
	var keyBytes *[SecretKeyLength]byte
	var metricsClient metrics.Client

	// read in key from keypath or if not given, try getting them from key bytes.
	if config.KeyPath != "" {
		keyBytes, err = ReadKeyFromDisk(config.KeyPath)
		if err != nil {
			return nil, err
		}
	} else {
		if config.KeyBytes == nil {
			return nil, fmt.Errorf("No key bytes provided.")
		}
		keyBytes = config.KeyBytes
	}

	// setup metrics service
	if config.EmitStats {
		// get hostname of box
		hostname, err := os.Hostname()
		if err != nil {
			return nil, fmt.Errorf("failed to obtain hostname: %v", err)
		}

		// build lemma prefix
		prefix := "lemma." + strings.Replace(hostname, ".", "_", -1)
		if config.StatsdPrefix != "" {
			prefix += "." + config.StatsdPrefix
		}

		// build metrics client
		hostport := fmt.Sprintf("%v:%v", config.StatsdHost, config.StatsdPort)
		metricsClient, err = metrics.NewWithOptions(hostport, prefix, metrics.Options{UseBuffering: true})
		if err != nil {
			return nil, err
		}
	} else {
		// if you don't want to emit stats, use the nop client
		metricsClient = metrics.NewNop()
	}

	return &Service{
		secretKey:     keyBytes,
		metricsClient: metricsClient,
	}, nil
}

// Seal takes plaintext and a key and returns encrypted and authenticated ciphertext.
// Allows passing in a key and useful for one off sealing purposes, otherwise
// create a secret.Service to seal multiple times.
func Seal(value []byte, secretKey *[SecretKeyLength]byte) (SealedData, error) {
	if secretKey == nil {
		return nil, fmt.Errorf("secret key is nil")
	}

	secretService, err := New(&Config{KeyBytes: secretKey})
	if err != nil {
		return nil, err
	}

	return secretService.Seal(value)
}

// Open authenticates the ciphertext and if valid, decrypts and returns plaintext.
// Allows passing in a key and useful for one off opening purposes, otherwise
// create a secret.Service to open multiple times.
func Open(e SealedData, secretKey *[SecretKeyLength]byte) ([]byte, error) {
	if secretKey == nil {
		return nil, fmt.Errorf("secret key is nil")
	}

	secretService, err := New(&Config{KeyBytes: secretKey})
	if err != nil {
		return nil, err
	}

	return secretService.Open(e)
}

// Seal takes plaintext and returns encrypted and authenticated ciphertext.
func (s *Service) Seal(value []byte) (SealedData, error) {
	// generate nonce
	nonce, err := generateNonce()
	if err != nil {
		return nil, fmt.Errorf("unable to generate nonce: %v", err)
	}

	// use nacl secret box to encrypt plaintext
	var encrypted []byte
	encrypted = secretbox.Seal(encrypted, value, nonce, s.secretKey)

	// return sealed ciphertext
	return &SealedBytes{
		Ciphertext: encrypted,
		Nonce:      nonce[:],
	}, nil
}

// Open authenticates the ciphertext and if valid, decrypts and returns plaintext.
func (s *Service) Open(e SealedData) (byt []byte, err error) {
	// once function is complete, check if we are returning err or not.
	// if we are, return emit a failure metric, if not a success metric.
	defer func() {
		if err == nil {
			s.metricsClient.Inc("success", 1, 1)
		} else {
			s.metricsClient.Inc("failure", 1, 1)
		}
	}()

	// convert nonce to an array
	nonce, err := nonceSliceToArray(e.NonceBytes())
	if err != nil {
		return nil, err
	}

	// decrypt
	var decrypted []byte
	decrypted, ok := secretbox.Open(decrypted, e.CiphertextBytes(), nonce, s.secretKey)
	if !ok {
		return nil, fmt.Errorf("unable to decrypt message")
	}

	return decrypted, nil
}

func ReadKeyFromDisk(keypath string) (*[SecretKeyLength]byte, error) {
	// load key from disk
	keyBytes, err := ioutil.ReadFile(keypath)
	if err != nil {
		return nil, err
	}

	// strip newline (\n or 0x0a) if it's at the end
	keyBytes = bytes.TrimSuffix(keyBytes, []byte("\n"))

	// decode string and convert to array and return it
	return EncodedStringToKey(string(keyBytes))
}

func KeySliceToArray(bytes []byte) (*[SecretKeyLength]byte, error) {
	// check that the lengths match
	if len(bytes) != SecretKeyLength {
		return nil, fmt.Errorf("wrong key length: %v", len(bytes))
	}

	// copy bytes into array
	var keyBytes [SecretKeyLength]byte
	copy(keyBytes[:], bytes)

	return &keyBytes, nil
}

func nonceSliceToArray(bytes []byte) (*[NonceLength]byte, error) {
	// check that the lengths match
	if len(bytes) != NonceLength {
		return nil, fmt.Errorf("wrong nonce length: %v", len(bytes))
	}

	// copy bytes into array
	var nonceBytes [NonceLength]byte
	copy(nonceBytes[:], bytes)

	return &nonceBytes, nil
}

func generateNonce() (*[NonceLength]byte, error) {
	// get b-bytes of random from /dev/urandom
	bytes, err := randomProvider.Bytes(NonceLength)
	if err != nil {
		return nil, err
	}

	return nonceSliceToArray(bytes)
}

var randomProvider random.RandomProvider

// init sets the package level randomProvider to be a real csprng. this is done
// so during tests, we can use a fake random number generator.
func init() {
	randomProvider = &random.CSPRNG{}
}
<|endoftext|>


<commit_before>package watch

import (
	"io/ioutil"
	"os"
	"path/filepath"
)

///////////////////////////////////////////////////////////////////////////////

type FileSystemItem struct {
	Root     string
	Path     string
	Name     string
	Size     int64
	Modified int64
	IsFolder bool

	ProfileDisabled  bool
	ProfileTags      []string
	ProfileArguments []string
}

///////////////////////////////////////////////////////////////////////////////

func YieldFileSystemItems(root string) chan *FileSystemItem {
	items := make(chan *FileSystemItem)

	go func() {
		filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return filepath.SkipDir
			}

			items <- &FileSystemItem{
				Root:     root,
				Path:     path,
				Name:     info.Name(),
				Size:     info.Size(),
				Modified: info.ModTime().Unix(),
				IsFolder: info.IsDir(),
			}

			return nil
		})
		close(items)
	}()

	return items
}

///////////////////////////////////////////////////////////////////////////////

// ReadContents reads files wholesale. This function is only called on files
// that end in '.goconvey'. These files should be very small, probably not
// ever more than a few hundred bytes. The ignored errors are ok because in
// the event of an IO error all that need be returned is an empty string.
func ReadContents(path string) string {
	file, _ := os.Open(path)
	defer file.Close()
	content, _ := ioutil.ReadAll(file)
	return string(content)
}

///////////////////////////////////////////////////////////////////////////////
<commit_msg>Limiting the size of profiles.<commit_after>package watch

import (
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

///////////////////////////////////////////////////////////////////////////////

type FileSystemItem struct {
	Root     string
	Path     string
	Name     string
	Size     int64
	Modified int64
	IsFolder bool

	ProfileDisabled  bool
	ProfileTags      []string
	ProfileArguments []string
}

///////////////////////////////////////////////////////////////////////////////

func YieldFileSystemItems(root string) chan *FileSystemItem {
	items := make(chan *FileSystemItem)

	go func() {
		filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return filepath.SkipDir
			}

			items <- &FileSystemItem{
				Root:     root,
				Path:     path,
				Name:     info.Name(),
				Size:     info.Size(),
				Modified: info.ModTime().Unix(),
				IsFolder: info.IsDir(),
			}

			return nil
		})
		close(items)
	}()

	return items
}

///////////////////////////////////////////////////////////////////////////////

// ReadContents reads files wholesale. This function is only called on files
// that end in '.goconvey'. These files should be very small, probably not
// ever more than a few hundred bytes. The ignored errors are ok because in
// the event of an IO error all that need be returned is an empty string.
func ReadContents(path string) string {
	file, err := os.Open(path)
	if err != nil {
		return ""
	}
	defer file.Close()
	reader := io.LimitReader(file, 1024*4)
	content, _ := ioutil.ReadAll(reader)
	return string(content)
}

///////////////////////////////////////////////////////////////////////////////
<|endoftext|>


<commit_before>package sender
// A multicast sender, i.e. publisher.
// I'm using the name sender here because this is mostly an experiment while learning Go.
// However, I do intend to implement a message bus publisher based on what I learn here.

//--------------------------------------------------------------------------------------------------

import (
	"net"
	"fmt"
	"github.com/jimlloyd/mbus/header"
	"github.com/jimlloyd/mbus/packet"
	"github.com/jimlloyd/mbus/utils"
	"github.com/jimlloyd/mbus/sender/history"
)

type Sender struct {
	conn *net.UDPConn
	mcast *net.UDPAddr
	sentTo	uint64

	history	*history.History
}

func NewSender(mcastAddress string) (*Sender, error) {
	var err error

	sender := new(Sender)

	// One connection is used both to send multicasts and to receive command packets.
	// We create the connection by setting up listening for commands packets,
	// but can also use the connection to send multicasts.
	sender.conn, err = utils.ListenUDP4()
	if err != nil {
		return nil, err
	}

	sender.mcast, err = net.ResolveUDPAddr("udp4", mcastAddress)
	if err != nil {
		return nil, err
	}

	minAgeSeconds := int32(10)
	maxAgeSeconds := int32(20)
	maxPayloadMB := uint32(50)
	sender.history = history.NewHistory(minAgeSeconds, maxAgeSeconds, maxPayloadMB)

	commands := make(chan packet.Packet, 10)
	go packet.Listen(sender.conn, commands)
	go sender.serveCommand(commands)

	return sender, nil
}

func (sender *Sender) Close() error {
	return sender.conn.Close()
}

func (sender *Sender) Send(payload []byte) (int, error) {

	h := header.MakeMessageHeader(sender.sentTo)
	buf, err := h.Encode()
	if err != nil {
		return 0, err
	}

	buf.Write(payload)
	message := buf.Bytes()

	sender.history.Add(sender.sentTo, message)

	n, err := sender.conn.WriteToUDP(message, sender.mcast)
	if err != nil {
		return 0, err
	}
	sender.sentTo += uint64(len(payload))

	return n, err
}

func (sender *Sender) ChannelSender(payloads <-chan []byte) {
	for {
		payload := <- payloads
		_, err := sender.Send(payload)
		if err!=nil {
			panic(err)
		}
	}
}

func (sender *Sender) serveCommand(commands <-chan packet.Packet) {
	for {
		packet := <-commands

		// We'll eventually respond to meaningful commands, but for now just log them.
		fmt.Println("Received command", len(packet.Data), "bytes:", string(packet.Data),
			"Remote:", packet.Remote())
	}
}

<commit_msg>Prepare for handling requests and responses on command interface.<commit_after>package sender
// A multicast sender, i.e. publisher.
// I'm using the name sender here because this is mostly an experiment while learning Go.
// However, I do intend to implement a message bus publisher based on what I learn here.

//--------------------------------------------------------------------------------------------------

import (
	"net"
	"fmt"
	"github.com/jimlloyd/mbus/header"
	"github.com/jimlloyd/mbus/packet"
	"github.com/jimlloyd/mbus/utils"
	"github.com/jimlloyd/mbus/sender/history"
)

type Sender struct {
	conn *net.UDPConn
	mcast *net.UDPAddr
	sentTo	uint64

	history	*history.History
}

func NewSender(mcastAddress string) (*Sender, error) {
	var err error

	sender := new(Sender)

	// One connection is used both to send multicasts and to receive command packets.
	// We create the connection by setting up listening for commands packets,
	// but can also use the connection to send multicasts.
	sender.conn, err = utils.ListenUDP4()
	if err != nil {
		return nil, err
	}

	sender.mcast, err = net.ResolveUDPAddr("udp4", mcastAddress)
	if err != nil {
		return nil, err
	}

	minAgeSeconds := int32(10)
	maxAgeSeconds := int32(20)
	maxPayloadMB := uint32(50)
	sender.history = history.NewHistory(minAgeSeconds, maxAgeSeconds, maxPayloadMB)

	commands := make(chan packet.Packet, 10)
	go packet.Listen(sender.conn, commands)
	go sender.serveCommand(commands)

	return sender, nil
}

func (sender *Sender) Close() error {
	return sender.conn.Close()
}

func (sender *Sender) Send(payload []byte) (int, error) {

	h := header.MakeMessageHeader(sender.sentTo)
	buf, err := h.Encode()
	if err != nil {
		return 0, err
	}

	buf.Write(payload)
	message := buf.Bytes()

	sender.history.Add(sender.sentTo, message)

	n, err := sender.conn.WriteToUDP(message, sender.mcast)
	if err != nil {
		return 0, err
	}
	sender.sentTo += uint64(len(payload))

	return n, err
}

func (sender *Sender) ChannelSender(payloads <-chan []byte) {
	for {
		payload := <- payloads
		_, err := sender.Send(payload)
		if err!=nil {
			panic(err)
		}
	}
}

func (sender *Sender) serveCommand(commands <-chan packet.Packet) {
	for {
		packet := <-commands

		messageType := header.PeekMessageType(packet.Data)
		switch messageType {
		case header.Invalid:
			fmt.Println("Ignoring invalid packet received on sender command interface.")
		case header.Request:
			sender.serveRequest(packet)
		case header.Response:
			sender.serveResponse(packet)
		default:
			fmt.Println("Message type", messageType, "not handled in sender command handler.")
		}

	}
}

func (sender *Sender) serveRequest(request packet.Packet) {
	// We'll eventually respond to requests, but for now just log them.
	fmt.Println("Received request from remote:", request.Remote())
}

func (sender *Sender) serveResponse(response packet.Packet) {
	// We'll eventually respond to responses, but for now just log them.
	fmt.Println("Received response from remote:", response.Remote())
}

<|endoftext|>


<commit_before>package serf

import (
	"bytes"
	"github.com/ugorji/go/codec"
)

// messageType are the types of gossip messages Serf will send along
// memberlist.
type messageType uint8

const (
	messageLeaveType messageType = iota
	messageJoinType
	messagePushPullType
	messageUserEventType
	messageQueryType
	messageQueryResponseType
)

// filterType is used with a queryFilter to specify the type of
// filter we are sending
type filterType uint8

const (
	filterNodeType filterType = iota
	filterTagType
)

// messageJoin is the message broadcasted after we join to
// associated the node with a lamport clock
type messageJoin struct {
	LTime LamportTime
	Node  string
}

// messageLeave is the message broadcasted to signal the intentional to
// leave.
type messageLeave struct {
	LTime LamportTime
	Node  string
}

// messagePushPullType is used when doing a state exchange. This
// is a relatively large message, but is sent infrequently
type messagePushPull struct {
	LTime        LamportTime            // Current node lamport time
	StatusLTimes map[string]LamportTime // Maps the node to its status time
	LeftMembers  []string               // List of left nodes
	EventLTime   LamportTime            // Lamport time for event clock
	Events       []*userEvents          // Recent events
}

// messageUserEvent is used for user-generated events
type messageUserEvent struct {
	LTime   LamportTime
	Name    string
	Payload []byte
	CC      bool // "Can Coalesce". Zero value is compatible with Serf 0.1
}

// messageQuery is used for query events
type messageQuery struct {
	LTime   LamportTime    // Event lamport time
	ID      uint32         // Query ID, randomly generated
	Addr    []byte         // Source address, used for a direct reply
	Port    uint16         // Source port, used for a direct reply
	Filters []*queryFilter // Potential query filters
	Ack     bool           // True if requesting an ack
	Name    string         // Query name
	Payload []byte         // Query payload
}

// queryFilter is set with a messageUserQuery to filter
// the set of nodes that need to reply
type queryFilter struct {
	Type filterType // Specify the type of filter
	Val  []byte     // Opaque blob, depends on the filter type
}

// messageQueryResponse is used to respond to a query
type messageQueryResponse struct {
	LTime   LamportTime // Event lamport time
	ID      uint32      // Query ID
	From    string      // Node name
	Payload []byte      // Optional response payload
}

func decodeMessage(buf []byte, out interface{}) error {
	var handle codec.MsgpackHandle
	return codec.NewDecoder(bytes.NewBuffer(buf), &handle).Decode(out)
}

func encodeMessage(t messageType, msg interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	buf.WriteByte(uint8(t))

	handle := codec.MsgpackHandle{}
	encoder := codec.NewEncoder(buf, &handle)
	err := encoder.Encode(msg)
	return buf.Bytes(), err
}
<commit_msg>serf: Adding field to distinguish ack from response<commit_after>package serf

import (
	"bytes"
	"github.com/ugorji/go/codec"
)

// messageType are the types of gossip messages Serf will send along
// memberlist.
type messageType uint8

const (
	messageLeaveType messageType = iota
	messageJoinType
	messagePushPullType
	messageUserEventType
	messageQueryType
	messageQueryResponseType
)

// filterType is used with a queryFilter to specify the type of
// filter we are sending
type filterType uint8

const (
	filterNodeType filterType = iota
	filterTagType
)

// messageJoin is the message broadcasted after we join to
// associated the node with a lamport clock
type messageJoin struct {
	LTime LamportTime
	Node  string
}

// messageLeave is the message broadcasted to signal the intentional to
// leave.
type messageLeave struct {
	LTime LamportTime
	Node  string
}

// messagePushPullType is used when doing a state exchange. This
// is a relatively large message, but is sent infrequently
type messagePushPull struct {
	LTime        LamportTime            // Current node lamport time
	StatusLTimes map[string]LamportTime // Maps the node to its status time
	LeftMembers  []string               // List of left nodes
	EventLTime   LamportTime            // Lamport time for event clock
	Events       []*userEvents          // Recent events
}

// messageUserEvent is used for user-generated events
type messageUserEvent struct {
	LTime   LamportTime
	Name    string
	Payload []byte
	CC      bool // "Can Coalesce". Zero value is compatible with Serf 0.1
}

// messageQuery is used for query events
type messageQuery struct {
	LTime   LamportTime    // Event lamport time
	ID      uint32         // Query ID, randomly generated
	Addr    []byte         // Source address, used for a direct reply
	Port    uint16         // Source port, used for a direct reply
	Filters []*queryFilter // Potential query filters
	Ack     bool           // True if requesting an ack
	Name    string         // Query name
	Payload []byte         // Query payload
}

// queryFilter is set with a messageUserQuery to filter
// the set of nodes that need to reply
type queryFilter struct {
	Type filterType // Specify the type of filter
	Val  []byte     // Opaque blob, depends on the filter type
}

// messageQueryResponse is used to respond to a query
type messageQueryResponse struct {
	LTime   LamportTime // Event lamport time
	ID      uint32      // Query ID
	From    string      // Node name
	Ack     bool        // Is this an Ack, or reply
	Payload []byte      // Optional response payload
}

func decodeMessage(buf []byte, out interface{}) error {
	var handle codec.MsgpackHandle
	return codec.NewDecoder(bytes.NewBuffer(buf), &handle).Decode(out)
}

func encodeMessage(t messageType, msg interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	buf.WriteByte(uint8(t))

	handle := codec.MsgpackHandle{}
	encoder := codec.NewEncoder(buf, &handle)
	err := encoder.Encode(msg)
	return buf.Bytes(), err
}
<|endoftext|>