<commit_before>package bufferedwriter

import (
  "io"
  "os"
  "log"
  "time"
  "strconv"
)

type Worker struct {
  length int
  data []byte
  capacity int
  fileRoot string
  fileTemp string
  logger *log.Logger
  permission os.FileMode
  channel chan io.ReadCloser
}

func newWorker(id int, channel chan io.ReadCloser, configuration *Configuration) (*Worker) {
  idString := strconv.Itoa(id)
  w := &Worker{
    channel: channel,
    logger: configuration.logger,
    capacity: configuration.size,
    permission: configuration.permission,
    data: make([]byte, configuration.size),
    fileRoot: configuration.path,
    fileTemp: configuration.temp,
  }

  if w.fileRoot[len(w.fileRoot)-1:] != "/" {
    w.fileRoot += "/"
  }
  w.fileRoot += configuration.prefix + idString + "_"
  w.fileTemp += configuration.prefix + idString + ".tmp"
  return w
}

func (w *Worker) work() {
  os.Remove(w.fileTemp)
  for {
    message := <- w.channel
    w.process(message)
  }
}

func (w *Worker) process(message io.ReadCloser) {
  defer message.Close()
  var swapped bool
  for {
    read, err := message.Read(w.data[w.length:])
    w.length += read
    if err == io.EOF {
      if w.length == w.capacity || swapped {
        w.swap()
        w.save()
      }
      break
    }
    if err != nil {
      w.length -= read
      if w.logger != nil {
        w.logger.Print("Failed to read reader", err)
      }
    }
    if w.length == w.capacity { swapped = w.swap() }
  }
}

func (w *Worker) swap() bool {
  if w.length == 0 { return false }
  defer func() { w.length = 0 }()
  f, err := os.OpenFile(w.fileTemp, os.O_WRONLY | os.O_APPEND | os.O_CREATE, 0600)
  if err != nil {
    if w.logger != nil {
      w.logger.Print("Failed to create temporary buffered file", err)
    }
    return false
  }
  defer f.Close()
  f.Write(w.data[0:w.length])
  stat, _ := f.Stat()
  return stat.Size() >= int64(w.capacity)
}

func (w *Worker) save() {
  target := w.fileRoot + strconv.FormatInt(time.Now().UnixNano(), 10) + ".log"
  if err := os.Rename(w.fileTemp, target); err != nil {
    if w.logger != nil {
      w.logger.Printf("Failed to rename %v to %v. Error: %v", w.fileTemp, target, err)
    }
    return
  }
  os.Chmod(target, w.permission)
}
<commit_msg>make sure temp path has a trailing slash<commit_after>package bufferedwriter

import (
  "io"
  "os"
  "log"
  "time"
  "strconv"
)

type Worker struct {
  length int
  data []byte
  capacity int
  fileRoot string
  fileTemp string
  logger *log.Logger
  permission os.FileMode
  channel chan io.ReadCloser
}

func newWorker(id int, channel chan io.ReadCloser, configuration *Configuration) (*Worker) {
  idString := strconv.Itoa(id)
  w := &Worker{
    channel: channel,
    logger: configuration.logger,
    capacity: configuration.size,
    permission: configuration.permission,
    data: make([]byte, configuration.size),
    fileRoot: configuration.path,
    fileTemp: configuration.temp,
  }

  if w.fileRoot[len(w.fileRoot)-1:] != "/" {
    w.fileRoot += "/"
  }

  if w.fileTemp[len(w.fileTemp)-1:] != "/" {
    w.fileTemp += "/"
  }
  w.fileRoot += configuration.prefix + idString + "_"
  w.fileTemp += configuration.prefix + idString + ".tmp"
  return w
}

func (w *Worker) work() {
  os.Remove(w.fileTemp)
  for {
    message := <- w.channel
    w.process(message)
  }
}

func (w *Worker) process(message io.ReadCloser) {
  defer message.Close()
  var swapped bool
  for {
    read, err := message.Read(w.data[w.length:])
    w.length += read
    if err == io.EOF {
      if w.length == w.capacity || swapped {
        w.swap()
        w.save()
      }
      break
    }
    if err != nil {
      w.length -= read
      if w.logger != nil {
        w.logger.Print("Failed to read reader", err)
      }
    }
    if w.length == w.capacity { swapped = w.swap() }
  }
}

func (w *Worker) swap() bool {
  if w.length == 0 { return false }
  defer func() { w.length = 0 }()
  f, err := os.OpenFile(w.fileTemp, os.O_WRONLY | os.O_APPEND | os.O_CREATE, 0600)
  if err != nil {
    if w.logger != nil {
      w.logger.Print("Failed to create temporary buffered file", err)
    }
    return false
  }
  defer f.Close()
  f.Write(w.data[0:w.length])
  stat, _ := f.Stat()
  return stat.Size() >= int64(w.capacity)
}

func (w *Worker) save() {
  target := w.fileRoot + strconv.FormatInt(time.Now().UnixNano(), 10) + ".log"
  if err := os.Rename(w.fileTemp, target); err != nil {
    if w.logger != nil {
      w.logger.Printf("Failed to rename %v to %v. Error: %v", w.fileTemp, target, err)
    }
    return
  }
  os.Chmod(target, w.permission)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"time"
)

// A Worker runs a job.
type Worker struct {
	Name       string
	vmProvider VMProvider
	mb         MessageBroker
	logger     *log.Logger
	payload    Payload
	reporter   *Reporter
	tw         *TimeoutWriter
	lw         *LimitWriter
	timeouts   TimeoutsConfig
	logLimits  LogLimitsConfig
}

// NewWorker creates a new worker with the given parameters. The worker assumes
// ownership of the given API, payload and reporter and they should not be
// reused for other workers.
func NewWorker(name string, VMProvider VMProvider, mb MessageBroker, timeouts TimeoutsConfig, logLimits LogLimitsConfig) *Worker {
	return &Worker{
		Name:       name,
		vmProvider: VMProvider,
		mb:         mb,
		logger:     log.New(os.Stdout, fmt.Sprintf("%s: ", name), log.Ldate|log.Ltime),
		timeouts:   timeouts,
		logLimits:  logLimits,
	}
}

// Process actually runs the job. It returns an error if an error occurred that
// should cause the job to be requeued.
func (w *Worker) Process(payload Payload) error {
	defer w.logger.Println("Finishing job")

	var err error
	w.payload = payload
	w.reporter, err = NewReporter(w.mb, payload.Job.ID)
	if err != nil {
		return err
	}

	w.logger.Printf("Starting job slug:%s id:%d\n", w.payload.Repository.Slug, w.payload.Job.ID)

	server, err := w.bootServer()
	if err != nil {
		w.logger.Printf("Booting a VM failed with the following errors: %v\n", err)
		return err
	}
	defer server.Destroy()

	fmt.Fprintf(w.reporter.Log, "Using worker: %s\n\n", w.Name)
	defer w.reporter.Log.Close()

	w.logger.Println("Opening SSH connection")
	ssh, err := NewSSHConnection(server, w.Name)
	if err != nil {
		w.logger.Printf("Couldn't connect to SSH: %v\n", err)
		fmt.Fprintf(w.reporter.Log, connectionErrorMessage)
		return err
	}
	defer ssh.Close()

	w.logger.Println("Uploading build script")
	err = w.uploadScript(ssh)
	if err != nil {
		w.logger.Printf("Couldn't upload script to SSH: %v\n", err)
		fmt.Fprintf(w.reporter.Log, connectionErrorMessage)
		return err
	}

	err = w.reporter.NotifyJobStarted()
	if err != nil {
		w.logger.Printf("Couldn't notify about job start: %v\n", err)
		return err
	}

	w.logger.Println("Running the build")
	exitCodeChan, err := w.runScript(ssh)
	if err != nil {
		w.logger.Printf("Failed to run build script: %v\n", err)
		fmt.Fprintf(w.reporter.Log, connectionErrorMessage)
		return err
	}

	select {
	case exitCode := <-exitCodeChan:
		w.logger.Println("Build finished.")
		switch exitCode {
		case 0:
			w.reporter.NotifyJobFinished("passed")
		case 1:
			w.reporter.NotifyJobFinished("failed")
		default:
			w.reporter.NotifyJobFinished("errored")
		case -1:
			fmt.Fprintf(w.reporter.Log, connectionErrorMessage)
			return fmt.Errorf("an error occurred with the SSH connection")
		}
		return nil
	case <-w.tw.Timeout:
		fmt.Fprintf(w.reporter.Log, stalledBuildMessage, w.timeouts.LogInactivity/60)
		return nil
	case <-w.lw.LimitReached:
		fmt.Fprintf(w.reporter.Log, logTooLongMessage, w.logLimits.MaxLogLength/1024/1024)
		return nil
	case <-time.After(time.Duration(w.timeouts.HardLimit) * time.Second):
		fmt.Fprintf(w.reporter.Log, noLogOutputMessage, w.timeouts.HardLimit/60)
		w.reporter.NotifyJobFinished("errored")
		return nil
	}
}

func (w *Worker) jobID() int64 {
	return w.payload.Job.ID
}

func (w *Worker) bootServer() (VM, error) {
	startTime := time.Now()
	hostname := fmt.Sprintf("testing-worker-go-%d-%s-%d", os.Getpid(), w.Name, w.jobID())
	w.logger.Printf("Booting %s\n", hostname)
	server, err := w.vmProvider.Start(hostname, time.Duration(w.timeouts.VMBoot)*time.Second)
	if err != nil {
		return nil, err
	}

	w.logger.Printf("VM provisioned in %.2f seconds\n", time.Now().Sub(startTime).Seconds())

	return server, nil
}

func (w *Worker) uploadScript(ssh *SSHConnection) error {
	err := ssh.UploadFile("~/build.sh", []byte(fmt.Sprintf("#!/bin/bash --login\n\necho This is build id %d\nfor i in {1..200}; do echo -n \"$i \"; sleep 1; done", w.jobID())))
	if err != nil {
		return err
	}

	return ssh.Run("chmod +x ~/build.sh")
}

func (w *Worker) runScript(ssh *SSHConnection) (<-chan int, error) {
	w.tw = NewTimeoutWriter(w.reporter.Log, time.Duration(w.timeouts.LogInactivity)*time.Second)
	w.lw = NewLimitWriter(w.tw, w.logLimits.MaxLogLength)
	return ssh.Start("~/build.sh", w.lw)
}
<commit_msg>worker: Print out the right error message for the right error<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"time"
)

// A Worker runs a job.
type Worker struct {
	Name       string
	vmProvider VMProvider
	mb         MessageBroker
	logger     *log.Logger
	payload    Payload
	reporter   *Reporter
	tw         *TimeoutWriter
	lw         *LimitWriter
	timeouts   TimeoutsConfig
	logLimits  LogLimitsConfig
}

// NewWorker creates a new worker with the given parameters. The worker assumes
// ownership of the given API, payload and reporter and they should not be
// reused for other workers.
func NewWorker(name string, VMProvider VMProvider, mb MessageBroker, timeouts TimeoutsConfig, logLimits LogLimitsConfig) *Worker {
	return &Worker{
		Name:       name,
		vmProvider: VMProvider,
		mb:         mb,
		logger:     log.New(os.Stdout, fmt.Sprintf("%s: ", name), log.Ldate|log.Ltime),
		timeouts:   timeouts,
		logLimits:  logLimits,
	}
}

// Process actually runs the job. It returns an error if an error occurred that
// should cause the job to be requeued.
func (w *Worker) Process(payload Payload) error {
	defer w.logger.Println("Finishing job")

	var err error
	w.payload = payload
	w.reporter, err = NewReporter(w.mb, payload.Job.ID)
	if err != nil {
		return err
	}

	w.logger.Printf("Starting job slug:%s id:%d\n", w.payload.Repository.Slug, w.payload.Job.ID)

	server, err := w.bootServer()
	if err != nil {
		w.logger.Printf("Booting a VM failed with the following errors: %v\n", err)
		return err
	}
	defer server.Destroy()

	fmt.Fprintf(w.reporter.Log, "Using worker: %s\n\n", w.Name)
	defer w.reporter.Log.Close()

	w.logger.Println("Opening SSH connection")
	ssh, err := NewSSHConnection(server, w.Name)
	if err != nil {
		w.logger.Printf("Couldn't connect to SSH: %v\n", err)
		fmt.Fprintf(w.reporter.Log, connectionErrorMessage)
		return err
	}
	defer ssh.Close()

	w.logger.Println("Uploading build script")
	err = w.uploadScript(ssh)
	if err != nil {
		w.logger.Printf("Couldn't upload script to SSH: %v\n", err)
		fmt.Fprintf(w.reporter.Log, connectionErrorMessage)
		return err
	}

	err = w.reporter.NotifyJobStarted()
	if err != nil {
		w.logger.Printf("Couldn't notify about job start: %v\n", err)
		return err
	}

	w.logger.Println("Running the build")
	exitCodeChan, err := w.runScript(ssh)
	if err != nil {
		w.logger.Printf("Failed to run build script: %v\n", err)
		fmt.Fprintf(w.reporter.Log, connectionErrorMessage)
		return err
	}

	select {
	case exitCode := <-exitCodeChan:
		w.logger.Println("Build finished.")
		switch exitCode {
		case 0:
			w.reporter.NotifyJobFinished("passed")
		case 1:
			w.reporter.NotifyJobFinished("failed")
		default:
			w.reporter.NotifyJobFinished("errored")
		case -1:
			fmt.Fprintf(w.reporter.Log, connectionErrorMessage)
			return fmt.Errorf("an error occurred with the SSH connection")
		}
		return nil
	case <-w.tw.Timeout:
		fmt.Fprintf(w.reporter.Log, noLogOutputMessage, w.timeouts.LogInactivity/60)
		return nil
	case <-w.lw.LimitReached:
		fmt.Fprintf(w.reporter.Log, logTooLongMessage, w.logLimits.MaxLogLength/1024/1024)
		return nil
	case <-time.After(time.Duration(w.timeouts.HardLimit) * time.Second):
		fmt.Fprintf(w.reporter.Log, stalledBuildMessage, w.timeouts.HardLimit/60)
		w.reporter.NotifyJobFinished("errored")
		return nil
	}
}

func (w *Worker) jobID() int64 {
	return w.payload.Job.ID
}

func (w *Worker) bootServer() (VM, error) {
	startTime := time.Now()
	hostname := fmt.Sprintf("testing-worker-go-%d-%s-%d", os.Getpid(), w.Name, w.jobID())
	w.logger.Printf("Booting %s\n", hostname)
	server, err := w.vmProvider.Start(hostname, time.Duration(w.timeouts.VMBoot)*time.Second)
	if err != nil {
		return nil, err
	}

	w.logger.Printf("VM provisioned in %.2f seconds\n", time.Now().Sub(startTime).Seconds())

	return server, nil
}

func (w *Worker) uploadScript(ssh *SSHConnection) error {
	err := ssh.UploadFile("~/build.sh", []byte(fmt.Sprintf("#!/bin/bash --login\n\necho This is build id %d\nfor i in {1..200}; do echo -n \"$i \"; sleep 1; done", w.jobID())))
	if err != nil {
		return err
	}

	return ssh.Run("chmod +x ~/build.sh")
}

func (w *Worker) runScript(ssh *SSHConnection) (<-chan int, error) {
	w.tw = NewTimeoutWriter(w.reporter.Log, time.Duration(w.timeouts.LogInactivity)*time.Second)
	w.lw = NewLimitWriter(w.tw, w.logLimits.MaxLogLength)
	return ssh.Start("~/build.sh", w.lw)
}
<|endoftext|>


<commit_before>package main

import (
    "log"
    "container/list"
    "huabot-sched/db"
    "strconv"
    "bytes"
)

const (
    NOOP = iota // server
    // for job
    GRAB_JOB    // client
    SCHED_LATER // client
    JOB_DONE    // client
    JOB_FAIL    // client
    WAIT_JOB    // server
    NO_JOB      // server
    // for func
    CAN_DO      // client
    CANT_DO     // client
    // for test
    PING        // client
    PONG        // server
    // other
    SLEEP       // client
    UNKNOWN     // server
)

type Worker struct {
    jobQueue *list.List
    conn     Conn
    sched    *Sched
    alive    bool
    Funcs    []string
}

func NewWorker(sched *Sched, conn Conn) (worker *Worker) {
    worker = new(Worker)
    worker.conn = conn
    worker.jobQueue = list.New()
    worker.sched = sched
    worker.Funcs = make([]string, 0)
    worker.alive = true
    return
}

func (worker *Worker) IsAlive() bool {
    return worker.alive
}

func (worker *Worker) HandleDo(job db.Job) (err error){
    log.Printf("HandleDo: %d\n", job.Id)
    worker.jobQueue.PushBack(job)
    pack, err := packJob(job)
    if err != nil {
        log.Printf("Error: packJob %d %s\n", job.Id, err.Error())
        return nil
    }
    err = worker.conn.Send(pack)
    return
}

func (worker *Worker) HandleCanDo(Func string) error {
    log.Printf("HandleCanDo: %s\n", Func)
    for _, f := range worker.Funcs {
        if f == Func {
            return nil
        }
    }
    worker.Funcs = append(worker.Funcs, Func)
    worker.sched.IncrStatFunc(Func)
    return nil
}

func (worker *Worker) HandleCanNoDo(Func string) error {
    log.Printf("HandleCanDo: %s\n", Func)
    newFuncs := make([]string, 0)
    for _, f := range worker.Funcs {
        if f == Func {
            continue
        }
        newFuncs = append(newFuncs, f)
    }
    worker.Funcs = newFuncs
    return nil
}

func (worker *Worker) HandleDone(jobId int64) (err error) {
    log.Printf("HandleDone: %d\n", jobId)
    worker.sched.Done(jobId)
    removeListJob(worker.jobQueue, jobId)
    return nil
}

func (worker *Worker) HandleFail(jobId int64) (err error) {
    log.Printf("HandleFail: %d\n", jobId)
    worker.sched.Fail(jobId)
    removeListJob(worker.jobQueue, jobId)
    return nil
}

func (worker *Worker) HandleWaitForJob() (err error) {
    log.Printf("HandleWaitForJob\n")
    err = worker.conn.Send(packCmd(WAIT_JOB))
    return nil
}

func (worker *Worker) HandleSchedLater(jobId, delay int64) (err error){
    log.Printf("HandleSchedLater: %d %d\n", jobId, delay)
    worker.sched.SchedLater(jobId, delay)
    removeListJob(worker.jobQueue, jobId)
    return nil
}

func (worker *Worker) HandleNoJob() (err error){
    log.Printf("HandleNoJob\n")
    err = worker.conn.Send(packCmd(NO_JOB))
    return
}

func (worker *Worker) HandleGrabJob() (err error){
    log.Printf("HandleGrabJob\n")
    worker.sched.grabQueue.PushBack(worker)
    worker.sched.Notify()
    return nil
}

func (worker *Worker) Handle() {
    var payload []byte
    var err error
    var conn = worker.conn
    for {
        payload, err = conn.Receive()
        if err != nil {
            log.Printf("Error: %s\n", err.Error())
            worker.sched.DieWorker(worker)
            return
        }

        switch payload[0] {
        case GRAB_JOB:
            err = worker.HandleGrabJob()
            break
        case JOB_DONE:
            jobId, _ := strconv.ParseInt(string(payload[2:]), 10, 0)
            err = worker.HandleDone(jobId)
            break
        case JOB_FAIL:
            jobId, _ := strconv.ParseInt(string(payload[2:]), 10, 0)
            err = worker.HandleFail(jobId)
            break
        case SCHED_LATER:
            parts := bytes.SplitN(payload[2:], NULL_CHAR, 2)
            if len(parts) != 2 {
                log.Printf("Error: invalid format.")
                break
            }
            jobId, _ := strconv.ParseInt(string(parts[0]), 10, 0)
            delay, _ := strconv.ParseInt(string(parts[1]), 10, 0)
            err = worker.HandleSchedLater(jobId, delay)
            break
        case SLEEP:
            err = conn.Send(packCmd(NOOP))
            break
        case PING:
            err = conn.Send(packCmd(PONG))
            break
        case CAN_DO:
            err = worker.HandleCanDo(string(payload[2:]))
            break
        case CANT_DO:
            err = worker.HandleCanNoDo(string(payload[2:]))
            break
        default:
            err = conn.Send(packCmd(UNKNOWN))
            break
        }
        if err != nil {
            log.Printf("Error: %s\n", err.Error())
            worker.alive = false
            worker.sched.DieWorker(worker)
            return
        }

        if !worker.alive {
            break
        }
    }
}

func (worker *Worker) Close() {
    worker.conn.Close()
    for e := worker.jobQueue.Front(); e != nil; e = e.Next() {
        worker.sched.Fail(e.Value.(db.Job).Id)
    }
}
<commit_msg>Fix. die worker not DecrStatFunc<commit_after>package main

import (
    "log"
    "container/list"
    "huabot-sched/db"
    "strconv"
    "bytes"
)

const (
    NOOP = iota // server
    // for job
    GRAB_JOB    // client
    SCHED_LATER // client
    JOB_DONE    // client
    JOB_FAIL    // client
    WAIT_JOB    // server
    NO_JOB      // server
    // for func
    CAN_DO      // client
    CANT_DO     // client
    // for test
    PING        // client
    PONG        // server
    // other
    SLEEP       // client
    UNKNOWN     // server
)

type Worker struct {
    jobQueue *list.List
    conn     Conn
    sched    *Sched
    alive    bool
    Funcs    []string
}

func NewWorker(sched *Sched, conn Conn) (worker *Worker) {
    worker = new(Worker)
    worker.conn = conn
    worker.jobQueue = list.New()
    worker.sched = sched
    worker.Funcs = make([]string, 0)
    worker.alive = true
    return
}

func (worker *Worker) IsAlive() bool {
    return worker.alive
}

func (worker *Worker) HandleDo(job db.Job) (err error){
    log.Printf("HandleDo: %d\n", job.Id)
    worker.jobQueue.PushBack(job)
    pack, err := packJob(job)
    if err != nil {
        log.Printf("Error: packJob %d %s\n", job.Id, err.Error())
        return nil
    }
    err = worker.conn.Send(pack)
    return
}

func (worker *Worker) HandleCanDo(Func string) error {
    log.Printf("HandleCanDo: %s\n", Func)
    for _, f := range worker.Funcs {
        if f == Func {
            return nil
        }
    }
    worker.Funcs = append(worker.Funcs, Func)
    worker.sched.IncrStatFunc(Func)
    return nil
}

func (worker *Worker) HandleCanNoDo(Func string) error {
    log.Printf("HandleCanDo: %s\n", Func)
    newFuncs := make([]string, 0)
    for _, f := range worker.Funcs {
        if f == Func {
            continue
        }
        newFuncs = append(newFuncs, f)
    }
    worker.Funcs = newFuncs
    return nil
}

func (worker *Worker) HandleDone(jobId int64) (err error) {
    log.Printf("HandleDone: %d\n", jobId)
    worker.sched.Done(jobId)
    removeListJob(worker.jobQueue, jobId)
    return nil
}

func (worker *Worker) HandleFail(jobId int64) (err error) {
    log.Printf("HandleFail: %d\n", jobId)
    worker.sched.Fail(jobId)
    removeListJob(worker.jobQueue, jobId)
    return nil
}

func (worker *Worker) HandleWaitForJob() (err error) {
    log.Printf("HandleWaitForJob\n")
    err = worker.conn.Send(packCmd(WAIT_JOB))
    return nil
}

func (worker *Worker) HandleSchedLater(jobId, delay int64) (err error){
    log.Printf("HandleSchedLater: %d %d\n", jobId, delay)
    worker.sched.SchedLater(jobId, delay)
    removeListJob(worker.jobQueue, jobId)
    return nil
}

func (worker *Worker) HandleNoJob() (err error){
    log.Printf("HandleNoJob\n")
    err = worker.conn.Send(packCmd(NO_JOB))
    return
}

func (worker *Worker) HandleGrabJob() (err error){
    log.Printf("HandleGrabJob\n")
    worker.sched.grabQueue.PushBack(worker)
    worker.sched.Notify()
    return nil
}

func (worker *Worker) Handle() {
    var payload []byte
    var err error
    var conn = worker.conn
    for {
        payload, err = conn.Receive()
        if err != nil {
            log.Printf("Error: %s\n", err.Error())
            worker.sched.DieWorker(worker)
            return
        }

        switch payload[0] {
        case GRAB_JOB:
            err = worker.HandleGrabJob()
            break
        case JOB_DONE:
            jobId, _ := strconv.ParseInt(string(payload[2:]), 10, 0)
            err = worker.HandleDone(jobId)
            break
        case JOB_FAIL:
            jobId, _ := strconv.ParseInt(string(payload[2:]), 10, 0)
            err = worker.HandleFail(jobId)
            break
        case SCHED_LATER:
            parts := bytes.SplitN(payload[2:], NULL_CHAR, 2)
            if len(parts) != 2 {
                log.Printf("Error: invalid format.")
                break
            }
            jobId, _ := strconv.ParseInt(string(parts[0]), 10, 0)
            delay, _ := strconv.ParseInt(string(parts[1]), 10, 0)
            err = worker.HandleSchedLater(jobId, delay)
            break
        case SLEEP:
            err = conn.Send(packCmd(NOOP))
            break
        case PING:
            err = conn.Send(packCmd(PONG))
            break
        case CAN_DO:
            err = worker.HandleCanDo(string(payload[2:]))
            break
        case CANT_DO:
            err = worker.HandleCanNoDo(string(payload[2:]))
            break
        default:
            err = conn.Send(packCmd(UNKNOWN))
            break
        }
        if err != nil {
            log.Printf("Error: %s\n", err.Error())
            worker.alive = false
            worker.sched.DieWorker(worker)
            return
        }

        if !worker.alive {
            break
        }
    }
}

func (worker *Worker) Close() {
    worker.conn.Close()
    for e := worker.jobQueue.Front(); e != nil; e = e.Next() {
        worker.sched.Fail(e.Value.(db.Job).Id)
    }
    for _, Func := range worker.Funcs {
        worker.sched.DecrStatFunc(Func)
    }
}
<|endoftext|>


<commit_before>package aws

import (
	"errors"
	"io/ioutil"
	"reflect"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go/aws/awserr"

	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

type mockS3API struct {
	s3iface.S3API
	getObjectOutput *s3.GetObjectOutput
	err             error
}

func (m *mockS3API) GetObject(in *s3.GetObjectInput) (*s3.GetObjectOutput, error) {
	return m.getObjectOutput, m.err
}

func (m *mockS3API) PutObject(in *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
	return nil, m.err
}

func TestStore_Get(t *testing.T) {
	type fields struct {
		s3     s3iface.S3API
		bucket string
	}
	type args struct {
		key string
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		want    []byte
		wantErr bool
	}{
		{
			name: "success",
			fields: fields{
				s3: &mockS3API{
					getObjectOutput: &s3.GetObjectOutput{
						Body: ioutil.NopCloser(strings.NewReader("my object body")),
					},
				},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
			want: []byte("my object body"),
		},
		{
			name: "failure",
			fields: fields{
				s3: &mockS3API{
					err: errors.New("an error"),
				},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
			wantErr: true,
		},
		{
			name: "not found",
			fields: fields{
				s3: &mockS3API{
					err: awserr.New(s3.ErrCodeNoSuchKey, "no such key", nil),
				},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
			wantErr: false,
			want:    nil,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Store{
				s3:     tt.fields.s3,
				bucket: tt.fields.bucket,
			}
			got, err := s.Get(tt.args.key)
			if (err != nil) != tt.wantErr {
				t.Errorf("Store.Get() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("Store.Get() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestStore_Set(t *testing.T) {
	type fields struct {
		s3     s3iface.S3API
		bucket string
	}
	type args struct {
		key   string
		value []byte
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		wantErr bool
	}{
		{
			name: "success",
			fields: fields{
				s3:     &mockS3API{},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
		},
		{
			name: "failure",
			fields: fields{
				s3: &mockS3API{
					err: errors.New("an error"),
				},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Store{
				s3:     tt.fields.s3,
				bucket: tt.fields.bucket,
			}
			if err := s.Set(tt.args.key, tt.args.value); (err != nil) != tt.wantErr {
				t.Errorf("Store.Set() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}
<commit_msg>ConfigureDirectorCloudConfig() is tested<commit_after>package aws

import (
	"errors"
	"io/ioutil"
	"reflect"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go/aws/awserr"

	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

type mockS3API struct {
	s3iface.S3API
	getObjectOutput *s3.GetObjectOutput
	err             error
}

func (m *mockS3API) GetObject(in *s3.GetObjectInput) (*s3.GetObjectOutput, error) {
	return m.getObjectOutput, m.err
}

func (m *mockS3API) PutObject(in *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
	return nil, m.err
}

func TestStore_Get(t *testing.T) {
	type fields struct {
		s3     s3iface.S3API
		bucket string
	}
	type args struct {
		key string
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		want    []byte
		wantErr bool
	}{
		{
			name: "success",
			fields: fields{
				s3: &mockS3API{
					getObjectOutput: &s3.GetObjectOutput{
						Body: ioutil.NopCloser(strings.NewReader("my object body")),
					},
				},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
			want: []byte("my object body"),
		},
		{
			name: "failure",
			fields: fields{
				s3: &mockS3API{
					err: errors.New("an error"),
				},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
			wantErr: true,
		},
		{
			name: "not found",
			fields: fields{
				s3: &mockS3API{
					err: awserr.New(s3.ErrCodeNoSuchKey, "no such key", nil),
				},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
			wantErr: false,
			want:    nil,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Store{
				s3:     tt.fields.s3,
				bucket: tt.fields.bucket,
			}
			got, err := s.Get(tt.args.key)
			if (err != nil) != tt.wantErr {
				t.Errorf("Store.Get() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("Store.Get() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestStore_Set(t *testing.T) {
	type fields struct {
		s3     s3iface.S3API
		bucket string
	}
	type args struct {
		key   string
		value []byte
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		wantErr bool
	}{
		{
			name: "success",
			fields: fields{
				s3:     &mockS3API{},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
		},
		{
			name: "failure",
			fields: fields{
				s3: &mockS3API{
					err: errors.New("an error"),
				},
				bucket: "my bucket",
			},
			args: args{
				key: "state.json",
			},
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &Store{
				s3:     tt.fields.s3,
				bucket: tt.fields.bucket,
			}
			if err := s.Set(tt.args.key, tt.args.value); (err != nil) != tt.wantErr {
				t.Errorf("Store.Set() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func TestEnvironment_ConfigureDirectorCloudConfig(t *testing.T) {
	type fields struct {
		AZ               string
		PublicSubnetID   string
		PrivateSubnetID  string
		ATCSecurityGroup string
		VMSecurityGroup  string
		Spot             bool
	}
	tests := []struct {
		name        string
		fields      fields
		cloudConfig string
		want        string
		wantErr     bool
	}{
		{
			name: "Success- template rendered",
			fields: fields{
				AZ:               "eu-west-1",
				PublicSubnetID:   "12345",
				PrivateSubnetID:  "67890",
				ATCSecurityGroup: "00000",
			},
			cloudConfig: "availability_zone: <% .AvailabilityZone %>\n public_subnet_id: <% .PublicSubnetID %>\n private_subnet_id: <% .PrivateSubnetID %>\n atc_security_group: <% .ATCSecurityGroupID %>\n spot: <% .Spot %>",
			want:        "availability_zone: eu-west-1\n public_subnet_id: 12345\n private_subnet_id: 67890\n atc_security_group: 00000\n spot: false",
			wantErr:     false,
		},
		{
			name:        "Failure- template not rendered",
			cloudConfig: "non_existant_key: <% .NonExistantKey %>",
			want:        "",
			wantErr:     true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := Environment{
				AZ:               tt.fields.AZ,
				PublicSubnetID:   tt.fields.PublicSubnetID,
				PrivateSubnetID:  tt.fields.PrivateSubnetID,
				ATCSecurityGroup: tt.fields.ATCSecurityGroup,
				VMSecurityGroup:  tt.fields.VMSecurityGroup,
				Spot:             tt.fields.Spot,
			}
			got, err := e.ConfigureDirectorCloudConfig(tt.cloudConfig)
			if (err != nil) != tt.wantErr {
				t.Errorf("Environment.ConfigureDirectorCloudConfig()\nerror expected:  %v\nreceived error:  %v", tt.wantErr, err)
				return
			}
			if got != tt.want {
				t.Errorf("Environment.ConfigureDirectorCloudConfig()\nexpected %v\nreceived %v", tt.want, got)
			}
		})
	}
}
<|endoftext|>


<commit_before>package builtins

import (
	"runtime"

	"github.com/dotcloud/docker/api"
	apiserver "github.com/dotcloud/docker/api/server"
	"github.com/dotcloud/docker/daemon/networkdriver/bridge"
	"github.com/dotcloud/docker/dockerversion"
	"github.com/dotcloud/docker/engine"
	"github.com/dotcloud/docker/registry"
	"github.com/dotcloud/docker/server"
	"github.com/dotcloud/docker/utils"
)

func Register(eng *engine.Engine) error {
	if err := daemon(eng); err != nil {
		return err
	}
	if err := remote(eng); err != nil {
		return err
	}
	if err := eng.Register("version", dockerVersion); err != nil {
		return err
	}
	return registry.NewService().Install(eng)
}

// remote: a RESTful api for cross-docker communication
func remote(eng *engine.Engine) error {
	if err := eng.Register("serveapi", apiserver.ServeApi); err != nil {
		return err
	}
	return eng.Register("acceptconnections", apiserver.AcceptConnections)
}

// daemon: a default execution and storage backend for Docker on Linux,
// with the following underlying components:
//
// * Pluggable storage drivers including aufs, vfs, lvm and btrfs.
// * Pluggable execution drivers including lxc and chroot.
//
// In practice `daemon` still includes most core Docker components, including:
//
// * The reference registry client implementation
// * Image management
// * The build facility
// * Logging
//
// These components should be broken off into plugins of their own.
//
func daemon(eng *engine.Engine) error {
	if err := eng.Register("initserver", server.InitServer); err != nil {
		return err
	}
	return eng.Register("init_networkdriver", bridge.InitDriver)
}

// builtins jobs independent of any subsystem
func dockerVersion(job *engine.Job) engine.Status {
	v := &engine.Env{}
	v.Set("Version", dockerversion.VERSION)
	v.SetJson("ApiVersion", api.APIVERSION)
	v.Set("GitCommit", dockerversion.GITCOMMIT)
	v.Set("GoVersion", runtime.Version())
	v.Set("Os", runtime.GOOS)
	v.Set("Arch", runtime.GOARCH)
	if kernelVersion, err := utils.GetKernelVersion(); err == nil {
		v.Set("KernelVersion", kernelVersion.String())
	}
	if _, err := v.WriteTo(job.Stdout); err != nil {
		return job.Error(err)
	}
	return engine.StatusOK
}
<commit_msg>fix version numbering<commit_after>package builtins

import (
	"runtime"

	"github.com/dotcloud/docker/api"
	apiserver "github.com/dotcloud/docker/api/server"
	"github.com/dotcloud/docker/daemon/networkdriver/bridge"
	"github.com/dotcloud/docker/dockerversion"
	"github.com/dotcloud/docker/engine"
	"github.com/dotcloud/docker/registry"
	"github.com/dotcloud/docker/server"
	"github.com/dotcloud/docker/utils"
)

func Register(eng *engine.Engine) error {
	if err := daemon(eng); err != nil {
		return err
	}
	if err := remote(eng); err != nil {
		return err
	}
	if err := eng.Register("version", dockerVersion); err != nil {
		return err
	}
	return registry.NewService().Install(eng)
}

// remote: a RESTful api for cross-docker communication
func remote(eng *engine.Engine) error {
	if err := eng.Register("serveapi", apiserver.ServeApi); err != nil {
		return err
	}
	return eng.Register("acceptconnections", apiserver.AcceptConnections)
}

// daemon: a default execution and storage backend for Docker on Linux,
// with the following underlying components:
//
// * Pluggable storage drivers including aufs, vfs, lvm and btrfs.
// * Pluggable execution drivers including lxc and chroot.
//
// In practice `daemon` still includes most core Docker components, including:
//
// * The reference registry client implementation
// * Image management
// * The build facility
// * Logging
//
// These components should be broken off into plugins of their own.
//
func daemon(eng *engine.Engine) error {
	if err := eng.Register("initserver", server.InitServer); err != nil {
		return err
	}
	return eng.Register("init_networkdriver", bridge.InitDriver)
}

// builtins jobs independent of any subsystem
func dockerVersion(job *engine.Job) engine.Status {
	v := &engine.Env{}
	v.SetJson("Version", dockerversion.VERSION)
	v.SetJson("ApiVersion", api.APIVERSION)
	v.Set("GitCommit", dockerversion.GITCOMMIT)
	v.Set("GoVersion", runtime.Version())
	v.Set("Os", runtime.GOOS)
	v.Set("Arch", runtime.GOARCH)
	if kernelVersion, err := utils.GetKernelVersion(); err == nil {
		v.Set("KernelVersion", kernelVersion.String())
	}
	if _, err := v.WriteTo(job.Stdout); err != nil {
		return job.Error(err)
	}
	return engine.StatusOK
}
<|endoftext|>


<commit_before>// +build linux

package gopsutil

import (
	"fmt"
	"os/exec"
	"strconv"
	"strings"
)

const (
	SectorSize = 512
)

// Get disk partitions.
// should use setmntent(3) but this implement use /etc/mtab file
func DiskPartitions(all bool) ([]DiskPartitionStat, error) {

	filename := "/etc/mtab"
	lines, err := readLines(filename)
	if err != nil {
		return nil, err
	}

	ret := make([]DiskPartitionStat, 0, len(lines))

	for _, line := range lines {
		fields := strings.Fields(line)
		d := DiskPartitionStat{
			Device:     fields[0],
			Mountpoint: fields[1],
			Fstype:     fields[2],
			Opts:       fields[3],
		}
		ret = append(ret, d)
	}

	return ret, nil
}

func DiskIOCounters() (map[string]DiskIOCountersStat, error) {
	filename := "/proc/diskstats"
	lines, err := readLines(filename)
	if err != nil {
		return nil, err
	}
	ret := make(map[string]DiskIOCountersStat, 0)
	empty := DiskIOCountersStat{}

	for _, line := range lines {
		fields := strings.Fields(line)
		name := fields[2]
		reads, err := strconv.ParseUint((fields[3]), 10, 64)
		if err != nil {
			return ret, err
		}
		rbytes, err := strconv.ParseUint((fields[5]), 10, 64)
		if err != nil {
			return ret, err
		}
		rtime, err := strconv.ParseUint((fields[6]), 10, 64)
		if err != nil {
			return ret, err
		}
		writes, err := strconv.ParseUint((fields[7]), 10, 64)
		if err != nil {
			return ret, err
		}
		wbytes, err := strconv.ParseUint((fields[9]), 10, 64)
		if err != nil {
			return ret, err
		}
		wtime, err := strconv.ParseUint((fields[10]), 10, 64)
		if err != nil {
			return ret, err
		}
		iotime, err := strconv.ParseUint((fields[12]), 10, 64)
		if err != nil {
			return ret, err
		}
		d := DiskIOCountersStat{
			ReadBytes:  rbytes * SectorSize,
			WriteBytes: wbytes * SectorSize,
			ReadCount:  reads,
			WriteCount: writes,
			ReadTime:   rtime,
			WriteTime:  wtime,
			IoTime:     iotime,
		}
		if d == empty {
			continue
		}
		d.Name = name

		d.SerialNumber = GetDiskSerialNumber(name)
		ret[name] = d
	}
	return ret, nil
}

func GetDiskSerialNumber(name string) (string, error) {
	n := fmt.Sprintf("--name=%s", name)
	out, err := exec.Command("/sbin/udevadm", "info", "--query=property", n).Output()

	// does not return error, just an empty string
	if err != nil {
		return "", nil
	}
	lines := strings.Split(string(out), "\n")
	for _, line := range lines {
		values := strings.Split(line, "=")
		if len(values) < 2 || values[0] != "ID_SERIAL" {
			// only get ID_SERIAL, not ID_SERIAL_SHORT
			continue
		}
		return values[1], nil
	}
	return "", nil
}
<commit_msg>sorry, I committed an old file.<commit_after>// +build linux

package gopsutil

import (
	"fmt"
	"os/exec"
	"strconv"
	"strings"
)

const (
	SectorSize = 512
)

// Get disk partitions.
// should use setmntent(3) but this implement use /etc/mtab file
func DiskPartitions(all bool) ([]DiskPartitionStat, error) {

	filename := "/etc/mtab"
	lines, err := readLines(filename)
	if err != nil {
		return nil, err
	}

	ret := make([]DiskPartitionStat, 0, len(lines))

	for _, line := range lines {
		fields := strings.Fields(line)
		d := DiskPartitionStat{
			Device:     fields[0],
			Mountpoint: fields[1],
			Fstype:     fields[2],
			Opts:       fields[3],
		}
		ret = append(ret, d)
	}

	return ret, nil
}

func DiskIOCounters() (map[string]DiskIOCountersStat, error) {
	filename := "/proc/diskstats"
	lines, err := readLines(filename)
	if err != nil {
		return nil, err
	}
	ret := make(map[string]DiskIOCountersStat, 0)
	empty := DiskIOCountersStat{}

	for _, line := range lines {
		fields := strings.Fields(line)
		name := fields[2]
		reads, err := strconv.ParseUint((fields[3]), 10, 64)
		if err != nil {
			return ret, err
		}
		rbytes, err := strconv.ParseUint((fields[5]), 10, 64)
		if err != nil {
			return ret, err
		}
		rtime, err := strconv.ParseUint((fields[6]), 10, 64)
		if err != nil {
			return ret, err
		}
		writes, err := strconv.ParseUint((fields[7]), 10, 64)
		if err != nil {
			return ret, err
		}
		wbytes, err := strconv.ParseUint((fields[9]), 10, 64)
		if err != nil {
			return ret, err
		}
		wtime, err := strconv.ParseUint((fields[10]), 10, 64)
		if err != nil {
			return ret, err
		}
		iotime, err := strconv.ParseUint((fields[12]), 10, 64)
		if err != nil {
			return ret, err
		}
		d := DiskIOCountersStat{
			ReadBytes:  rbytes * SectorSize,
			WriteBytes: wbytes * SectorSize,
			ReadCount:  reads,
			WriteCount: writes,
			ReadTime:   rtime,
			WriteTime:  wtime,
			IoTime:     iotime,
		}
		if d == empty {
			continue
		}
		d.Name = name

		d.SerialNumber = GetDiskSerialNumber(name)
		ret[name] = d
	}
	return ret, nil
}

func GetDiskSerialNumber(name string) string {
	n := fmt.Sprintf("--name=%s", name)
	out, err := exec.Command("/sbin/udevadm", "info", "--query=property", n).Output()

	// does not return error, just an empty string
	if err != nil {
		return ""
	}
	lines := strings.Split(string(out), "\n")
	for _, line := range lines {
		values := strings.Split(line, "=")
		if len(values) < 2 || values[0] != "ID_SERIAL" {
			// only get ID_SERIAL, not ID_SERIAL_SHORT
			continue
		}
		return values[1]
	}
	return ""
}
<|endoftext|>


<commit_before>package agent

import (
	"bytes"
	"encoding/json"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/testutil"
)

func TestRexecWriter(t *testing.T) {
	writer := &rexecWriter{
		BufCh:    make(chan []byte, 16),
		BufSize:  16,
		BufIdle:  10 * time.Millisecond,
		CancelCh: make(chan struct{}),
	}

	// Write short, wait for idle
	start := time.Now()
	n, err := writer.Write([]byte("test"))
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if n != 4 {
		t.Fatalf("bad: %v", n)
	}

	select {
	case b := <-writer.BufCh:
		if len(b) != 4 {
			t.Fatalf("Bad: %v", b)
		}
		if time.Now().Sub(start) < writer.BufIdle {
			t.Fatalf("too early")
		}
	case <-time.After(2 * writer.BufIdle):
		t.Fatalf("timeout")
	}

	// Write in succession to prevent the timeout
	writer.Write([]byte("test"))
	time.Sleep(writer.BufIdle / 2)
	writer.Write([]byte("test"))
	time.Sleep(writer.BufIdle / 2)
	start = time.Now()
	writer.Write([]byte("test"))

	select {
	case b := <-writer.BufCh:
		if len(b) != 12 {
			t.Fatalf("Bad: %v", b)
		}
		if time.Now().Sub(start) < writer.BufIdle {
			t.Fatalf("too early")
		}
	case <-time.After(2 * writer.BufIdle):
		t.Fatalf("timeout")
	}

	// Write large values, multiple flushes required
	writer.Write([]byte("01234567890123456789012345678901"))

	select {
	case b := <-writer.BufCh:
		if string(b) != "0123456789012345" {
			t.Fatalf("bad: %s", b)
		}
	default:
		t.Fatalf("should have buf")
	}
	select {
	case b := <-writer.BufCh:
		if string(b) != "6789012345678901" {
			t.Fatalf("bad: %s", b)
		}
	default:
		t.Fatalf("should have buf")
	}
}

func TestRemoteExecGetSpec(t *testing.T) {
	dir, agent := makeAgent(t, nextConfig())
	defer os.RemoveAll(dir)
	defer agent.Shutdown()
	testutil.WaitForLeader(t, agent.RPC, "dc1")

	event := &remoteExecEvent{
		Prefix:  "_rexec",
		Session: makeRexecSession(t, agent),
	}
	defer destroySession(t, agent, event.Session)

	spec := &remoteExecSpec{
		Command: "uptime",
		Script:  []byte("#!/bin/bash"),
		Wait:    time.Second,
	}
	buf, err := json.Marshal(spec)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	key := "_rexec/" + event.Session + "/job"
	setKV(t, agent, key, buf)

	var out remoteExecSpec
	if !agent.remoteExecGetSpec(event, &out) {
		t.Fatalf("bad")
	}
	if !reflect.DeepEqual(spec, &out) {
		t.Fatalf("bad spec")
	}
}

func TestRemoteExecWrites(t *testing.T) {
	dir, agent := makeAgent(t, nextConfig())
	defer os.RemoveAll(dir)
	defer agent.Shutdown()
	testutil.WaitForLeader(t, agent.RPC, "dc1")

	event := &remoteExecEvent{
		Prefix:  "_rexec",
		Session: makeRexecSession(t, agent),
	}
	defer destroySession(t, agent, event.Session)

	if !agent.remoteExecWriteAck(event) {
		t.Fatalf("bad")
	}

	output := []byte("testing")
	if !agent.remoteExecWriteOutput(event, 0, output) {
		t.Fatalf("bad")
	}
	if !agent.remoteExecWriteOutput(event, 10, output) {
		t.Fatalf("bad")
	}

	exitCode := 1
	if !agent.remoteExecWriteExitCode(event, &exitCode) {
		t.Fatalf("bad")
	}

	key := "_rexec/" + event.Session + "/" + agent.config.NodeName + "/ack"
	d := getKV(t, agent, key)
	if d == nil || d.Session != event.Session {
		t.Fatalf("bad ack: %#v", d)
	}

	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/out/00000"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || !bytes.Equal(d.Value, output) {
		t.Fatalf("bad output: %#v", d)
	}

	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/out/0000a"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || !bytes.Equal(d.Value, output) {
		t.Fatalf("bad output: %#v", d)
	}

	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/exit"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || string(d.Value) != "1" {
		t.Fatalf("bad output: %#v", d)
	}
}

func TestHandleRemoteExec(t *testing.T) {
	dir, agent := makeAgent(t, nextConfig())
	defer os.RemoveAll(dir)
	defer agent.Shutdown()
	testutil.WaitForLeader(t, agent.RPC, "dc1")

	event := &remoteExecEvent{
		Prefix:  "_rexec",
		Session: makeRexecSession(t, agent),
	}
	defer destroySession(t, agent, event.Session)

	spec := &remoteExecSpec{
		Command: "uptime",
		Wait:    time.Second,
	}
	buf, err := json.Marshal(spec)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	key := "_rexec/" + event.Session + "/job"
	setKV(t, agent, key, buf)

	buf, err = json.Marshal(event)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	msg := &UserEvent{
		ID:      generateUUID(),
		Payload: buf,
	}

	// Handle the event...
	agent.handleRemoteExec(msg)

	// Verify we have an ack
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/ack"
	d := getKV(t, agent, key)
	if d == nil || d.Session != event.Session {
		t.Fatalf("bad ack: %#v", d)
	}

	// Verify we have output
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/out/00000"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session ||
		!bytes.Contains(d.Value, []byte("load")) {
		t.Fatalf("bad output: %#v", d)
	}

	// Verify we have an exit code
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/exit"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || string(d.Value) != "0" {
		t.Fatalf("bad output: %#v", d)
	}
}

func TestHandleRemoteExecFailed(t *testing.T) {
	dir, agent := makeAgent(t, nextConfig())
	defer os.RemoveAll(dir)
	defer agent.Shutdown()
	testutil.WaitForLeader(t, agent.RPC, "dc1")

	event := &remoteExecEvent{
		Prefix:  "_rexec",
		Session: makeRexecSession(t, agent),
	}
	defer destroySession(t, agent, event.Session)

	spec := &remoteExecSpec{
		Command: "echo failing;exit 2",
		Wait:    time.Second,
	}
	buf, err := json.Marshal(spec)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	key := "_rexec/" + event.Session + "/job"
	setKV(t, agent, key, buf)

	buf, err = json.Marshal(event)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	msg := &UserEvent{
		ID:      generateUUID(),
		Payload: buf,
	}

	// Handle the event...
	agent.handleRemoteExec(msg)

	// Verify we have an ack
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/ack"
	d := getKV(t, agent, key)
	if d == nil || d.Session != event.Session {
		t.Fatalf("bad ack: %#v", d)
	}

	// Verify we have output
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/out/00000"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session ||
		!bytes.Contains(d.Value, []byte("failing")) {
		t.Fatalf("bad output: %#v", d)
	}

	// Verify we have an exit code
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/exit"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || string(d.Value) != "2" {
		t.Fatalf("bad output: %#v", d)
	}
}

func makeRexecSession(t *testing.T, agent *Agent) string {
	args := structs.SessionRequest{
		Datacenter: agent.config.Datacenter,
		Op:         structs.SessionCreate,
		Session: structs.Session{
			Node:      agent.config.NodeName,
			LockDelay: 15 * time.Second,
		},
	}
	var out string
	if err := agent.RPC("Session.Apply", &args, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	return out
}

func destroySession(t *testing.T, agent *Agent, session string) {
	args := structs.SessionRequest{
		Datacenter: agent.config.Datacenter,
		Op:         structs.SessionDestroy,
		Session: structs.Session{
			ID: session,
		},
	}
	var out string
	if err := agent.RPC("Session.Apply", &args, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func setKV(t *testing.T, agent *Agent, key string, val []byte) {
	write := structs.KVSRequest{
		Datacenter: agent.config.Datacenter,
		Op:         structs.KVSSet,
		DirEnt: structs.DirEntry{
			Key:   key,
			Value: val,
		},
	}
	var success bool
	if err := agent.RPC("KVS.Apply", &write, &success); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func getKV(t *testing.T, agent *Agent, key string) *structs.DirEntry {
	req := structs.KeyRequest{
		Datacenter: agent.config.Datacenter,
		Key:        key,
	}
	var out structs.IndexedDirEntries
	if err := agent.RPC("KVS.Get", &req, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	if len(out.Entries) > 0 {
		return out.Entries[0]
	}
	return nil
}
<commit_msg>Remove test code duplication<commit_after>package agent

import (
	"bytes"
	"encoding/json"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/testutil"
)

func TestRexecWriter(t *testing.T) {
	writer := &rexecWriter{
		BufCh:    make(chan []byte, 16),
		BufSize:  16,
		BufIdle:  10 * time.Millisecond,
		CancelCh: make(chan struct{}),
	}

	// Write short, wait for idle
	start := time.Now()
	n, err := writer.Write([]byte("test"))
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if n != 4 {
		t.Fatalf("bad: %v", n)
	}

	select {
	case b := <-writer.BufCh:
		if len(b) != 4 {
			t.Fatalf("Bad: %v", b)
		}
		if time.Now().Sub(start) < writer.BufIdle {
			t.Fatalf("too early")
		}
	case <-time.After(2 * writer.BufIdle):
		t.Fatalf("timeout")
	}

	// Write in succession to prevent the timeout
	writer.Write([]byte("test"))
	time.Sleep(writer.BufIdle / 2)
	writer.Write([]byte("test"))
	time.Sleep(writer.BufIdle / 2)
	start = time.Now()
	writer.Write([]byte("test"))

	select {
	case b := <-writer.BufCh:
		if len(b) != 12 {
			t.Fatalf("Bad: %v", b)
		}
		if time.Now().Sub(start) < writer.BufIdle {
			t.Fatalf("too early")
		}
	case <-time.After(2 * writer.BufIdle):
		t.Fatalf("timeout")
	}

	// Write large values, multiple flushes required
	writer.Write([]byte("01234567890123456789012345678901"))

	select {
	case b := <-writer.BufCh:
		if string(b) != "0123456789012345" {
			t.Fatalf("bad: %s", b)
		}
	default:
		t.Fatalf("should have buf")
	}
	select {
	case b := <-writer.BufCh:
		if string(b) != "6789012345678901" {
			t.Fatalf("bad: %s", b)
		}
	default:
		t.Fatalf("should have buf")
	}
}

func TestRemoteExecGetSpec(t *testing.T) {
	dir, agent := makeAgent(t, nextConfig())
	defer os.RemoveAll(dir)
	defer agent.Shutdown()
	testutil.WaitForLeader(t, agent.RPC, "dc1")

	event := &remoteExecEvent{
		Prefix:  "_rexec",
		Session: makeRexecSession(t, agent),
	}
	defer destroySession(t, agent, event.Session)

	spec := &remoteExecSpec{
		Command: "uptime",
		Script:  []byte("#!/bin/bash"),
		Wait:    time.Second,
	}
	buf, err := json.Marshal(spec)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	key := "_rexec/" + event.Session + "/job"
	setKV(t, agent, key, buf)

	var out remoteExecSpec
	if !agent.remoteExecGetSpec(event, &out) {
		t.Fatalf("bad")
	}
	if !reflect.DeepEqual(spec, &out) {
		t.Fatalf("bad spec")
	}
}

func TestRemoteExecWrites(t *testing.T) {
	dir, agent := makeAgent(t, nextConfig())
	defer os.RemoveAll(dir)
	defer agent.Shutdown()
	testutil.WaitForLeader(t, agent.RPC, "dc1")

	event := &remoteExecEvent{
		Prefix:  "_rexec",
		Session: makeRexecSession(t, agent),
	}
	defer destroySession(t, agent, event.Session)

	if !agent.remoteExecWriteAck(event) {
		t.Fatalf("bad")
	}

	output := []byte("testing")
	if !agent.remoteExecWriteOutput(event, 0, output) {
		t.Fatalf("bad")
	}
	if !agent.remoteExecWriteOutput(event, 10, output) {
		t.Fatalf("bad")
	}

	exitCode := 1
	if !agent.remoteExecWriteExitCode(event, &exitCode) {
		t.Fatalf("bad")
	}

	key := "_rexec/" + event.Session + "/" + agent.config.NodeName + "/ack"
	d := getKV(t, agent, key)
	if d == nil || d.Session != event.Session {
		t.Fatalf("bad ack: %#v", d)
	}

	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/out/00000"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || !bytes.Equal(d.Value, output) {
		t.Fatalf("bad output: %#v", d)
	}

	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/out/0000a"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || !bytes.Equal(d.Value, output) {
		t.Fatalf("bad output: %#v", d)
	}

	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/exit"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || string(d.Value) != "1" {
		t.Fatalf("bad output: %#v", d)
	}
}

func _TestHandleRemoteExec(t *testing.T, command string, expectedSubstring string, expectedReturnCode string) {
	dir, agent := makeAgent(t, nextConfig())
	defer os.RemoveAll(dir)
	defer agent.Shutdown()
	testutil.WaitForLeader(t, agent.RPC, "dc1")

	event := &remoteExecEvent{
		Prefix:  "_rexec",
		Session: makeRexecSession(t, agent),
	}
	defer destroySession(t, agent, event.Session)

	spec := &remoteExecSpec{
		Command: command,
		Wait:    time.Second,
	}
	buf, err := json.Marshal(spec)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	key := "_rexec/" + event.Session + "/job"
	setKV(t, agent, key, buf)

	buf, err = json.Marshal(event)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	msg := &UserEvent{
		ID:      generateUUID(),
		Payload: buf,
	}

	// Handle the event...
	agent.handleRemoteExec(msg)

	// Verify we have an ack
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/ack"
	d := getKV(t, agent, key)
	if d == nil || d.Session != event.Session {
		t.Fatalf("bad ack: %#v", d)
	}

	// Verify we have output
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/out/00000"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session ||
		!bytes.Contains(d.Value, []byte(expectedSubstring)) {
		t.Fatalf("bad output: %#v", d)
	}

	// Verify we have an exit code
	key = "_rexec/" + event.Session + "/" + agent.config.NodeName + "/exit"
	d = getKV(t, agent, key)
	if d == nil || d.Session != event.Session || string(d.Value) != expectedReturnCode {
		t.Fatalf("bad output: %#v", d)
	}
}

func TestHandleRemoteExec(t *testing.T) {
	_TestHandleRemoteExec(t, "uptime", "load", "0")
}

func TestHandleRemoteExecFailed(t *testing.T) {
	_TestHandleRemoteExec(t, "echo failing;exit 2", "failing", "2")
}

func makeRexecSession(t *testing.T, agent *Agent) string {
	args := structs.SessionRequest{
		Datacenter: agent.config.Datacenter,
		Op:         structs.SessionCreate,
		Session: structs.Session{
			Node:      agent.config.NodeName,
			LockDelay: 15 * time.Second,
		},
	}
	var out string
	if err := agent.RPC("Session.Apply", &args, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	return out
}

func destroySession(t *testing.T, agent *Agent, session string) {
	args := structs.SessionRequest{
		Datacenter: agent.config.Datacenter,
		Op:         structs.SessionDestroy,
		Session: structs.Session{
			ID: session,
		},
	}
	var out string
	if err := agent.RPC("Session.Apply", &args, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func setKV(t *testing.T, agent *Agent, key string, val []byte) {
	write := structs.KVSRequest{
		Datacenter: agent.config.Datacenter,
		Op:         structs.KVSSet,
		DirEnt: structs.DirEntry{
			Key:   key,
			Value: val,
		},
	}
	var success bool
	if err := agent.RPC("KVS.Apply", &write, &success); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func getKV(t *testing.T, agent *Agent, key string) *structs.DirEntry {
	req := structs.KeyRequest{
		Datacenter: agent.config.Datacenter,
		Key:        key,
	}
	var out structs.IndexedDirEntries
	if err := agent.RPC("KVS.Get", &req, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	if len(out.Entries) > 0 {
		return out.Entries[0]
	}
	return nil
}
<|endoftext|>


<commit_before>package goproxy

import (
	"bytes"
	"io/ioutil"
	"net"
	"net/http"
	"regexp"
	"strings"
)

// ReqCondition.HandleReq will decide whether or not to use the ReqHandler on an HTTP request
// before sending it to the remote server
type ReqCondition interface {
	RespCondition
	HandleReq(req *http.Request, ctx *ProxyCtx) bool
}

// RespCondition.HandleReq will decide whether or not to use the RespHandler on an HTTP response
// before sending it to the proxy client. Note that resp might be nil, in case there was an
// error sending the request.
type RespCondition interface {
	HandleResp(resp *http.Response, ctx *ProxyCtx) bool
}

// ReqConditionFunc.HandleReq(req,ctx) <=> ReqConditionFunc(req,ctx)
type ReqConditionFunc func(req *http.Request, ctx *ProxyCtx) bool

// RespConditionFunc.HandleResp(resp,ctx) <=> RespConditionFunc(resp,ctx)
type RespConditionFunc func(resp *http.Response, ctx *ProxyCtx) bool

func (c ReqConditionFunc) HandleReq(req *http.Request, ctx *ProxyCtx) bool {
	return c(req, ctx)
}

// ReqConditionFunc cannot test responses. It only satisfies RespCondition interface so that
// to be usable as RespCondition.
func (c ReqConditionFunc) HandleResp(resp *http.Response, ctx *ProxyCtx) bool {
	return c(ctx.Req, ctx)
}

func (c RespConditionFunc) HandleResp(resp *http.Response, ctx *ProxyCtx) bool {
	return c(resp, ctx)
}

// UrlHasPrefix returns a ReqCondition checking wether the destination URL the proxy client has requested
// has the given prefix, with or without the host.
// For example UrlHasPrefix("host/x") will match requests of the form 'GET host/x', and will match
// requests to url 'http://host/x'
func UrlHasPrefix(prefix string) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		return strings.HasPrefix(req.URL.Path, prefix) ||
			strings.HasPrefix(req.URL.Host+req.URL.Path, prefix) ||
			strings.HasPrefix(req.URL.Scheme+req.URL.Host+req.URL.Path, prefix)
	}
}

// UrlIs returns a ReqCondition, testing whether or not the request URL is one of the given strings
// with or without the host prefix.
// UrlIs("google.com/","foo") will match requests 'GET /' to 'google.com', requests `'GET google.com/' to
// any host, and requests of the form 'GET foo'.
func UrlIs(urls ...string) ReqConditionFunc {
	urlSet := make(map[string]bool)
	for _, u := range urls {
		urlSet[u] = true
	}
	return func(req *http.Request, ctx *ProxyCtx) bool {
		_, pathOk := urlSet[req.URL.Path]
		_, hostAndOk := urlSet[req.URL.Host+req.URL.Path]
		return pathOk || hostAndOk
	}
}

// ReqHostMatches returns a ReqCondition, testing whether the host to which the request was directed to matches
// any of the given regular expressions.
func ReqHostMatches(regexps ...*regexp.Regexp) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		for _, re := range regexps {
			if re.MatchString(req.Host) {
				return true
			}
		}
		return false
	}
}

// ReqHostIs returns a ReqCondition, testing whether the host to which the request is directed to equal
// to one of the given strings
func ReqHostIs(hosts ...string) ReqConditionFunc {
	hostSet := make(map[string]bool)
	for _, h := range hosts {
		hostSet[h] = true
	}
	return func(req *http.Request, ctx *ProxyCtx) bool {
		_, ok := hostSet[req.URL.Host]
		return ok
	}
}

var localHostIpv4 = regexp.MustCompile(`127\.0\.0\.\d+`)

// IsLocalHost checks whether the destination host is explicitly local host
// (buggy, there can be IPv6 addresses it doesn't catch)
var IsLocalHost ReqConditionFunc = func(req *http.Request, ctx *ProxyCtx) bool {
	return req.URL.Host == "::1" ||
		req.URL.Host == "0:0:0:0:0:0:0:1" ||
		localHostIpv4.MatchString(req.URL.Host) ||
		req.URL.Host == "localhost"
}

// UrlMatches returns a ReqCondition testing whether the destination URL
// of the request matches the given regexp, with or without prefix
func UrlMatches(re *regexp.Regexp) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		return re.MatchString(req.URL.Path) ||
			re.MatchString(req.URL.Host+req.URL.Path)
	}
}

// DstHostIs returns a ReqCondition testing wether the host in the request url is the given string
func DstHostIs(host string) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		return req.URL.Host == host
	}
}

// SrcIpIs returns a ReqCondition testing whether the source IP of the request is one of the given strings
func SrcIpIs(ips ...string) ReqCondition {
	return ReqConditionFunc(func(req *http.Request, ctx *ProxyCtx) bool {
		for _, ip := range ips {
			if strings.HasPrefix(req.RemoteAddr, ip+":") {
				return true
			}
		}
		return false
	})
}

// Not returns a ReqCondition negating the given ReqCondition
func Not(r ReqCondition) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		return !r.HandleReq(req, ctx)
	}
}

// ContentTypeIs returns a RespCondition testing whether the HTTP response has Content-Type header equal
// to one of the given strings.
func ContentTypeIs(typ string, types ...string) RespCondition {
	types = append(types, typ)
	return RespConditionFunc(func(resp *http.Response, ctx *ProxyCtx) bool {
		if resp == nil {
			return false
		}
		contentType := resp.Header.Get("Content-Type")
		for _, typ := range types {
			if contentType == typ || strings.HasPrefix(contentType, typ+";") {
				return true
			}
		}
		return false
	})
}

// StatusCodeIs returns a RespCondition, testing whether or not the HTTP status
// code is one of the given ints
func StatusCodeIs(codes ...int) goproxy.RespCondition {
	codeSet := make(map[int]bool)
	for _, c := range codes {
		codeSet[c] = true
	}
	return goproxy.RespConditionFunc(func(resp *http.Response, ctx *goproxy.ProxyCtx) bool {
		if resp == nil {
			return false
		}
		_, codeMatch := codeSet[resp.StatusCode]
		return codeMatch
	})
}

// ProxyHttpServer.OnRequest Will return a temporary ReqProxyConds struct, aggregating the given condtions.
// You will use the ReqProxyConds struct to register a ReqHandler, that would filter
// the request, only if all the given ReqCondition matched.
// Typical usage:
//	proxy.OnRequest(UrlIs("example.com/foo"),UrlMatches(regexp.MustParse(`.*\.exampl.\com\./.*`)).Do(...)
func (proxy *ProxyHttpServer) OnRequest(conds ...ReqCondition) *ReqProxyConds {
	return &ReqProxyConds{proxy, conds}
}

// ReqProxyConds aggregate ReqConditions for a ProxyHttpServer. Upon calling Do, it will register a ReqHandler that would
// handle the request if all conditions on the HTTP request are met.
type ReqProxyConds struct {
	proxy    *ProxyHttpServer
	reqConds []ReqCondition
}

// DoFunc is equivalent to proxy.OnRequest().Do(FuncReqHandler(f))
func (pcond *ReqProxyConds) DoFunc(f func(req *http.Request, ctx *ProxyCtx) (*http.Request, *http.Response)) {
	pcond.Do(FuncReqHandler(f))
}

// ReqProxyConds.Do will register the ReqHandler on the proxy,
// the ReqHandler will handle the HTTP request if all the conditions
// aggregated in the ReqProxyConds are met. Typical usage:
//	proxy.OnRequest().Do(handler) // will call handler.Handle(req,ctx) on every request to the proxy
//	proxy.OnRequest(cond1,cond2).Do(handler)
//	// given request to the proxy, will test if cond1.HandleReq(req,ctx) && cond2.HandleReq(req,ctx) are true
//	// if they are, will call handler.Handle(req,ctx)
func (pcond *ReqProxyConds) Do(h ReqHandler) {
	pcond.proxy.reqHandlers = append(pcond.proxy.reqHandlers,
		FuncReqHandler(func(r *http.Request, ctx *ProxyCtx) (*http.Request, *http.Response) {
			for _, cond := range pcond.reqConds {
				if !cond.HandleReq(r, ctx) {
					return r, nil
				}
			}
			return h.Handle(r, ctx)
		}))
}

// HandleConnect is used when proxy receives an HTTP CONNECT request,
// it'll then use the HttpsHandler to determine what should it
// do with this request. The handler returns a ConnectAction struct, the Action field in the ConnectAction
// struct returned will determine what to do with this request. ConnectAccept will simply accept the request
// forwarding all bytes from the client to the remote host, ConnectReject will close the connection with the
// client, and ConnectMitm, will assume the underlying connection is an HTTPS connection, and will use Man
// in the Middle attack to eavesdrop the connection. All regular handler will be active on this eavesdropped
// connection.
// The ConnectAction struct contains possible tlsConfig that will be used for eavesdropping. If nil, the proxy
// will use the default tls configuration.
//	proxy.OnRequest().HandleConnect(goproxy.AlwaysReject) // rejects all CONNECT requests
func (pcond *ReqProxyConds) HandleConnect(h HttpsHandler) {
	pcond.proxy.httpsHandlers = append(pcond.proxy.httpsHandlers,
		FuncHttpsHandler(func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
			for _, cond := range pcond.reqConds {
				if !cond.HandleReq(ctx.Req, ctx) {
					return nil, ""
				}
			}
			return h.HandleConnect(host, ctx)
		}))
}

// HandleConnectFunc is equivalent to HandleConnect,
// for example, accepting CONNECT request if they contain a password in header
//	io.WriteString(h,password)
//	passHash := h.Sum(nil)
//	proxy.OnRequest().HandleConnectFunc(func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
//		c := sha1.New()
//		io.WriteString(c,ctx.Req.Header.Get("X-GoProxy-Auth"))
//		if c.Sum(nil) == passHash {
//			return OkConnect, host
//		}
//		return RejectConnect, host
//	})
func (pcond *ReqProxyConds) HandleConnectFunc(f func(host string, ctx *ProxyCtx) (*ConnectAction, string)) {
	pcond.HandleConnect(FuncHttpsHandler(f))
}

func (pcond *ReqProxyConds) HijackConnect(f func(req *http.Request, client net.Conn, ctx *ProxyCtx)) {
	pcond.proxy.httpsHandlers = append(pcond.proxy.httpsHandlers,
		FuncHttpsHandler(func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
			for _, cond := range pcond.reqConds {
				if !cond.HandleReq(ctx.Req, ctx) {
					return nil, ""
				}
			}
			return &ConnectAction{Action: ConnectHijack, Hijack: f}, host
		}))
}

// ProxyConds is used to aggregate RespConditions for a ProxyHttpServer.
// Upon calling ProxyConds.Do, it will register a RespHandler that would
// handle the HTTP response from remote server if all conditions on the HTTP response are met.
type ProxyConds struct {
	proxy    *ProxyHttpServer
	reqConds []ReqCondition
	respCond []RespCondition
}

// ProxyConds.DoFunc is equivalent to proxy.OnResponse().Do(FuncRespHandler(f))
func (pcond *ProxyConds) DoFunc(f func(resp *http.Response, ctx *ProxyCtx) *http.Response) {
	pcond.Do(FuncRespHandler(f))
}

// ProxyConds.Do will register the RespHandler on the proxy, h.Handle(resp,ctx) will be called on every
// request that matches the conditions aggregated in pcond.
func (pcond *ProxyConds) Do(h RespHandler) {
	pcond.proxy.respHandlers = append(pcond.proxy.respHandlers,
		FuncRespHandler(func(resp *http.Response, ctx *ProxyCtx) *http.Response {
			for _, cond := range pcond.reqConds {
				if !cond.HandleReq(ctx.Req, ctx) {
					return resp
				}
			}
			for _, cond := range pcond.respCond {
				if !cond.HandleResp(resp, ctx) {
					return resp
				}
			}
			return h.Handle(resp, ctx)
		}))
}

// OnResponse is used when adding a response-filter to the HTTP proxy, usual pattern is
//	proxy.OnResponse(cond1,cond2).Do(handler) // handler.Handle(resp,ctx) will be used
//				// if cond1.HandleResp(resp) && cond2.HandleResp(resp)
func (proxy *ProxyHttpServer) OnResponse(conds ...RespCondition) *ProxyConds {
	return &ProxyConds{proxy, make([]ReqCondition, 0), conds}
}

// AlwaysMitm is a HttpsHandler that always eavesdrop https connections, for example to
// eavesdrop all https connections to www.google.com, we can use
//	proxy.OnRequest(goproxy.ReqHostIs("www.google.com")).HandleConnect(goproxy.AlwaysMitm)
var AlwaysMitm FuncHttpsHandler = func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
	return MitmConnect, host
}

// AlwaysReject is a HttpsHandler that drops any CONNECT request, for example, this code will disallow
// connections to hosts on any other port than 443
//	proxy.OnRequest(goproxy.Not(goproxy.ReqHostMatches(regexp.MustCompile(":443$"))).
//		HandleConnect(goproxy.AlwaysReject)
var AlwaysReject FuncHttpsHandler = func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
	return RejectConnect, host
}

// HandleBytes will return a RespHandler that read the entire body of the request
// to a byte array in memory, would run the user supplied f function on the byte arra,
// and will replace the body of the original response with the resulting byte array.
func HandleBytes(f func(b []byte, ctx *ProxyCtx) []byte) RespHandler {
	return FuncRespHandler(func(resp *http.Response, ctx *ProxyCtx) *http.Response {
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			ctx.Warnf("Cannot read response %s", err)
			return resp
		}
		resp.Body.Close()

		resp.Body = ioutil.NopCloser(bytes.NewBuffer(f(b, ctx)))
		return resp
	})
}
<commit_msg>Forgot that this isn't a library like it was in my fork<commit_after>package goproxy

import (
	"bytes"
	"io/ioutil"
	"net"
	"net/http"
	"regexp"
	"strings"
)

// ReqCondition.HandleReq will decide whether or not to use the ReqHandler on an HTTP request
// before sending it to the remote server
type ReqCondition interface {
	RespCondition
	HandleReq(req *http.Request, ctx *ProxyCtx) bool
}

// RespCondition.HandleReq will decide whether or not to use the RespHandler on an HTTP response
// before sending it to the proxy client. Note that resp might be nil, in case there was an
// error sending the request.
type RespCondition interface {
	HandleResp(resp *http.Response, ctx *ProxyCtx) bool
}

// ReqConditionFunc.HandleReq(req,ctx) <=> ReqConditionFunc(req,ctx)
type ReqConditionFunc func(req *http.Request, ctx *ProxyCtx) bool

// RespConditionFunc.HandleResp(resp,ctx) <=> RespConditionFunc(resp,ctx)
type RespConditionFunc func(resp *http.Response, ctx *ProxyCtx) bool

func (c ReqConditionFunc) HandleReq(req *http.Request, ctx *ProxyCtx) bool {
	return c(req, ctx)
}

// ReqConditionFunc cannot test responses. It only satisfies RespCondition interface so that
// to be usable as RespCondition.
func (c ReqConditionFunc) HandleResp(resp *http.Response, ctx *ProxyCtx) bool {
	return c(ctx.Req, ctx)
}

func (c RespConditionFunc) HandleResp(resp *http.Response, ctx *ProxyCtx) bool {
	return c(resp, ctx)
}

// UrlHasPrefix returns a ReqCondition checking wether the destination URL the proxy client has requested
// has the given prefix, with or without the host.
// For example UrlHasPrefix("host/x") will match requests of the form 'GET host/x', and will match
// requests to url 'http://host/x'
func UrlHasPrefix(prefix string) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		return strings.HasPrefix(req.URL.Path, prefix) ||
			strings.HasPrefix(req.URL.Host+req.URL.Path, prefix) ||
			strings.HasPrefix(req.URL.Scheme+req.URL.Host+req.URL.Path, prefix)
	}
}

// UrlIs returns a ReqCondition, testing whether or not the request URL is one of the given strings
// with or without the host prefix.
// UrlIs("google.com/","foo") will match requests 'GET /' to 'google.com', requests `'GET google.com/' to
// any host, and requests of the form 'GET foo'.
func UrlIs(urls ...string) ReqConditionFunc {
	urlSet := make(map[string]bool)
	for _, u := range urls {
		urlSet[u] = true
	}
	return func(req *http.Request, ctx *ProxyCtx) bool {
		_, pathOk := urlSet[req.URL.Path]
		_, hostAndOk := urlSet[req.URL.Host+req.URL.Path]
		return pathOk || hostAndOk
	}
}

// ReqHostMatches returns a ReqCondition, testing whether the host to which the request was directed to matches
// any of the given regular expressions.
func ReqHostMatches(regexps ...*regexp.Regexp) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		for _, re := range regexps {
			if re.MatchString(req.Host) {
				return true
			}
		}
		return false
	}
}

// ReqHostIs returns a ReqCondition, testing whether the host to which the request is directed to equal
// to one of the given strings
func ReqHostIs(hosts ...string) ReqConditionFunc {
	hostSet := make(map[string]bool)
	for _, h := range hosts {
		hostSet[h] = true
	}
	return func(req *http.Request, ctx *ProxyCtx) bool {
		_, ok := hostSet[req.URL.Host]
		return ok
	}
}

var localHostIpv4 = regexp.MustCompile(`127\.0\.0\.\d+`)

// IsLocalHost checks whether the destination host is explicitly local host
// (buggy, there can be IPv6 addresses it doesn't catch)
var IsLocalHost ReqConditionFunc = func(req *http.Request, ctx *ProxyCtx) bool {
	return req.URL.Host == "::1" ||
		req.URL.Host == "0:0:0:0:0:0:0:1" ||
		localHostIpv4.MatchString(req.URL.Host) ||
		req.URL.Host == "localhost"
}

// UrlMatches returns a ReqCondition testing whether the destination URL
// of the request matches the given regexp, with or without prefix
func UrlMatches(re *regexp.Regexp) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		return re.MatchString(req.URL.Path) ||
			re.MatchString(req.URL.Host+req.URL.Path)
	}
}

// DstHostIs returns a ReqCondition testing wether the host in the request url is the given string
func DstHostIs(host string) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		return req.URL.Host == host
	}
}

// SrcIpIs returns a ReqCondition testing whether the source IP of the request is one of the given strings
func SrcIpIs(ips ...string) ReqCondition {
	return ReqConditionFunc(func(req *http.Request, ctx *ProxyCtx) bool {
		for _, ip := range ips {
			if strings.HasPrefix(req.RemoteAddr, ip+":") {
				return true
			}
		}
		return false
	})
}

// Not returns a ReqCondition negating the given ReqCondition
func Not(r ReqCondition) ReqConditionFunc {
	return func(req *http.Request, ctx *ProxyCtx) bool {
		return !r.HandleReq(req, ctx)
	}
}

// ContentTypeIs returns a RespCondition testing whether the HTTP response has Content-Type header equal
// to one of the given strings.
func ContentTypeIs(typ string, types ...string) RespCondition {
	types = append(types, typ)
	return RespConditionFunc(func(resp *http.Response, ctx *ProxyCtx) bool {
		if resp == nil {
			return false
		}
		contentType := resp.Header.Get("Content-Type")
		for _, typ := range types {
			if contentType == typ || strings.HasPrefix(contentType, typ+";") {
				return true
			}
		}
		return false
	})
}

// StatusCodeIs returns a RespCondition, testing whether or not the HTTP status
// code is one of the given ints
func StatusCodeIs(codes ...int) RespCondition {
	codeSet := make(map[int]bool)
	for _, c := range codes {
		codeSet[c] = true
	}
	return RespConditionFunc(func(resp *http.Response, ctx *ProxyCtx) bool {
		if resp == nil {
			return false
		}
		_, codeMatch := codeSet[resp.StatusCode]
		return codeMatch
	})
}

// ProxyHttpServer.OnRequest Will return a temporary ReqProxyConds struct, aggregating the given condtions.
// You will use the ReqProxyConds struct to register a ReqHandler, that would filter
// the request, only if all the given ReqCondition matched.
// Typical usage:
//	proxy.OnRequest(UrlIs("example.com/foo"),UrlMatches(regexp.MustParse(`.*\.exampl.\com\./.*`)).Do(...)
func (proxy *ProxyHttpServer) OnRequest(conds ...ReqCondition) *ReqProxyConds {
	return &ReqProxyConds{proxy, conds}
}

// ReqProxyConds aggregate ReqConditions for a ProxyHttpServer. Upon calling Do, it will register a ReqHandler that would
// handle the request if all conditions on the HTTP request are met.
type ReqProxyConds struct {
	proxy    *ProxyHttpServer
	reqConds []ReqCondition
}

// DoFunc is equivalent to proxy.OnRequest().Do(FuncReqHandler(f))
func (pcond *ReqProxyConds) DoFunc(f func(req *http.Request, ctx *ProxyCtx) (*http.Request, *http.Response)) {
	pcond.Do(FuncReqHandler(f))
}

// ReqProxyConds.Do will register the ReqHandler on the proxy,
// the ReqHandler will handle the HTTP request if all the conditions
// aggregated in the ReqProxyConds are met. Typical usage:
//	proxy.OnRequest().Do(handler) // will call handler.Handle(req,ctx) on every request to the proxy
//	proxy.OnRequest(cond1,cond2).Do(handler)
//	// given request to the proxy, will test if cond1.HandleReq(req,ctx) && cond2.HandleReq(req,ctx) are true
//	// if they are, will call handler.Handle(req,ctx)
func (pcond *ReqProxyConds) Do(h ReqHandler) {
	pcond.proxy.reqHandlers = append(pcond.proxy.reqHandlers,
		FuncReqHandler(func(r *http.Request, ctx *ProxyCtx) (*http.Request, *http.Response) {
			for _, cond := range pcond.reqConds {
				if !cond.HandleReq(r, ctx) {
					return r, nil
				}
			}
			return h.Handle(r, ctx)
		}))
}

// HandleConnect is used when proxy receives an HTTP CONNECT request,
// it'll then use the HttpsHandler to determine what should it
// do with this request. The handler returns a ConnectAction struct, the Action field in the ConnectAction
// struct returned will determine what to do with this request. ConnectAccept will simply accept the request
// forwarding all bytes from the client to the remote host, ConnectReject will close the connection with the
// client, and ConnectMitm, will assume the underlying connection is an HTTPS connection, and will use Man
// in the Middle attack to eavesdrop the connection. All regular handler will be active on this eavesdropped
// connection.
// The ConnectAction struct contains possible tlsConfig that will be used for eavesdropping. If nil, the proxy
// will use the default tls configuration.
//	proxy.OnRequest().HandleConnect(goproxy.AlwaysReject) // rejects all CONNECT requests
func (pcond *ReqProxyConds) HandleConnect(h HttpsHandler) {
	pcond.proxy.httpsHandlers = append(pcond.proxy.httpsHandlers,
		FuncHttpsHandler(func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
			for _, cond := range pcond.reqConds {
				if !cond.HandleReq(ctx.Req, ctx) {
					return nil, ""
				}
			}
			return h.HandleConnect(host, ctx)
		}))
}

// HandleConnectFunc is equivalent to HandleConnect,
// for example, accepting CONNECT request if they contain a password in header
//	io.WriteString(h,password)
//	passHash := h.Sum(nil)
//	proxy.OnRequest().HandleConnectFunc(func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
//		c := sha1.New()
//		io.WriteString(c,ctx.Req.Header.Get("X-GoProxy-Auth"))
//		if c.Sum(nil) == passHash {
//			return OkConnect, host
//		}
//		return RejectConnect, host
//	})
func (pcond *ReqProxyConds) HandleConnectFunc(f func(host string, ctx *ProxyCtx) (*ConnectAction, string)) {
	pcond.HandleConnect(FuncHttpsHandler(f))
}

func (pcond *ReqProxyConds) HijackConnect(f func(req *http.Request, client net.Conn, ctx *ProxyCtx)) {
	pcond.proxy.httpsHandlers = append(pcond.proxy.httpsHandlers,
		FuncHttpsHandler(func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
			for _, cond := range pcond.reqConds {
				if !cond.HandleReq(ctx.Req, ctx) {
					return nil, ""
				}
			}
			return &ConnectAction{Action: ConnectHijack, Hijack: f}, host
		}))
}

// ProxyConds is used to aggregate RespConditions for a ProxyHttpServer.
// Upon calling ProxyConds.Do, it will register a RespHandler that would
// handle the HTTP response from remote server if all conditions on the HTTP response are met.
type ProxyConds struct {
	proxy    *ProxyHttpServer
	reqConds []ReqCondition
	respCond []RespCondition
}

// ProxyConds.DoFunc is equivalent to proxy.OnResponse().Do(FuncRespHandler(f))
func (pcond *ProxyConds) DoFunc(f func(resp *http.Response, ctx *ProxyCtx) *http.Response) {
	pcond.Do(FuncRespHandler(f))
}

// ProxyConds.Do will register the RespHandler on the proxy, h.Handle(resp,ctx) will be called on every
// request that matches the conditions aggregated in pcond.
func (pcond *ProxyConds) Do(h RespHandler) {
	pcond.proxy.respHandlers = append(pcond.proxy.respHandlers,
		FuncRespHandler(func(resp *http.Response, ctx *ProxyCtx) *http.Response {
			for _, cond := range pcond.reqConds {
				if !cond.HandleReq(ctx.Req, ctx) {
					return resp
				}
			}
			for _, cond := range pcond.respCond {
				if !cond.HandleResp(resp, ctx) {
					return resp
				}
			}
			return h.Handle(resp, ctx)
		}))
}

// OnResponse is used when adding a response-filter to the HTTP proxy, usual pattern is
//	proxy.OnResponse(cond1,cond2).Do(handler) // handler.Handle(resp,ctx) will be used
//				// if cond1.HandleResp(resp) && cond2.HandleResp(resp)
func (proxy *ProxyHttpServer) OnResponse(conds ...RespCondition) *ProxyConds {
	return &ProxyConds{proxy, make([]ReqCondition, 0), conds}
}

// AlwaysMitm is a HttpsHandler that always eavesdrop https connections, for example to
// eavesdrop all https connections to www.google.com, we can use
//	proxy.OnRequest(goproxy.ReqHostIs("www.google.com")).HandleConnect(goproxy.AlwaysMitm)
var AlwaysMitm FuncHttpsHandler = func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
	return MitmConnect, host
}

// AlwaysReject is a HttpsHandler that drops any CONNECT request, for example, this code will disallow
// connections to hosts on any other port than 443
//	proxy.OnRequest(goproxy.Not(goproxy.ReqHostMatches(regexp.MustCompile(":443$"))).
//		HandleConnect(goproxy.AlwaysReject)
var AlwaysReject FuncHttpsHandler = func(host string, ctx *ProxyCtx) (*ConnectAction, string) {
	return RejectConnect, host
}

// HandleBytes will return a RespHandler that read the entire body of the request
// to a byte array in memory, would run the user supplied f function on the byte arra,
// and will replace the body of the original response with the resulting byte array.
func HandleBytes(f func(b []byte, ctx *ProxyCtx) []byte) RespHandler {
	return FuncRespHandler(func(resp *http.Response, ctx *ProxyCtx) *http.Response {
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			ctx.Warnf("Cannot read response %s", err)
			return resp
		}
		resp.Body.Close()

		resp.Body = ioutil.NopCloser(bytes.NewBuffer(f(b, ctx)))
		return resp
	})
}
<|endoftext|>


<commit_before>package events

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/fsouza/go-dockerclient"
	"github.com/rancherio/go-machine-service/locks"
)

const RancherIPLabelKey = "io.rancher.container.ip"
const RancherSystemLabelKey = "io.rancher.container.system"
const RancherIPEnvKey = "RANCHER_IP="
const RancherNameserver = "169.254.169.250"
const RancherDomain = "rancher.internal"
const RancherDns = "io.rancher.container.dns"

type StartHandler struct {
	Client            SimpleDockerClient
	ContainerStateDir string
}

func setupResolvConf(container *docker.Container) error {
	if _, ok := container.Config.Labels[RancherSystemLabelKey]; ok {
		return nil
	}

	p := container.ResolvConfPath
	input, err := os.Open(p)
	if err != nil {
		return err
	}

	defer input.Close()

	var buffer bytes.Buffer
	scanner := bufio.NewScanner(input)
	for scanner.Scan() {
		text := scanner.Text()
		if strings.Contains(text, RancherNameserver) {
			continue
		}

		if strings.HasPrefix(text, "nameserver") {
			text = "# " + text
		}

		if strings.HasPrefix(text, "search") && !strings.Contains(text, RancherDomain) {
			text = text + " " + RancherDomain
		}

		if _, err := buffer.Write([]byte(text)); err != nil {
			return err
		}

		if _, err := buffer.Write([]byte("\n")); err != nil {
			return err
		}
	}

	buffer.Write([]byte("nameserver "))
	buffer.Write([]byte(RancherNameserver))
	buffer.Write([]byte("\n"))

	input.Close()
	return ioutil.WriteFile(p, buffer.Bytes(), 0666)
}

func (h *StartHandler) Handle(event *docker.APIEvents) error {
	// Note: event.ID == container's ID
	lock := locks.Lock("start." + event.ID)
	if lock == nil {
		log.Debugf("Container locked. Can't run StartHandler. ID: [%s]", event.ID)
		return nil
	}
	defer lock.Unlock()

	c, err := h.Client.InspectContainer(event.ID)
	if err != nil {
		return err
	}

	rancherIP, err := h.getRancherIP(c)
	if err != nil {
		return err
	}

	if rancherIP == "" {
		if c.Config.Labels[RancherDns] == "true" {
			return setupResolvConf(c)
		}
		return nil
	}

	if !c.State.Running {
		log.Infof("Container [%s] not running. Can't assign IP [%s].", c.ID, rancherIP)
		return nil
	}

	pid := c.State.Pid
	log.Infof("Assigning IP [%s], ContainerId [%s], Pid [%v]", rancherIP, event.ID, pid)
	if err := configureIP(strconv.Itoa(pid), rancherIP); err != nil {
		// If it stopped running, don't return error
		c, inspectErr := h.Client.InspectContainer(event.ID)
		if inspectErr != nil {
			log.Warn("Failed to inspect container: ", event.ID, inspectErr)
			return err
		}

		if !c.State.Running {
			log.Infof("Container [%s] not running. Can't assign IP [%s].", c.ID, rancherIP)
			return nil
		}

		return err
	}

	return setupResolvConf(c)
}

func (h *StartHandler) getRancherIP(c *docker.Container) (string, error) {
	if ip, ok := c.Config.Labels[RancherIPLabelKey]; ok {
		return ip, nil
	}
	for _, env := range c.Config.Env {
		if strings.HasPrefix(env, RancherIPEnvKey) {
			return strings.TrimPrefix(env, RancherIPEnvKey), nil
		}
	}

	filePath := path.Join(h.ContainerStateDir, c.ID)
	if _, err := os.Stat(filePath); err == nil {
		file, e := ioutil.ReadFile(filePath)
		if e != nil {
			return "", fmt.Errorf("Error reading file for container %s: %v", c.ID, e)
		}
		var instanceData instance
		jerr := json.Unmarshal(file, &instanceData)
		if jerr != nil {
			return "", fmt.Errorf("Error unmarshalling json for container %s: %v", c.ID, jerr)
		}

		if len(instanceData.Nics) > 0 {
			nic := instanceData.Nics[0]
			var ipData ipAddress
			for _, i := range nic.IpAddresses {
				if i.Role == "primary" {
					ipData = i
					break
				}
			}
			if ipData.Address != "" {
				return ipData.Address + "/" + strconv.Itoa(ipData.Subnet.CidrSize), nil
			}
		}
	}

	return "", nil
}

func configureIP(pid string, ip string) error {
	command := buildCommand(pid, ip)

	output, err := command.CombinedOutput()
	log.Debugf(string(output))
	if err != nil {
		return err
	}

	return nil
}

var buildCommand = func(pid string, ip string) *exec.Cmd {
	return exec.Command("net-util.sh", "-p", pid, "-i", ip)
}

type subnet struct {
	CidrSize int
}

type ipAddress struct {
	Address string
	Subnet  subnet
	Role    string
}

type nic struct {
	IpAddresses []ipAddress `json:"ipAddresses"`
}

type instance struct {
	Nics []nic
}
<commit_msg>Set namespace from svc/stack labels<commit_after>package events

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/fsouza/go-dockerclient"
	"github.com/rancherio/go-machine-service/locks"
)

const RancherIPLabelKey = "io.rancher.container.ip"
const RancherSystemLabelKey = "io.rancher.container.system"
const RancherIPEnvKey = "RANCHER_IP="
const RancherNameserver = "169.254.169.250"
const RancherDomain = "rancher.internal"
const RancherDns = "io.rancher.container.dns"

type StartHandler struct {
	Client            SimpleDockerClient
	ContainerStateDir string
}

func getDnsSearch(container *docker.Container) []string {
	var defaultDomains []string
	var svcNameSpace string
	var stackNameSpace string

	//from labels - for upgraded systems
	if container.Config.Labels != nil {
		if value, ok := container.Config.Labels["io.rancher.stack_service.name"]; ok {
			splitted := strings.Split(value, "/")
			svc := strings.ToLower(splitted[1])
			stack := strings.ToLower(splitted[0])
			svcNameSpace = svc + "." + stack + "." + RancherDomain
			stackNameSpace = stack + "." + RancherDomain
			defaultDomains = append(defaultDomains, svcNameSpace)
			defaultDomains = append(defaultDomains, stackNameSpace)
		}
	}

	//from search domains
	if container.HostConfig.DNSSearch != nil {
		for _, domain := range container.HostConfig.DNSSearch {
			if domain != svcNameSpace && domain != stackNameSpace {
				defaultDomains = append(defaultDomains, domain)
			}
		}
	}

	// default rancher domain
	defaultDomains = append(defaultDomains, RancherDomain)
	return defaultDomains
}

func setupResolvConf(container *docker.Container) error {
	if _, ok := container.Config.Labels[RancherSystemLabelKey]; ok {
		return nil
	}

	p := container.ResolvConfPath
	input, err := os.Open(p)
	if err != nil {
		return err
	}

	defer input.Close()

	var buffer bytes.Buffer
	scanner := bufio.NewScanner(input)
	searchSet := false
	nameserverSet := false
	for scanner.Scan() {
		text := scanner.Text()

		if strings.Contains(text, RancherNameserver) {
			nameserverSet = true
		} else if strings.HasPrefix(text, "nameserver") {
			text = "# " + text
		}

		if strings.HasPrefix(text, "search") {
			for _, domain := range getDnsSearch(container) {
				if strings.Contains(text, " "+domain) {
					continue
				}
				text = text + " " + domain
			}
			searchSet = true
		}

		if _, err := buffer.Write([]byte(text)); err != nil {
			return err
		}

		if _, err := buffer.Write([]byte("\n")); err != nil {
			return err
		}
	}

	if !searchSet {
		buffer.Write([]byte("search " + strings.ToLower(strings.Join(getDnsSearch(container), " "))))
		buffer.Write([]byte("\n"))
	}

	if !nameserverSet {
		buffer.Write([]byte("nameserver "))
		buffer.Write([]byte(RancherNameserver))
		buffer.Write([]byte("\n"))
	}

	input.Close()
	return ioutil.WriteFile(p, buffer.Bytes(), 0666)
}

func (h *StartHandler) Handle(event *docker.APIEvents) error {
	// Note: event.ID == container's ID
	lock := locks.Lock("start." + event.ID)
	if lock == nil {
		log.Debugf("Container locked. Can't run StartHandler. ID: [%s]", event.ID)
		return nil
	}
	defer lock.Unlock()

	c, err := h.Client.InspectContainer(event.ID)
	if err != nil {
		return err
	}

	rancherIP, err := h.getRancherIP(c)
	if err != nil {
		return err
	}

	if rancherIP == "" {
		if c.Config.Labels[RancherDns] == "true" {
			return setupResolvConf(c)
		}
		return nil
	}

	if !c.State.Running {
		log.Infof("Container [%s] not running. Can't assign IP [%s].", c.ID, rancherIP)
		return nil
	}

	pid := c.State.Pid
	log.Infof("Assigning IP [%s], ContainerId [%s], Pid [%v]", rancherIP, event.ID, pid)
	if err := configureIP(strconv.Itoa(pid), rancherIP); err != nil {
		// If it stopped running, don't return error
		c, inspectErr := h.Client.InspectContainer(event.ID)
		if inspectErr != nil {
			log.Warn("Failed to inspect container: ", event.ID, inspectErr)
			return err
		}

		if !c.State.Running {
			log.Infof("Container [%s] not running. Can't assign IP [%s].", c.ID, rancherIP)
			return nil
		}

		return err
	}

	return setupResolvConf(c)
}

func (h *StartHandler) getRancherIP(c *docker.Container) (string, error) {
	if ip, ok := c.Config.Labels[RancherIPLabelKey]; ok {
		return ip, nil
	}
	for _, env := range c.Config.Env {
		if strings.HasPrefix(env, RancherIPEnvKey) {
			return strings.TrimPrefix(env, RancherIPEnvKey), nil
		}
	}

	filePath := path.Join(h.ContainerStateDir, c.ID)
	if _, err := os.Stat(filePath); err == nil {
		file, e := ioutil.ReadFile(filePath)
		if e != nil {
			return "", fmt.Errorf("Error reading file for container %s: %v", c.ID, e)
		}
		var instanceData instance
		jerr := json.Unmarshal(file, &instanceData)
		if jerr != nil {
			return "", fmt.Errorf("Error unmarshalling json for container %s: %v", c.ID, jerr)
		}

		if len(instanceData.Nics) > 0 {
			nic := instanceData.Nics[0]
			var ipData ipAddress
			for _, i := range nic.IpAddresses {
				if i.Role == "primary" {
					ipData = i
					break
				}
			}
			if ipData.Address != "" {
				return ipData.Address + "/" + strconv.Itoa(ipData.Subnet.CidrSize), nil
			}
		}
	}

	return "", nil
}

func configureIP(pid string, ip string) error {
	command := buildCommand(pid, ip)

	output, err := command.CombinedOutput()
	log.Debugf(string(output))
	if err != nil {
		return err
	}

	return nil
}

var buildCommand = func(pid string, ip string) *exec.Cmd {
	return exec.Command("net-util.sh", "-p", pid, "-i", ip)
}

type subnet struct {
	CidrSize int
}

type ipAddress struct {
	Address string
	Subnet  subnet
	Role    string
}

type nic struct {
	IpAddresses []ipAddress `json:"ipAddresses"`
}

type instance struct {
	Nics []nic
}
<|endoftext|>


<commit_before>package link

import (
	"fmt"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/internal"
)

type RawAttachProgramOptions struct {
	// File descriptor to attach to. This differs for each attach type.
	Target int
	// Program to attach.
	Program *ebpf.Program
	// Program to replace (cgroups).
	Replace *ebpf.Program
	// Attach must match the attach type of Program (and Replace).
	Attach ebpf.AttachType
	// Flags control the attach behaviour. This differs for each attach type.
	Flags uint32
}

// RawAttachProgram is a low level wrapper around BPF_PROG_ATTACH.
//
// You should use one of the higher level abstractions available in this
// package if possible.
func RawAttachProgram(opts RawAttachProgramOptions) error {
	if err := haveProgAttach(); err != nil {
		return err
	}

	var replaceFd uint32
	if opts.Replace != nil {
		replaceFd = uint32(opts.Replace.FD())
	}

	attr := internal.BPFProgAttachAttr{
		TargetFd:     uint32(opts.Target),
		AttachBpfFd:  uint32(opts.Program.FD()),
		ReplaceBpfFd: replaceFd,
		AttachType:   uint32(opts.Attach),
		AttachFlags:  uint32(opts.Flags),
	}

	if err := internal.BPFProgAttach(&attr); err != nil {
		return fmt.Errorf("can't attach program: %s", err)
	}
	return nil
}

type RawDetachProgramOptions struct {
	Target  int
	Program *ebpf.Program
	Attach  ebpf.AttachType
}

// RawDetachProgram is a low level wrapper around BPF_PROG_DETACH.
//
// You should use one of the higher level abstractions available in this
// package if possible.
func RawDetachProgram(opts RawDetachProgramOptions) error {
	if err := haveProgAttach(); err != nil {
		return err
	}

	attr := internal.BPFProgDetachAttr{
		TargetFd:    uint32(opts.Target),
		AttachBpfFd: uint32(opts.Program.FD()),
		AttachType:  uint32(opts.Attach),
	}
	if err := internal.BPFProgDetach(&attr); err != nil {
		return fmt.Errorf("can't detach program: %s", err)
	}

	return nil
}
<commit_msg>link: wrap errors for Raw{At,De}tachProgram<commit_after>package link

import (
	"fmt"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/internal"
)

type RawAttachProgramOptions struct {
	// File descriptor to attach to. This differs for each attach type.
	Target int
	// Program to attach.
	Program *ebpf.Program
	// Program to replace (cgroups).
	Replace *ebpf.Program
	// Attach must match the attach type of Program (and Replace).
	Attach ebpf.AttachType
	// Flags control the attach behaviour. This differs for each attach type.
	Flags uint32
}

// RawAttachProgram is a low level wrapper around BPF_PROG_ATTACH.
//
// You should use one of the higher level abstractions available in this
// package if possible.
func RawAttachProgram(opts RawAttachProgramOptions) error {
	if err := haveProgAttach(); err != nil {
		return err
	}

	var replaceFd uint32
	if opts.Replace != nil {
		replaceFd = uint32(opts.Replace.FD())
	}

	attr := internal.BPFProgAttachAttr{
		TargetFd:     uint32(opts.Target),
		AttachBpfFd:  uint32(opts.Program.FD()),
		ReplaceBpfFd: replaceFd,
		AttachType:   uint32(opts.Attach),
		AttachFlags:  uint32(opts.Flags),
	}

	if err := internal.BPFProgAttach(&attr); err != nil {
		return fmt.Errorf("can't attach program: %w", err)
	}
	return nil
}

type RawDetachProgramOptions struct {
	Target  int
	Program *ebpf.Program
	Attach  ebpf.AttachType
}

// RawDetachProgram is a low level wrapper around BPF_PROG_DETACH.
//
// You should use one of the higher level abstractions available in this
// package if possible.
func RawDetachProgram(opts RawDetachProgramOptions) error {
	if err := haveProgAttach(); err != nil {
		return err
	}

	attr := internal.BPFProgDetachAttr{
		TargetFd:    uint32(opts.Target),
		AttachBpfFd: uint32(opts.Program.FD()),
		AttachType:  uint32(opts.Attach),
	}
	if err := internal.BPFProgDetach(&attr); err != nil {
		return fmt.Errorf("can't detach program: %w", err)
	}

	return nil
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"os"
)

func (pindex *PIndex) Run() {
	for m := range pindex.stream {
		// TODO: probably need things like stream reset/rollback
		// and snapshot kinds of ops here, too.

		// TODO: maybe need a more batchy API?  Perhaps, yet another
		// goroutine that clumps up up updates into bigger batches?

		switch m := m.(type) {
		case *StreamUpdate:
			pindex.bindex.Index(string(m.Id()), m.Body())
		case *StreamDelete:
			pindex.bindex.Delete(string(m.Id()))
		}
	}

	// TODO: what about any inflight queries or ops?

	pindex.BIndex().Close()

	os.RemoveAll(pindex.Path())
}
<commit_msg>comment on how bleve handles Close() and concurrency<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"os"
)

func (pindex *PIndex) Run() {
	for m := range pindex.stream {
		// TODO: probably need things like stream reset/rollback
		// and snapshot kinds of ops here, too.

		// TODO: maybe need a more batchy API?  Perhaps, yet another
		// goroutine that clumps up up updates into bigger batches?

		switch m := m.(type) {
		case *StreamUpdate:
			pindex.bindex.Index(string(m.Id()), m.Body())
		case *StreamDelete:
			pindex.bindex.Delete(string(m.Id()))
		}
	}

	// The bleve.Index.Close() handles any inflight, concurrent
	// queries with its own locking.
	pindex.BIndex().Close()

	os.RemoveAll(pindex.Path())
}
<|endoftext|>