<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/pkg/jsonmessage"
	"github.com/fsouza/go-dockerclient"
	"github.com/pwaller/barrier"
)

type Container struct {
	Name      string
	Args, Env []string
	StatusURI string

	client    *docker.Client
	container *docker.Container

	Failed, Superceded, Ready, Closing barrier.Barrier

	wg *sync.WaitGroup

	Errors  <-chan error
	errorsW chan<- error
}

type SourceType int

const (
	// Run build with current directory with context
	BuildCwd            SourceType = iota
	BuildTarballContent            // Build with specified io.Reader as context
	BuildTarballURL                // Build with specified remote URL as context
	DockerPull                     // Run a docker pull to obtain the image
)

type ContainerSource struct {
	Type                SourceType
	buildTarballContent io.Reader
	buildTarballURL     string
	dockerImageName     string
}

// Construct a *Container. When the `wg` WaitGroup is zero, there is nothing
// outstanding (such as firewall rules which need garbage collecting).
func NewContainer(client *docker.Client, name string, wg *sync.WaitGroup) *Container {

	errors := make(chan error)

	c := &Container{
		Name:    name,
		client:  client,
		wg:      wg,
		Errors:  errors,
		errorsW: errors,
	}

	// If the container fails we should assume it should be torn down.
	c.Failed.Forward(&c.Closing)

	return c
}

// Generate a docker image. This can be done through various mechanisms in
// response to an UpdateEvent (see SourceType constant declarations).
func (c *Container) Build(config UpdateEvent) error {
	if config.BuildComplete != nil {
		defer close(config.BuildComplete)
	}

	var err error
	bo := docker.BuildImageOptions{}
	bo.Name = c.Name
	bo.OutputStream = config.OutputStream
	if bo.OutputStream == nil {
		bo.OutputStream = os.Stderr
	}

	switch config.Source.Type {
	case BuildCwd:
		bo.ContextDir, err = os.Getwd()
		if err != nil {
			return err
		}
	case BuildTarballContent:
		bo.InputStream = config.Source.buildTarballContent
	default:
		return fmt.Errorf("Unimplemented ContainerSource: %v", config.Source.Type)
	}

	return c.client.BuildImage(bo)
}

func PullProgressCopier(target io.Writer) (io.WriteCloser, <-chan error) {
	reader, wrappedWriter := io.Pipe()
	errorC := make(chan error)
	go func() {
		finish := make(chan struct{})
		defer close(finish)
		defer close(errorC)

		mu := sync.Mutex{}
		lastMessage := jsonmessage.JSONMessage{}
		newMessage := false

		printMessage := func(m *jsonmessage.JSONMessage) {
			if m.ProgressMessage != "" {
				fmt.Fprintln(target, m.ID[:8], m.Status, m.ProgressMessage)
			} else if m.Progress != nil {
				fmt.Fprintln(target, m.ID[:8], m.Status, m.Progress.String())
			} else {
				m.Display(target, false)
			}
		}

		go func() {
			tick := time.Tick(1 * time.Second)
			for {
				select {
				case <-tick:
					mu.Lock()
					if newMessage {
						printMessage(&lastMessage)
						newMessage = false
					}
					mu.Unlock()

				case <-finish:
					return
				}
			}
		}()

		dec := json.NewDecoder(reader)
		for {
			tmp := jsonmessage.JSONMessage{}
			err := dec.Decode(&tmp)

			mu.Lock()
			if tmp.Error != nil || tmp.ErrorMessage != "" {
				tmp.Display(target, false)
				if tmp.Error != nil {
					errorC <- tmp.Error
				} else {
					errorC <- fmt.Errorf("%s", tmp.ErrorMessage)
				}
				return
			} else if tmp.Status != "Downloading" && tmp.Status != "Extracting" {
				printMessage(&tmp)
			} else {
				newMessage = true
				lastMessage = tmp
			}
			mu.Unlock()

			if err == io.EOF {
				return
			}
			if err != nil {
				log.Print("decode failure in  ", err)
				return
			}
		}
	}()
	return wrappedWriter, errorC
}

// Pull an image from a docker repository.
func (c *Container) Pull(config UpdateEvent) error {
	if config.BuildComplete != nil {
		defer close(config.BuildComplete)
	}

	parts := strings.SplitN(config.Source.dockerImageName, ":", 2)

	pio := docker.PullImageOptions{}

	pio.Repository = parts[0]
	pio.Tag = "latest"
	if len(parts) == 2 {
		pio.Tag = parts[1]
	}
	pio.Registry = ""
	pio.RawJSONStream = true

	log.Println("Pulling", pio.Repository)

	target := config.OutputStream
	if target == nil {
		target = os.Stderr
	}

	outputStream, errorC := PullProgressCopier(target)
	pio.OutputStream = outputStream

	pullImageErr := c.client.PullImage(pio, docker.AuthConfiguration{})

	outputStream.Close()

	if pullImageErr != nil {
		return pullImageErr
	}

	return <-errorC
}

// `docker create` the container.
func (c *Container) Create(source ContainerSource) error {
	opts := docker.CreateContainerOptions{
		Name: c.Name,
		Config: &docker.Config{
			Hostname:     c.Name,
			AttachStdout: true,
			AttachStderr: true,
			Env:          c.Env,
			Cmd:          c.Args,
		},
	}

	switch source.Type {
	case DockerPull:
		opts.Config.Image = source.dockerImageName
	case BuildCwd, BuildTarballContent:
		opts.Config.Image = c.Name
	}

	var err error
	c.container, err = c.client.CreateContainer(opts)

	return err
}

// CopyOutput copies the output of the container to `w` and blocks until
// completion
func (c *Container) CopyOutput() error {

	// TODO(pwaller): at some point move this on to 'c' for configurability?
	w := os.Stderr
	// Blocks until stream closed
	return c.client.AttachToContainer(docker.AttachToContainerOptions{
		Container:    c.container.ID,
		OutputStream: w,
		ErrorStream:  w,
		Logs:         true,
		Stdout:       true,
		Stderr:       true,
		Stream:       true,
	})
}

// :todo(drj): May want to return errors for truly broken containers (timeout).
// Poll for the program inside the container being ready to accept connections
// Returns `true` for success and `false` for failure.
func (c *Container) AwaitListening() bool {

	for _, port := range c.container.NetworkSettings.PortMappingAPI() {
		url := fmt.Sprint("http://", port.IP, ":", port.PublicPort, c.StatusURI)
		for {
			response, err := http.Get(url)
			if response != nil && response.Body != nil {
				response.Body.Close()
			}
			if err == nil && response.StatusCode == http.StatusOK {
				return true
			}
			time.Sleep(50 * time.Millisecond)

			select {
			case <-c.Closing.Barrier():
				// If the container has closed, cease waiting
				return false
			default:
			}
		}
	}

	return false
}

// Start the container (and notify it if c.Closing falls)
func (c *Container) Start() error {
	hc := &docker.HostConfig{
		PublishAllPorts: true,
	}
	err := c.client.StartContainer(c.container.ID, hc)
	if err != nil {
		return err
	}

	// Load container.NetworkSettings
	c.container, err = c.client.InspectContainer(c.container.ID)
	if err != nil {
		return err
	}

	// Listen on the Closing barrier and send a kill to the container if it
	// falls.
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		<-c.Closing.Barrier()
		// If the container is signaled to close, send a kill signal
		err := c.client.KillContainer(docker.KillContainerOptions{
			ID: c.container.ID,
		})
		if err == nil {
			return
		}
		switch err := err.(type) {
		case *docker.NoSuchContainer:
			// The container already went away, who cares.
			return
		default:
			log.Println("Killing container failed:", c.container.ID, err)
		}
	}()
	return nil
}

// Wait until container exits
func (c *Container) Wait() (int, error) {
	return c.client.WaitContainer(c.container.ID)
}

// Internal function for raising an error.
func (c *Container) err(err error) {
	c.errorsW <- err
	c.Closing.Fall()
}

// Manage the whole lifecycle of the container in response to a request to
// start it.
func (c *Container) Run(event UpdateEvent) (int, error) {

	defer c.Closing.Fall()
	defer close(c.errorsW)

	switch event.Source.Type {
	case DockerPull:
		err := c.Pull(event)
		if err != nil {
			return -2, err
		}
	case BuildTarballContent, BuildCwd:
		err := c.Build(event)
		if err != nil {
			return -2, err
		}
	}

	err := c.Create(event.Source)
	if err != nil {
		return -1, err
	}
	defer c.Delete()

	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		err := c.CopyOutput()
		if err != nil {
			c.err(err)
		}
	}()

	err = c.Start()
	if err != nil {
		return -1, err
	}

	go func() {
		if !c.AwaitListening() {
			c.Failed.Fall()
			return
		}
		c.Ready.Fall()
	}()

	return c.Wait()
}

func (c *Container) Delete() {
	err := c.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            c.container.ID,
		RemoveVolumes: true,
		Force:         true,
	})
	if err != nil {
		log.Println("Warn: failed to delete container:", err)
	}
}
<commit_msg>Add 5 minute start deadline<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/pkg/jsonmessage"
	"github.com/fsouza/go-dockerclient"
	"github.com/pwaller/barrier"
)

type Container struct {
	Name      string
	Args, Env []string
	StatusURI string

	client    *docker.Client
	container *docker.Container

	Failed, Superceded, Ready, Closing barrier.Barrier

	wg *sync.WaitGroup

	Errors  <-chan error
	errorsW chan<- error
}

type SourceType int

const (
	// Run build with current directory with context
	BuildCwd            SourceType = iota
	BuildTarballContent            // Build with specified io.Reader as context
	BuildTarballURL                // Build with specified remote URL as context
	DockerPull                     // Run a docker pull to obtain the image
)

type ContainerSource struct {
	Type                SourceType
	buildTarballContent io.Reader
	buildTarballURL     string
	dockerImageName     string
}

// Construct a *Container. When the `wg` WaitGroup is zero, there is nothing
// outstanding (such as firewall rules which need garbage collecting).
func NewContainer(client *docker.Client, name string, wg *sync.WaitGroup) *Container {

	errors := make(chan error)

	c := &Container{
		Name:    name,
		client:  client,
		wg:      wg,
		Errors:  errors,
		errorsW: errors,
	}

	// If the container fails we should assume it should be torn down.
	c.Failed.Forward(&c.Closing)

	return c
}

// Generate a docker image. This can be done through various mechanisms in
// response to an UpdateEvent (see SourceType constant declarations).
func (c *Container) Build(config UpdateEvent) error {
	if config.BuildComplete != nil {
		defer close(config.BuildComplete)
	}

	var err error
	bo := docker.BuildImageOptions{}
	bo.Name = c.Name
	bo.OutputStream = config.OutputStream
	if bo.OutputStream == nil {
		bo.OutputStream = os.Stderr
	}

	switch config.Source.Type {
	case BuildCwd:
		bo.ContextDir, err = os.Getwd()
		if err != nil {
			return err
		}
	case BuildTarballContent:
		bo.InputStream = config.Source.buildTarballContent
	default:
		return fmt.Errorf("Unimplemented ContainerSource: %v", config.Source.Type)
	}

	return c.client.BuildImage(bo)
}

func PullProgressCopier(target io.Writer) (io.WriteCloser, <-chan error) {
	reader, wrappedWriter := io.Pipe()
	errorC := make(chan error)
	go func() {
		finish := make(chan struct{})
		defer close(finish)
		defer close(errorC)

		mu := sync.Mutex{}
		lastMessage := jsonmessage.JSONMessage{}
		newMessage := false

		printMessage := func(m *jsonmessage.JSONMessage) {
			if m.ProgressMessage != "" {
				fmt.Fprintln(target, m.ID[:8], m.Status, m.ProgressMessage)
			} else if m.Progress != nil {
				fmt.Fprintln(target, m.ID[:8], m.Status, m.Progress.String())
			} else {
				m.Display(target, false)
			}
		}

		go func() {
			tick := time.Tick(1 * time.Second)
			for {
				select {
				case <-tick:
					mu.Lock()
					if newMessage {
						printMessage(&lastMessage)
						newMessage = false
					}
					mu.Unlock()

				case <-finish:
					return
				}
			}
		}()

		dec := json.NewDecoder(reader)
		for {
			tmp := jsonmessage.JSONMessage{}
			err := dec.Decode(&tmp)

			mu.Lock()
			if tmp.Error != nil || tmp.ErrorMessage != "" {
				tmp.Display(target, false)
				if tmp.Error != nil {
					errorC <- tmp.Error
				} else {
					errorC <- fmt.Errorf("%s", tmp.ErrorMessage)
				}
				return
			} else if tmp.Status != "Downloading" && tmp.Status != "Extracting" {
				printMessage(&tmp)
			} else {
				newMessage = true
				lastMessage = tmp
			}
			mu.Unlock()

			if err == io.EOF {
				return
			}
			if err != nil {
				log.Print("decode failure in  ", err)
				return
			}
		}
	}()
	return wrappedWriter, errorC
}

// Pull an image from a docker repository.
func (c *Container) Pull(config UpdateEvent) error {
	if config.BuildComplete != nil {
		defer close(config.BuildComplete)
	}

	parts := strings.SplitN(config.Source.dockerImageName, ":", 2)

	pio := docker.PullImageOptions{}

	pio.Repository = parts[0]
	pio.Tag = "latest"
	if len(parts) == 2 {
		pio.Tag = parts[1]
	}
	pio.Registry = ""
	pio.RawJSONStream = true

	log.Println("Pulling", pio.Repository)

	target := config.OutputStream
	if target == nil {
		target = os.Stderr
	}

	outputStream, errorC := PullProgressCopier(target)
	pio.OutputStream = outputStream

	pullImageErr := c.client.PullImage(pio, docker.AuthConfiguration{})

	outputStream.Close()

	if pullImageErr != nil {
		return pullImageErr
	}

	return <-errorC
}

// `docker create` the container.
func (c *Container) Create(source ContainerSource) error {
	opts := docker.CreateContainerOptions{
		Name: c.Name,
		Config: &docker.Config{
			Hostname:     c.Name,
			AttachStdout: true,
			AttachStderr: true,
			Env:          c.Env,
			Cmd:          c.Args,
		},
	}

	switch source.Type {
	case DockerPull:
		opts.Config.Image = source.dockerImageName
	case BuildCwd, BuildTarballContent:
		opts.Config.Image = c.Name
	}

	var err error
	c.container, err = c.client.CreateContainer(opts)

	return err
}

// CopyOutput copies the output of the container to `w` and blocks until
// completion
func (c *Container) CopyOutput() error {

	// TODO(pwaller): at some point move this on to 'c' for configurability?
	w := os.Stderr
	// Blocks until stream closed
	return c.client.AttachToContainer(docker.AttachToContainerOptions{
		Container:    c.container.ID,
		OutputStream: w,
		ErrorStream:  w,
		Logs:         true,
		Stdout:       true,
		Stderr:       true,
		Stream:       true,
	})
}

// :todo(drj): May want to return errors for truly broken containers (timeout).
// Poll for the program inside the container being ready to accept connections
// Returns `true` for success and `false` for failure.
func (c *Container) AwaitListening() bool {

	const (
		DefaultTimeout = 5 * time.Minute
		PollFrequency  = 10 // times per second (via integer division of ns)
	)

	startDeadline := time.Now().Add(DefaultTimeout)

	for _, port := range c.container.NetworkSettings.PortMappingAPI() {
		url := fmt.Sprint("http://", port.IP, ":", port.PublicPort, c.StatusURI)
		for {
			response, err := http.Get(url)
			if response != nil && response.Body != nil {
				response.Body.Close()
			}
			if err == nil && response.StatusCode == http.StatusOK {
				return true
			}

			if time.Now().After(startDeadline) {
				log.Printf("Took longer than %v to start, giving up", DefaultTimeout)
				return false
			}

			time.Sleep(time.Second / PollFrequency)

			select {
			case <-c.Closing.Barrier():
				// If the container has closed, cease waiting
				return false
			default:
			}
		}
	}

	return false
}

// Start the container (and notify it if c.Closing falls)
func (c *Container) Start() error {
	hc := &docker.HostConfig{
		PublishAllPorts: true,
	}
	err := c.client.StartContainer(c.container.ID, hc)
	if err != nil {
		return err
	}

	// Load container.NetworkSettings
	c.container, err = c.client.InspectContainer(c.container.ID)
	if err != nil {
		return err
	}

	// Listen on the Closing barrier and send a kill to the container if it
	// falls.
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		<-c.Closing.Barrier()
		// If the container is signaled to close, send a kill signal
		err := c.client.KillContainer(docker.KillContainerOptions{
			ID: c.container.ID,
		})
		if err == nil {
			return
		}
		switch err := err.(type) {
		case *docker.NoSuchContainer:
			// The container already went away, who cares.
			return
		default:
			log.Println("Killing container failed:", c.container.ID, err)
		}
	}()
	return nil
}

// Wait until container exits
func (c *Container) Wait() (int, error) {
	return c.client.WaitContainer(c.container.ID)
}

// Internal function for raising an error.
func (c *Container) err(err error) {
	c.errorsW <- err
	c.Closing.Fall()
}

// Manage the whole lifecycle of the container in response to a request to
// start it.
func (c *Container) Run(event UpdateEvent) (int, error) {

	defer c.Closing.Fall()
	defer close(c.errorsW)

	switch event.Source.Type {
	case DockerPull:
		err := c.Pull(event)
		if err != nil {
			return -2, err
		}
	case BuildTarballContent, BuildCwd:
		err := c.Build(event)
		if err != nil {
			return -2, err
		}
	}

	err := c.Create(event.Source)
	if err != nil {
		return -1, err
	}
	defer c.Delete()

	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		err := c.CopyOutput()
		if err != nil {
			c.err(err)
		}
	}()

	err = c.Start()
	if err != nil {
		return -1, err
	}

	go func() {
		if !c.AwaitListening() {
			c.Failed.Fall()
			return
		}
		c.Ready.Fall()
	}()

	return c.Wait()
}

func (c *Container) Delete() {
	err := c.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            c.container.ID,
		RemoveVolumes: true,
		Force:         true,
	})
	if err != nil {
		log.Println("Warn: failed to delete container:", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/fsouza/go-dockerclient"
	"github.com/pwaller/barrier"
)

type Container struct {
	Name string
	Env  []string

	client    *docker.Client
	container *docker.Container

	Failed, Superceded, Ready, Closing barrier.Barrier

	wg *sync.WaitGroup

	Errors  <-chan error
	errorsW chan<- error
}

type SourceType int

const (
	// Run build with current directory with context
	BuildCwd            SourceType = iota
	BuildTarballContent            // Build with specified io.Reader as context
	BuildTarballURL                // Build with specified remote URL as context
	DockerPull                     // Run a docker pull to obtain the image
)

type ContainerSource struct {
	Type                SourceType
	buildTarballContent io.Reader
	buildTarballURL     string
	dockerImageName     string
}

func NewContainer(client *docker.Client, name string, wg *sync.WaitGroup) *Container {

	errors := make(chan error)

	c := &Container{
		Name:    name,
		client:  client,
		wg:      wg,
		Errors:  errors,
		errorsW: errors,
	}

	// If the container fails we should assume it should be torn down.
	c.Failed.Forward(&c.Closing)

	return c
}

func (c *Container) Build(config UpdateEvent) error {
	if config.BuildComplete != nil {
		defer close(config.BuildComplete)
	}

	var err error
	bo := docker.BuildImageOptions{}
	bo.Name = c.Name
	bo.OutputStream = config.OutputStream
	if bo.OutputStream == nil {
		bo.OutputStream = os.Stderr
	}

	switch config.Source.Type {
	case BuildCwd:
		bo.ContextDir, err = os.Getwd()
		if err != nil {
			return err
		}
	case BuildTarballContent:
		bo.InputStream = config.Source.buildTarballContent
	default:
		return fmt.Errorf("Unimplemented ContainerSource: %v", config.Source.Type)
	}

	return c.client.BuildImage(bo)
}

func (c *Container) Create() error {
	opts := docker.CreateContainerOptions{
		Name: c.Name,
		Config: &docker.Config{
			Hostname:     c.Name,
			Image:        c.Name,
			AttachStdout: true,
			AttachStderr: true,
			Env:          c.Env,
		},
	}

	var err error
	c.container, err = c.client.CreateContainer(opts)

	return err
}

// CopyOutput copies the output of the container to `w` and blocks until
// completion
func (c *Container) CopyOutput() error {
	// TODO(pwaller): at some point move this on to 'c' for configurability?
	w := os.Stderr
	// Blocks until stream closed
	return c.client.AttachToContainer(docker.AttachToContainerOptions{
		Container:    c.container.ID,
		OutputStream: w,
		ErrorStream:  w,
		Logs:         true,
		Stdout:       true,
		Stderr:       true,
		Stream:       true,
	})
}

// :todo(drj): May want to return errors for truly broken containers (timeout).
func (c *Container) AwaitListening() {

	for _, port := range c.container.NetworkSettings.PortMappingAPI() {
		url := fmt.Sprint("http://", port.IP, ":", port.PublicPort, "/")
		for {
			response, err := http.Get(url)
			if err == nil && response.StatusCode == http.StatusOK {
				break
			}
			time.Sleep(50 * time.Millisecond)
		}
	}
	c.Ready.Fall()
}

func (c *Container) Start() error {
	hc := &docker.HostConfig{
		PublishAllPorts: true,
	}
	err := c.client.StartContainer(c.container.ID, hc)
	if err != nil {
		return err
	}

	// Load container.NetworkSettings
	c.container, err = c.client.InspectContainer(c.container.ID)
	if err != nil {
		return err
	}

	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		<-c.Closing.Barrier()
		// If the container is signaled to close, send a kill signal
		err := c.client.KillContainer(docker.KillContainerOptions{
			ID: c.container.ID,
		})
		if err == nil {
			return
		}
		switch err := err.(type) {
		case *docker.NoSuchContainer:
			// The container already went away, who cares.
			return
		default:
			log.Println("Killing container failed:", c.container.ID, err)
		}
	}()
	return nil
}

func (c *Container) Wait() (int, error) {
	return c.client.WaitContainer(c.container.ID)
}

func (c *Container) err(err error) {
	c.errorsW <- err
	c.Closing.Fall()
}

func (c *Container) Run(event UpdateEvent) (int, error) {

	defer close(c.errorsW)

	err := c.Build(event)
	if err != nil {
		return -2, err
	}

	err = c.Create()
	if err != nil {
		return -1, err
	}
	defer c.Delete()

	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		err := c.CopyOutput()
		if err != nil {
			c.err(err)
		}
	}()

	err = c.Start()
	if err != nil {
		return -1, err
	}

	go func() {
		c.AwaitListening()
		c.Ready.Fall()
		log.Println("Listening on", c.container.NetworkSettings.PortMappingAPI())
	}()

	return c.Wait()
}

func (c *Container) Delete() {
	c.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            c.container.ID,
		RemoveVolumes: true,
		Force:         true,
	})
}
<commit_msg>When Run() finishes, fell the Closing barrier<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/fsouza/go-dockerclient"
	"github.com/pwaller/barrier"
)

type Container struct {
	Name string
	Env  []string

	client    *docker.Client
	container *docker.Container

	Failed, Superceded, Ready, Closing barrier.Barrier

	wg *sync.WaitGroup

	Errors  <-chan error
	errorsW chan<- error
}

type SourceType int

const (
	// Run build with current directory with context
	BuildCwd            SourceType = iota
	BuildTarballContent            // Build with specified io.Reader as context
	BuildTarballURL                // Build with specified remote URL as context
	DockerPull                     // Run a docker pull to obtain the image
)

type ContainerSource struct {
	Type                SourceType
	buildTarballContent io.Reader
	buildTarballURL     string
	dockerImageName     string
}

func NewContainer(client *docker.Client, name string, wg *sync.WaitGroup) *Container {

	errors := make(chan error)

	c := &Container{
		Name:    name,
		client:  client,
		wg:      wg,
		Errors:  errors,
		errorsW: errors,
	}

	// If the container fails we should assume it should be torn down.
	c.Failed.Forward(&c.Closing)

	return c
}

func (c *Container) Build(config UpdateEvent) error {
	if config.BuildComplete != nil {
		defer close(config.BuildComplete)
	}

	var err error
	bo := docker.BuildImageOptions{}
	bo.Name = c.Name
	bo.OutputStream = config.OutputStream
	if bo.OutputStream == nil {
		bo.OutputStream = os.Stderr
	}

	switch config.Source.Type {
	case BuildCwd:
		bo.ContextDir, err = os.Getwd()
		if err != nil {
			return err
		}
	case BuildTarballContent:
		bo.InputStream = config.Source.buildTarballContent
	default:
		return fmt.Errorf("Unimplemented ContainerSource: %v", config.Source.Type)
	}

	return c.client.BuildImage(bo)
}

func (c *Container) Create() error {
	opts := docker.CreateContainerOptions{
		Name: c.Name,
		Config: &docker.Config{
			Hostname:     c.Name,
			Image:        c.Name,
			AttachStdout: true,
			AttachStderr: true,
			Env:          c.Env,
		},
	}

	var err error
	c.container, err = c.client.CreateContainer(opts)

	return err
}

// CopyOutput copies the output of the container to `w` and blocks until
// completion
func (c *Container) CopyOutput() error {
	// TODO(pwaller): at some point move this on to 'c' for configurability?
	w := os.Stderr
	// Blocks until stream closed
	return c.client.AttachToContainer(docker.AttachToContainerOptions{
		Container:    c.container.ID,
		OutputStream: w,
		ErrorStream:  w,
		Logs:         true,
		Stdout:       true,
		Stderr:       true,
		Stream:       true,
	})
}

// :todo(drj): May want to return errors for truly broken containers (timeout).
func (c *Container) AwaitListening() {

	for _, port := range c.container.NetworkSettings.PortMappingAPI() {
		url := fmt.Sprint("http://", port.IP, ":", port.PublicPort, "/")
		for {
			response, err := http.Get(url)
			if err == nil && response.StatusCode == http.StatusOK {
				break
			}
			time.Sleep(50 * time.Millisecond)
		}
	}
	c.Ready.Fall()
}

func (c *Container) Start() error {
	hc := &docker.HostConfig{
		PublishAllPorts: true,
	}
	err := c.client.StartContainer(c.container.ID, hc)
	if err != nil {
		return err
	}

	// Load container.NetworkSettings
	c.container, err = c.client.InspectContainer(c.container.ID)
	if err != nil {
		return err
	}

	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		<-c.Closing.Barrier()
		// If the container is signaled to close, send a kill signal
		err := c.client.KillContainer(docker.KillContainerOptions{
			ID: c.container.ID,
		})
		if err == nil {
			return
		}
		switch err := err.(type) {
		case *docker.NoSuchContainer:
			// The container already went away, who cares.
			return
		default:
			log.Println("Killing container failed:", c.container.ID, err)
		}
	}()
	return nil
}

func (c *Container) Wait() (int, error) {
	return c.client.WaitContainer(c.container.ID)
}

func (c *Container) err(err error) {
	c.errorsW <- err
	c.Closing.Fall()
}

func (c *Container) Run(event UpdateEvent) (int, error) {

	defer c.Closing.Fall()
	defer close(c.errorsW)

	err := c.Build(event)
	if err != nil {
		return -2, err
	}

	err = c.Create()
	if err != nil {
		return -1, err
	}
	defer c.Delete()

	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		err := c.CopyOutput()
		if err != nil {
			c.err(err)
		}
	}()

	err = c.Start()
	if err != nil {
		return -1, err
	}

	go func() {
		c.AwaitListening()
		c.Ready.Fall()
		log.Println("Listening on", c.container.NetworkSettings.PortMappingAPI())
	}()

	return c.Wait()
}

func (c *Container) Delete() {
	c.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            c.container.ID,
		RemoveVolumes: true,
		Force:         true,
	})
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"
	"strconv"

	"github.com/fsouza/go-dockerclient"
	ui "github.com/gizak/termui"
)

type Widgets struct {
	cid    *ui.Par
	cpu    *ui.Gauge
	memory *ui.Gauge
}

func NewWidgets(id string) *Widgets {
	cid := ui.NewPar(id)
	cid.Border = false
	cid.Height = 2
	cid.Width = 10
	cid.TextFgColor = ui.ColorWhite
	return &Widgets{cid, mkGauge(), mkGauge()}
}

type Container struct {
	id      string
	widgets *Widgets
	stats   chan *docker.Stats
	done    chan bool
}

func NewContainer(cid string) *Container {
	return &Container{
		id:      cid,
		widgets: NewWidgets(cid),
		stats:   make(chan *docker.Stats),
		done:    make(chan bool),
	}
}

func (c *Container) Collect(client *docker.Client) {

	go func() {
		fmt.Sprintf("starting collector for container: %s\n", c.id)
		opts := docker.StatsOptions{
			ID:     c.id,
			Stats:  c.stats,
			Stream: true,
			Done:   c.done,
		}
		client.Stats(opts)
		fmt.Sprintf("stopping collector for container: %s\n", c.id)
	}()

	go func() {
		for s := range c.stats {
			c.UpdateMem(s.MemoryStats.Usage, s.MemoryStats.Limit)
		}
	}()

}

func (c *Container) UpdateCPU(total uint64, system uint64) {
	c.widgets.cpu.BarColor = colorScale(n)
	c.widgets.cpu.Percent = n
}

func (c *Container) UpdateMem(cur uint64, limit uint64) {
	percent := round((float64(cur) / float64(limit)) * 100)
	if percent < 5 {
		percent = 5
	}
	c.widgets.memory.Percent = percent
	c.widgets.memory.Label = fmt.Sprintf("%s / %s", byteFormat(cur), byteFormat(limit))
}

func byteFormat(n uint64) string {
	if n < 1024 {
		return fmt.Sprintf("%sB", strconv.FormatUint(n, 10))
	}
	if n < 1048576 {
		n = n / 1024
		return fmt.Sprintf("%sK", strconv.FormatUint(n, 10))
	}
	if n < 1073741824 {
		n = n / 1048576
		return fmt.Sprintf("%sM", strconv.FormatUint(n, 10))
	}
	n = n / 1024000000
	return fmt.Sprintf("%sG", strconv.FormatUint(n, 10))
}

func round(num float64) int {
	return int(num + math.Copysign(0.5, num))
}

func colorScale(n int) ui.Attribute {
	if n > 70 {
		return ui.ColorRed
	}
	if n > 30 {
		return ui.ColorYellow
	}
	return ui.ColorGreen
}
<commit_msg>add cpucalc, updatecpu method<commit_after>package main

import (
	"fmt"
	"math"
	"strconv"

	"github.com/fsouza/go-dockerclient"
	ui "github.com/gizak/termui"
)

type Widgets struct {
	cid    *ui.Par
	cpu    *ui.Gauge
	memory *ui.Gauge
}

func NewWidgets(id string) *Widgets {
	cid := ui.NewPar(id)
	cid.Border = false
	cid.Height = 2
	cid.Width = 10
	cid.TextFgColor = ui.ColorWhite
	return &Widgets{cid, mkGauge(), mkGauge()}
}

type CpuCalc struct {
	lastCpu    uint64
	lastSysCpu uint64
}

func (c *CpuCalc) Utilization(cpu uint64, syscpu uint64, ncpus int) int {
	cpudiff := float64(cpu) - float64(c.lastCpu)
	syscpudiff := float64(syscpu) - float64(c.lastSysCpu)
	util := round((cpudiff / syscpudiff * 100) * float64(ncpus))
	c.lastCpu = cpu
	c.lastSysCpu = syscpu
	return util
}

type Container struct {
	id      string
	widgets *Widgets
	stats   chan *docker.Stats
	done    chan bool
	cpucalc *CpuCalc
}

func NewContainer(cid string) *Container {
	return &Container{
		id:      cid,
		widgets: NewWidgets(cid),
		stats:   make(chan *docker.Stats),
		done:    make(chan bool),
		cpucalc: &CpuCalc{},
	}
}

func (c *Container) Collect(client *docker.Client) {

	go func() {
		fmt.Sprintf("starting collector for container: %s\n", c.id)
		opts := docker.StatsOptions{
			ID:     c.id,
			Stats:  c.stats,
			Stream: true,
			Done:   c.done,
		}
		client.Stats(opts)
		fmt.Sprintf("stopping collector for container: %s\n", c.id)
	}()

	go func() {
		for s := range c.stats {
			c.UpdateMem(s.MemoryStats.Usage, s.MemoryStats.Limit)
			c.UpdateCPU(s.CPUStats.CPUUsage.TotalUsage, s.CPUStats.SystemCPUUsage, len(s.CPUStats.CPUUsage.PercpuUsage))
		}
	}()

}

func (c *Container) UpdateCPU(total uint64, system uint64, ncpus int) {
	util := c.cpucalc.Utilization(total, system, ncpus)
	c.widgets.cpu.Label = fmt.Sprintf("%s%%", strconv.Itoa(util))
	c.widgets.cpu.BarColor = colorScale(util)
	if util < 5 && util > 0 {
		util = 5
	}
	c.widgets.cpu.Percent = util
}

func (c *Container) UpdateMem(cur uint64, limit uint64) {
	percent := round((float64(cur) / float64(limit)) * 100)
	if percent < 5 {
		percent = 5
	}
	c.widgets.memory.Percent = percent
	c.widgets.memory.Label = fmt.Sprintf("%s / %s", byteFormat(cur), byteFormat(limit))
}

func byteFormat(n uint64) string {
	if n < 1024 {
		return fmt.Sprintf("%sB", strconv.FormatUint(n, 10))
	}
	if n < 1048576 {
		n = n / 1024
		return fmt.Sprintf("%sK", strconv.FormatUint(n, 10))
	}
	if n < 1073741824 {
		n = n / 1048576
		return fmt.Sprintf("%sM", strconv.FormatUint(n, 10))
	}
	n = n / 1024000000
	return fmt.Sprintf("%sG", strconv.FormatUint(n, 10))
}

func round(num float64) int {
	return int(num + math.Copysign(0.5, num))
}

func colorScale(n int) ui.Attribute {
	if n > 70 {
		return ui.ColorRed
	}
	if n > 30 {
		return ui.ColorYellow
	}
	return ui.ColorGreen
}
<|endoftext|>


<commit_before>package copy

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"reflect"
	"strings"

	pb "gopkg.in/cheggaaa/pb.v1"

	"github.com/Sirupsen/logrus"
	"github.com/containers/image/image"
	"github.com/containers/image/signature"
	"github.com/containers/image/transports"
	"github.com/containers/image/types"
)

// supportedDigests lists the supported blob digest types.
var supportedDigests = map[string]func() hash.Hash{
	"sha256": sha256.New,
}

type digestingReader struct {
	source           io.Reader
	digest           hash.Hash
	expectedDigest   []byte
	validationFailed bool
}

// newDigestingReader returns an io.Reader implementation with contents of source, which will eventually return a non-EOF error
// and set validationFailed to true if the source stream does not match expectedDigestString.
func newDigestingReader(source io.Reader, expectedDigestString string) (*digestingReader, error) {
	fields := strings.SplitN(expectedDigestString, ":", 2)
	if len(fields) != 2 {
		return nil, fmt.Errorf("Invalid digest specification %s", expectedDigestString)
	}
	fn, ok := supportedDigests[fields[0]]
	if !ok {
		return nil, fmt.Errorf("Invalid digest specification %s: unknown digest type %s", expectedDigestString, fields[0])
	}
	digest := fn()
	expectedDigest, err := hex.DecodeString(fields[1])
	if err != nil {
		return nil, fmt.Errorf("Invalid digest value %s: %v", expectedDigestString, err)
	}
	if len(expectedDigest) != digest.Size() {
		return nil, fmt.Errorf("Invalid digest specification %s: length %d does not match %d", expectedDigestString, len(expectedDigest), digest.Size())
	}
	return &digestingReader{
		source:           source,
		digest:           digest,
		expectedDigest:   expectedDigest,
		validationFailed: false,
	}, nil
}

func (d *digestingReader) Read(p []byte) (int, error) {
	n, err := d.source.Read(p)
	if n > 0 {
		if n2, err := d.digest.Write(p[:n]); n2 != n || err != nil {
			// Coverage: This should not happen, the hash.Hash interface requires
			// d.digest.Write to never return an error, and the io.Writer interface
			// requires n2 == len(input) if no error is returned.
			return 0, fmt.Errorf("Error updating digest during verification: %d vs. %d, %v", n2, n, err)
		}
	}
	if err == io.EOF {
		actualDigest := d.digest.Sum(nil)
		if subtle.ConstantTimeCompare(actualDigest, d.expectedDigest) != 1 {
			d.validationFailed = true
			return 0, fmt.Errorf("Digest did not match, expected %s, got %s", hex.EncodeToString(d.expectedDigest), hex.EncodeToString(actualDigest))
		}
	}
	return n, err
}

// Options allows supplying non-default configuration modifying the behavior of CopyImage.
type Options struct {
	RemoveSignatures bool   // Remove any pre-existing signatures. SignBy will still add a new signature.
	SignBy           string // If non-empty, asks for a signature to be added during the copy, and specifies a key ID, as accepted by signature.NewGPGSigningMechanism().SignDockerManifest(),
	ReportWriter     io.Writer
}

// Image copies image from srcRef to destRef, using policyContext to validate source image admissibility.
func Image(ctx *types.SystemContext, policyContext *signature.PolicyContext, destRef, srcRef types.ImageReference, options *Options) error {
	if options.ReportWriter == nil {
		options.ReportWriter = ioutil.Discard
	}
	writeReport := func(f string, a ...interface{}) {
		fmt.Fprintf(options.ReportWriter, f, a...)
	}
	dest, err := destRef.NewImageDestination(ctx)
	if err != nil {
		return fmt.Errorf("Error initializing destination %s: %v", transports.ImageName(destRef), err)
	}
	defer dest.Close()

	rawSource, err := srcRef.NewImageSource(ctx, dest.SupportedManifestMIMETypes())
	if err != nil {
		return fmt.Errorf("Error initializing source %s: %v", transports.ImageName(srcRef), err)
	}
	src := image.FromSource(rawSource)
	defer src.Close()

	// Please keep this policy check BEFORE reading any other information about the image.
	if allowed, err := policyContext.IsRunningImageAllowed(src); !allowed || err != nil { // Be paranoid and fail if either return value indicates so.
		return fmt.Errorf("Source image rejected: %v", err)
	}

	writeReport("Getting image source manifest\n")
	manifest, _, err := src.Manifest()
	if err != nil {
		return fmt.Errorf("Error reading manifest: %v", err)
	}

	var sigs [][]byte
	if options != nil && options.RemoveSignatures {
		sigs = [][]byte{}
	} else {
		writeReport("Getting image source signatures\n")
		s, err := src.Signatures()
		if err != nil {
			return fmt.Errorf("Error reading signatures: %v", err)
		}
		sigs = s
	}
	if len(sigs) != 0 {
		writeReport("Checking if image destination supports signatures\n")
		if err := dest.SupportsSignatures(); err != nil {
			return fmt.Errorf("Can not copy signatures: %v", err)
		}
	}
	canModifyManifest := len(sigs) == 0

	writeReport("Getting image source configuration\n")
	srcConfigInfo, err := src.ConfigInfo()
	if err != nil {
		return fmt.Errorf("Error parsing manifest: %v", err)
	}
	if srcConfigInfo.Digest != "" {
		writeReport("Uploading blob %s\n", srcConfigInfo.Digest)
		destConfigInfo, err := copyBlob(dest, rawSource, srcConfigInfo, false, options.ReportWriter)
		if err != nil {
			return err
		}
		if destConfigInfo.Digest != srcConfigInfo.Digest {
			return fmt.Errorf("Internal error: copying uncompressed config blob %s changed digest to %s", srcConfigInfo.Digest, destConfigInfo.Digest)
		}
	}

	srcLayerInfos, err := src.LayerInfos()
	if err != nil {
		return fmt.Errorf("Error parsing manifest: %v", err)
	}
	destLayerInfos := []types.BlobInfo{}
	copiedLayers := map[string]types.BlobInfo{}
	for _, srcLayer := range srcLayerInfos {
		destLayer, ok := copiedLayers[srcLayer.Digest]
		if !ok {
			writeReport("Uploading blob %s\n", srcLayer.Digest)
			destLayer, err = copyBlob(dest, rawSource, srcLayer, canModifyManifest, options.ReportWriter)
			if err != nil {
				return err
			}
			copiedLayers[srcLayer.Digest] = destLayer
		}
		destLayerInfos = append(destLayerInfos, destLayer)
	}

	manifestUpdates := types.ManifestUpdateOptions{}
	if layerDigestsDiffer(srcLayerInfos, destLayerInfos) {
		manifestUpdates.LayerInfos = destLayerInfos
	}

	if !reflect.DeepEqual(manifestUpdates, types.ManifestUpdateOptions{}) {
		if !canModifyManifest {
			return fmt.Errorf("Internal error: copy needs an updated manifest but that was known to be forbidden")
		}
		manifest, err = src.UpdatedManifest(manifestUpdates)
		if err != nil {
			return fmt.Errorf("Error creating an updated manifest: %v", err)
		}
	}

	if options != nil && options.SignBy != "" {
		mech, err := signature.NewGPGSigningMechanism()
		if err != nil {
			return fmt.Errorf("Error initializing GPG: %v", err)
		}
		dockerReference := dest.Reference().DockerReference()
		if dockerReference == nil {
			return fmt.Errorf("Cannot determine canonical Docker reference for destination %s", transports.ImageName(dest.Reference()))
		}

		writeReport("Signing manifest\n")
		newSig, err := signature.SignDockerManifest(manifest, dockerReference.String(), mech, options.SignBy)
		if err != nil {
			return fmt.Errorf("Error creating signature: %v", err)
		}
		sigs = append(sigs, newSig)
	}

	writeReport("Uploading manifest to image destination\n")
	if err := dest.PutManifest(manifest); err != nil {
		return fmt.Errorf("Error writing manifest: %v", err)
	}

	writeReport("Storing signatures\n")
	if err := dest.PutSignatures(sigs); err != nil {
		return fmt.Errorf("Error writing signatures: %v", err)
	}

	if err := dest.Commit(); err != nil {
		return fmt.Errorf("Error committing the finished image: %v", err)
	}

	return nil
}

// layerDigestsDiffer return true iff the digests in a and b differ (ignoring sizes and possible other fields)
func layerDigestsDiffer(a, b []types.BlobInfo) bool {
	if len(a) != len(b) {
		return true
	}
	for i := range a {
		if a[i].Digest != b[i].Digest {
			return true
		}
	}
	return false
}

// copyBlob copies a blob with srcInfo (with known Digest and possibly known Size) in src to dest, perhaps compressing it if canCompress,
// and returns a complete blobInfo of the copied blob.
func copyBlob(dest types.ImageDestination, src types.ImageSource, srcInfo types.BlobInfo, canCompress bool, reportWriter io.Writer) (types.BlobInfo, error) {
	srcStream, srcBlobSize, err := src.GetBlob(srcInfo.Digest) // We currently completely ignore srcInfo.Size throughout.
	if err != nil {
		return types.BlobInfo{}, fmt.Errorf("Error reading blob %s: %v", srcInfo.Digest, err)
	}
	defer srcStream.Close()

	// Be paranoid; in case PutBlob somehow managed to ignore an error from digestingReader,
	// use a separate validation failure indicator.
	// Note that we don't use a stronger "validationSucceeded" indicator, because
	// dest.PutBlob may detect that the layer already exists, in which case we don't
	// read stream to the end, and validation does not happen.
	digestingReader, err := newDigestingReader(srcStream, srcInfo.Digest)
	if err != nil {
		return types.BlobInfo{}, fmt.Errorf("Error preparing to verify blob %s: %v", srcInfo.Digest, err)
	}
	var destStream io.Reader = digestingReader
	isCompressed, destStream, err := isStreamCompressed(destStream) // We could skip this in some cases, but let's keep the code path uniform
	if err != nil {
		return types.BlobInfo{}, fmt.Errorf("Error reading blob %s: %v", srcInfo.Digest, err)
	}

	bar := pb.New(int(srcInfo.Size)).SetUnits(pb.U_BYTES)
	bar.Output = reportWriter
	bar.SetMaxWidth(80)
	bar.ShowTimeLeft = false
	bar.ShowPercent = false
	bar.Start()
	destStream = bar.NewProxyReader(destStream)

	defer func() {
		if err != nil {
			fmt.Fprint(reportWriter, "\n")
		}
	}()

	var inputInfo types.BlobInfo
	if !canCompress || isCompressed || !dest.ShouldCompressLayers() {
		logrus.Debugf("Using original blob without modification")
		inputInfo.Digest = srcInfo.Digest
		inputInfo.Size = srcBlobSize
	} else {
		logrus.Debugf("Compressing blob on the fly")
		pipeReader, pipeWriter := io.Pipe()
		defer pipeReader.Close()

		// If this fails while writing data, it will do pipeWriter.CloseWithError(); if it fails otherwise,
		// e.g. because we have exited and due to pipeReader.Close() above further writing to the pipe has failed,
		// we don’t care.
		go compressGoroutine(pipeWriter, destStream) // Closes pipeWriter
		destStream = pipeReader
		inputInfo.Digest = ""
		inputInfo.Size = -1
	}

	uploadedInfo, err := dest.PutBlob(destStream, inputInfo)
	if err != nil {
		return types.BlobInfo{}, fmt.Errorf("Error writing blob: %v", err)
	}
	if digestingReader.validationFailed { // Coverage: This should never happen.
		return types.BlobInfo{}, fmt.Errorf("Internal error uploading blob %s, digest verification failed but was ignored", srcInfo.Digest)
	}
	if inputInfo.Digest != "" && uploadedInfo.Digest != inputInfo.Digest {
		return types.BlobInfo{}, fmt.Errorf("Internal error uploading blob %s, blob with digest %s uploaded with digest %s", srcInfo.Digest, inputInfo.Digest, uploadedInfo.Digest)
	}
	return uploadedInfo, nil
}

// compressionPrefixes is an internal implementation detail of isStreamCompressed
var compressionPrefixes = map[string][]byte{
	"gzip":  {0x1F, 0x8B, 0x08},                   // gzip (RFC 1952)
	"bzip2": {0x42, 0x5A, 0x68},                   // bzip2 (decompress.c:BZ2_decompress)
	"xz":    {0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00}, // xz (/usr/share/doc/xz/xz-file-format.txt)
}

// isStreamCompressed returns true if input is recognized as a compressed format.
// Because it consumes the start of input, other consumers must use the returned io.Reader instead to also read from the beginning.
func isStreamCompressed(input io.Reader) (bool, io.Reader, error) {
	buffer := [8]byte{}

	n, err := io.ReadAtLeast(input, buffer[:], len(buffer))
	if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
		// This is a “real” error. We could just ignore it this time, process the data we have, and hope that the source will report the same error again.
		// Instead, fail immediately with the original error cause instead of a possibly secondary/misleading error returned later.
		return false, nil, err
	}

	isCompressed := false
	for algo, prefix := range compressionPrefixes {
		if bytes.HasPrefix(buffer[:n], prefix) {
			logrus.Debugf("Detected compression format %s", algo)
			isCompressed = true
			break
		}
	}
	if !isCompressed {
		logrus.Debugf("No compression detected")
	}

	return isCompressed, io.MultiReader(bytes.NewReader(buffer[:n]), input), nil
}

// compressGoroutine reads all input from src and writes its compressed equivalent to dest.
func compressGoroutine(dest *io.PipeWriter, src io.Reader) {
	err := errors.New("Internal error: unexpected panic in compressGoroutine")
	defer func() { // Note that this is not the same as {defer dest.CloseWithError(err)}; we need err to be evaluated lazily.
		dest.CloseWithError(err) // CloseWithError(nil) is equivalent to Close()
	}()

	zipper := gzip.NewWriter(dest)
	defer zipper.Close()

	_, err = io.Copy(zipper, src) // Sets err to nil, i.e. causes dest.Close()
}
<commit_msg>copy: progress fix<commit_after>package copy

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"reflect"
	"strings"

	pb "gopkg.in/cheggaaa/pb.v1"

	"github.com/Sirupsen/logrus"
	"github.com/containers/image/image"
	"github.com/containers/image/signature"
	"github.com/containers/image/transports"
	"github.com/containers/image/types"
)

// supportedDigests lists the supported blob digest types.
var supportedDigests = map[string]func() hash.Hash{
	"sha256": sha256.New,
}

type digestingReader struct {
	source           io.Reader
	digest           hash.Hash
	expectedDigest   []byte
	validationFailed bool
}

// newDigestingReader returns an io.Reader implementation with contents of source, which will eventually return a non-EOF error
// and set validationFailed to true if the source stream does not match expectedDigestString.
func newDigestingReader(source io.Reader, expectedDigestString string) (*digestingReader, error) {
	fields := strings.SplitN(expectedDigestString, ":", 2)
	if len(fields) != 2 {
		return nil, fmt.Errorf("Invalid digest specification %s", expectedDigestString)
	}
	fn, ok := supportedDigests[fields[0]]
	if !ok {
		return nil, fmt.Errorf("Invalid digest specification %s: unknown digest type %s", expectedDigestString, fields[0])
	}
	digest := fn()
	expectedDigest, err := hex.DecodeString(fields[1])
	if err != nil {
		return nil, fmt.Errorf("Invalid digest value %s: %v", expectedDigestString, err)
	}
	if len(expectedDigest) != digest.Size() {
		return nil, fmt.Errorf("Invalid digest specification %s: length %d does not match %d", expectedDigestString, len(expectedDigest), digest.Size())
	}
	return &digestingReader{
		source:           source,
		digest:           digest,
		expectedDigest:   expectedDigest,
		validationFailed: false,
	}, nil
}

func (d *digestingReader) Read(p []byte) (int, error) {
	n, err := d.source.Read(p)
	if n > 0 {
		if n2, err := d.digest.Write(p[:n]); n2 != n || err != nil {
			// Coverage: This should not happen, the hash.Hash interface requires
			// d.digest.Write to never return an error, and the io.Writer interface
			// requires n2 == len(input) if no error is returned.
			return 0, fmt.Errorf("Error updating digest during verification: %d vs. %d, %v", n2, n, err)
		}
	}
	if err == io.EOF {
		actualDigest := d.digest.Sum(nil)
		if subtle.ConstantTimeCompare(actualDigest, d.expectedDigest) != 1 {
			d.validationFailed = true
			return 0, fmt.Errorf("Digest did not match, expected %s, got %s", hex.EncodeToString(d.expectedDigest), hex.EncodeToString(actualDigest))
		}
	}
	return n, err
}

// Options allows supplying non-default configuration modifying the behavior of CopyImage.
type Options struct {
	RemoveSignatures bool   // Remove any pre-existing signatures. SignBy will still add a new signature.
	SignBy           string // If non-empty, asks for a signature to be added during the copy, and specifies a key ID, as accepted by signature.NewGPGSigningMechanism().SignDockerManifest(),
	ReportWriter     io.Writer
}

// Image copies image from srcRef to destRef, using policyContext to validate source image admissibility.
func Image(ctx *types.SystemContext, policyContext *signature.PolicyContext, destRef, srcRef types.ImageReference, options *Options) error {
	if options.ReportWriter == nil {
		options.ReportWriter = ioutil.Discard
	}
	writeReport := func(f string, a ...interface{}) {
		fmt.Fprintf(options.ReportWriter, f, a...)
	}
	dest, err := destRef.NewImageDestination(ctx)
	if err != nil {
		return fmt.Errorf("Error initializing destination %s: %v", transports.ImageName(destRef), err)
	}
	defer dest.Close()

	rawSource, err := srcRef.NewImageSource(ctx, dest.SupportedManifestMIMETypes())
	if err != nil {
		return fmt.Errorf("Error initializing source %s: %v", transports.ImageName(srcRef), err)
	}
	src := image.FromSource(rawSource)
	defer src.Close()

	// Please keep this policy check BEFORE reading any other information about the image.
	if allowed, err := policyContext.IsRunningImageAllowed(src); !allowed || err != nil { // Be paranoid and fail if either return value indicates so.
		return fmt.Errorf("Source image rejected: %v", err)
	}

	writeReport("Getting image source manifest\n")
	manifest, _, err := src.Manifest()
	if err != nil {
		return fmt.Errorf("Error reading manifest: %v", err)
	}

	var sigs [][]byte
	if options != nil && options.RemoveSignatures {
		sigs = [][]byte{}
	} else {
		writeReport("Getting image source signatures\n")
		s, err := src.Signatures()
		if err != nil {
			return fmt.Errorf("Error reading signatures: %v", err)
		}
		sigs = s
	}
	if len(sigs) != 0 {
		writeReport("Checking if image destination supports signatures\n")
		if err := dest.SupportsSignatures(); err != nil {
			return fmt.Errorf("Can not copy signatures: %v", err)
		}
	}
	canModifyManifest := len(sigs) == 0

	writeReport("Getting image source configuration\n")
	srcConfigInfo, err := src.ConfigInfo()
	if err != nil {
		return fmt.Errorf("Error parsing manifest: %v", err)
	}
	if srcConfigInfo.Digest != "" {
		writeReport("Uploading blob %s\n", srcConfigInfo.Digest)
		destConfigInfo, err := copyBlob(dest, rawSource, srcConfigInfo, false, options.ReportWriter)
		if err != nil {
			return err
		}
		if destConfigInfo.Digest != srcConfigInfo.Digest {
			return fmt.Errorf("Internal error: copying uncompressed config blob %s changed digest to %s", srcConfigInfo.Digest, destConfigInfo.Digest)
		}
	}

	srcLayerInfos, err := src.LayerInfos()
	if err != nil {
		return fmt.Errorf("Error parsing manifest: %v", err)
	}
	destLayerInfos := []types.BlobInfo{}
	copiedLayers := map[string]types.BlobInfo{}
	for _, srcLayer := range srcLayerInfos {
		destLayer, ok := copiedLayers[srcLayer.Digest]
		if !ok {
			writeReport("Uploading blob %s\n", srcLayer.Digest)
			destLayer, err = copyBlob(dest, rawSource, srcLayer, canModifyManifest, options.ReportWriter)
			if err != nil {
				return err
			}
			copiedLayers[srcLayer.Digest] = destLayer
		}
		destLayerInfos = append(destLayerInfos, destLayer)
	}

	manifestUpdates := types.ManifestUpdateOptions{}
	if layerDigestsDiffer(srcLayerInfos, destLayerInfos) {
		manifestUpdates.LayerInfos = destLayerInfos
	}

	if !reflect.DeepEqual(manifestUpdates, types.ManifestUpdateOptions{}) {
		if !canModifyManifest {
			return fmt.Errorf("Internal error: copy needs an updated manifest but that was known to be forbidden")
		}
		manifest, err = src.UpdatedManifest(manifestUpdates)
		if err != nil {
			return fmt.Errorf("Error creating an updated manifest: %v", err)
		}
	}

	if options != nil && options.SignBy != "" {
		mech, err := signature.NewGPGSigningMechanism()
		if err != nil {
			return fmt.Errorf("Error initializing GPG: %v", err)
		}
		dockerReference := dest.Reference().DockerReference()
		if dockerReference == nil {
			return fmt.Errorf("Cannot determine canonical Docker reference for destination %s", transports.ImageName(dest.Reference()))
		}

		writeReport("Signing manifest\n")
		newSig, err := signature.SignDockerManifest(manifest, dockerReference.String(), mech, options.SignBy)
		if err != nil {
			return fmt.Errorf("Error creating signature: %v", err)
		}
		sigs = append(sigs, newSig)
	}

	writeReport("Uploading manifest to image destination\n")
	if err := dest.PutManifest(manifest); err != nil {
		return fmt.Errorf("Error writing manifest: %v", err)
	}

	writeReport("Storing signatures\n")
	if err := dest.PutSignatures(sigs); err != nil {
		return fmt.Errorf("Error writing signatures: %v", err)
	}

	if err := dest.Commit(); err != nil {
		return fmt.Errorf("Error committing the finished image: %v", err)
	}

	return nil
}

// layerDigestsDiffer return true iff the digests in a and b differ (ignoring sizes and possible other fields)
func layerDigestsDiffer(a, b []types.BlobInfo) bool {
	if len(a) != len(b) {
		return true
	}
	for i := range a {
		if a[i].Digest != b[i].Digest {
			return true
		}
	}
	return false
}

// copyBlob copies a blob with srcInfo (with known Digest and possibly known Size) in src to dest, perhaps compressing it if canCompress,
// and returns a complete blobInfo of the copied blob.
func copyBlob(dest types.ImageDestination, src types.ImageSource, srcInfo types.BlobInfo, canCompress bool, reportWriter io.Writer) (types.BlobInfo, error) {
	srcStream, srcBlobSize, err := src.GetBlob(srcInfo.Digest) // We currently completely ignore srcInfo.Size throughout.
	if err != nil {
		return types.BlobInfo{}, fmt.Errorf("Error reading blob %s: %v", srcInfo.Digest, err)
	}
	defer srcStream.Close()

	// Be paranoid; in case PutBlob somehow managed to ignore an error from digestingReader,
	// use a separate validation failure indicator.
	// Note that we don't use a stronger "validationSucceeded" indicator, because
	// dest.PutBlob may detect that the layer already exists, in which case we don't
	// read stream to the end, and validation does not happen.
	digestingReader, err := newDigestingReader(srcStream, srcInfo.Digest)
	if err != nil {
		return types.BlobInfo{}, fmt.Errorf("Error preparing to verify blob %s: %v", srcInfo.Digest, err)
	}
	var destStream io.Reader = digestingReader
	isCompressed, destStream, err := isStreamCompressed(destStream) // We could skip this in some cases, but let's keep the code path uniform
	if err != nil {
		return types.BlobInfo{}, fmt.Errorf("Error reading blob %s: %v", srcInfo.Digest, err)
	}

	bar := pb.New(int(srcInfo.Size)).SetUnits(pb.U_BYTES)
	bar.Output = reportWriter
	bar.SetMaxWidth(80)
	bar.ShowTimeLeft = false
	bar.ShowPercent = false
	bar.Start()
	destStream = bar.NewProxyReader(destStream)

	defer fmt.Fprint(reportWriter, "\n")

	var inputInfo types.BlobInfo
	if !canCompress || isCompressed || !dest.ShouldCompressLayers() {
		logrus.Debugf("Using original blob without modification")
		inputInfo.Digest = srcInfo.Digest
		inputInfo.Size = srcBlobSize
	} else {
		logrus.Debugf("Compressing blob on the fly")
		pipeReader, pipeWriter := io.Pipe()
		defer pipeReader.Close()

		// If this fails while writing data, it will do pipeWriter.CloseWithError(); if it fails otherwise,
		// e.g. because we have exited and due to pipeReader.Close() above further writing to the pipe has failed,
		// we don’t care.
		go compressGoroutine(pipeWriter, destStream) // Closes pipeWriter
		destStream = pipeReader
		inputInfo.Digest = ""
		inputInfo.Size = -1
	}

	uploadedInfo, err := dest.PutBlob(destStream, inputInfo)
	if err != nil {
		return types.BlobInfo{}, fmt.Errorf("Error writing blob: %v", err)
	}
	if digestingReader.validationFailed { // Coverage: This should never happen.
		return types.BlobInfo{}, fmt.Errorf("Internal error uploading blob %s, digest verification failed but was ignored", srcInfo.Digest)
	}
	if inputInfo.Digest != "" && uploadedInfo.Digest != inputInfo.Digest {
		return types.BlobInfo{}, fmt.Errorf("Internal error uploading blob %s, blob with digest %s uploaded with digest %s", srcInfo.Digest, inputInfo.Digest, uploadedInfo.Digest)
	}
	return uploadedInfo, nil
}

// compressionPrefixes is an internal implementation detail of isStreamCompressed
var compressionPrefixes = map[string][]byte{
	"gzip":  {0x1F, 0x8B, 0x08},                   // gzip (RFC 1952)
	"bzip2": {0x42, 0x5A, 0x68},                   // bzip2 (decompress.c:BZ2_decompress)
	"xz":    {0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00}, // xz (/usr/share/doc/xz/xz-file-format.txt)
}

// isStreamCompressed returns true if input is recognized as a compressed format.
// Because it consumes the start of input, other consumers must use the returned io.Reader instead to also read from the beginning.
func isStreamCompressed(input io.Reader) (bool, io.Reader, error) {
	buffer := [8]byte{}

	n, err := io.ReadAtLeast(input, buffer[:], len(buffer))
	if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
		// This is a “real” error. We could just ignore it this time, process the data we have, and hope that the source will report the same error again.
		// Instead, fail immediately with the original error cause instead of a possibly secondary/misleading error returned later.
		return false, nil, err
	}

	isCompressed := false
	for algo, prefix := range compressionPrefixes {
		if bytes.HasPrefix(buffer[:n], prefix) {
			logrus.Debugf("Detected compression format %s", algo)
			isCompressed = true
			break
		}
	}
	if !isCompressed {
		logrus.Debugf("No compression detected")
	}

	return isCompressed, io.MultiReader(bytes.NewReader(buffer[:n]), input), nil
}

// compressGoroutine reads all input from src and writes its compressed equivalent to dest.
func compressGoroutine(dest *io.PipeWriter, src io.Reader) {
	err := errors.New("Internal error: unexpected panic in compressGoroutine")
	defer func() { // Note that this is not the same as {defer dest.CloseWithError(err)}; we need err to be evaluated lazily.
		dest.CloseWithError(err) // CloseWithError(nil) is equivalent to Close()
	}()

	zipper := gzip.NewWriter(dest)
	defer zipper.Close()

	_, err = io.Copy(zipper, src) // Sets err to nil, i.e. causes dest.Close()
}
<|endoftext|>


<commit_before>package ares

import (
	"fmt"
	"log"

	"github.com/nlopes/slack"
)

type Ares struct {
	SlackAppToken string
	SlackBotToken string
	SlackAppID    string
	ImgurClientID string
	BotUserID     string
	// Bot won't be added or re-added to these channels
	ManagedChannels []string
	// list of userIDs of admins
	Admins []string
	// maintain a user dict
	Users map[string]string
	// Muted users list
	MutedUsers map[string]bool
}

func (a *Ares) initBot() {
	api := slack.New(a.SlackAppToken)
	a.getBotandAdmin()

	channels, err := api.GetChannels(true)

	if err != nil {
		log.Fatal("Failed to get public channels list: ", err.Error())
	}

	for _, channel := range channels {
		a.addBotChannel(channel.ID)
	}

	groups, err := api.GetGroups(true)

	if err != nil {
		log.Fatal("Failed to get private channels list: ", err.Error())
	}

	for _, group := range groups {
		a.addBotGroup(group.ID)
	}
}

// Fetches the bot user ID and also saves the admin user ids
func (a *Ares) getBotandAdmin() {
	api := slack.New(a.SlackAppToken)
	users, err := api.GetUsers()

	if err != nil {
		log.Fatal("Failed to fetch slack users info", err.Error())
	}

	a.Users = make(map[string]string)
	a.MutedUsers = make(map[string]bool)

	for _, user := range users {

		if user.Profile.ApiAppID == a.SlackAppID {
			a.BotUserID = user.ID
		}

		if user.IsAdmin {
			a.Admins = append(a.Admins, user.ID)
		} else {
			a.Users[user.Name] = user.ID
		}
	}

	if a.BotUserID == "" {
		log.Fatal("Unable to find bot user on the Slack")
	}
}

func (a *Ares) deleteMsg(channel, msgTimestamp string) {
	api := slack.New(a.SlackAppToken)

	if _, _, err := api.DeleteMessage(channel, msgTimestamp); err != nil {
		log.Println("Failed to delete msg: ", err.Error())
	}
}

func (a *Ares) deleteFile(fileId string) {
	api := slack.New(a.SlackAppToken)

	if err := api.DeleteFile(fileId); err != nil {
		fmt.Println("Deletion of the file failed:", err.Error())
	}
}

func (a *Ares) addBotChannel(channelID string) {
	api := slack.New(a.SlackAppToken)
	if _, err := api.InviteUserToChannel(channelID, a.BotUserID); err != nil {
		if err.Error() != "already_in_channel" {
			log.Println(fmt.Sprintf("Failed to add bot to %s: %s", channelID, err.Error()))
		}
	}
}

func (a *Ares) addBotGroup(group string) {
	api := slack.New(a.SlackAppToken)
	if _, _, err := api.InviteUserToGroup(group, a.BotUserID); err != nil {
		log.Println(fmt.Sprintf("Failed to add bot to %s: %s", group, err.Error()))
	}
}

func (a *Ares) notifyUser(user, deleteLink string) {
	api := slack.New(a.SlackBotToken)
	params := slack.PostMessageParameters{}
	attachment := slack.Attachment{
		Color:   "danger",
		Pretext: "Hey there, image uploads to this slack is disabled. Please use external image upload services next time. This incident has been reported :squirrel:",
		Text:    "I have uploaded the image to Imgur for now.",
		Fields: []slack.AttachmentField{
			{
				Title: "If you wish to delete the image uploaded, click on the following link",
				Value: deleteLink,
			},
		},
	}
	params.Attachments = []slack.Attachment{attachment}
	_, _, err := api.PostMessage(user, "", params)
	if err != nil {
		log.Println("Failed to send DM to user", err.Error())
	}
}

func (a *Ares) sendImgToSlack(channel, user, url, commentText string) {
	api := slack.New(a.SlackBotToken)
	params := slack.PostMessageParameters{}
	attachment := slack.Attachment{
		Color:    "#D3D3D3",
		Text:     fmt.Sprintf("image originally uploaded by <@%s>", user),
		ImageURL: url,
		Fields: []slack.AttachmentField{
			{
				Title: "",
				Value: commentText,
			},
		},
	}
	params.Attachments = []slack.Attachment{attachment}
	_, _, err := api.PostMessage(channel, "", params)
	if err != nil {
		log.Println("Failed to send DM to user", err.Error())
	}
}

func (a *Ares) handleFile(file *slack.File, channel string) {
	commentText := file.InitialComment.Comment

	resp := uploadToImgur(file.URLPrivateDownload, a.SlackAppToken, a.ImgurClientID)
	if resp.Status == false {
		log.Println("Failed to download/upload: ", file)
		log.Println("Download url: ", file.URLPrivateDownload)
		return
	}

	a.notifyUser(file.User, resp.DeleteLink)
	a.sendImgToSlack(channel, file.User, resp.Link, commentText)
	a.deleteFile(file.ID)
}

func isImageFile(fileType string) bool {
	for _, fType := range []string{"jpg", "jpeg", "png", "gif"} {
		if fType == fileType {
			return true
		}
	}
	return false
}

func (a *Ares) Run() {
	api := slack.New(a.SlackBotToken)

	a.initBot()
	log.Println("Bot initialized. Starting moderation duty.")

	rtm := api.NewRTM()
	go rtm.ManageConnection()

	for msg := range rtm.IncomingEvents {
		switch ev := msg.Data.(type) {

		case *slack.MessageEvent:

			if a.isAdminUser(ev.Msg.User) {
				a.performMuteAction(ev.Msg.Text)
			} else if a.isMuted(ev.Msg.User) {
				go a.deleteMsg(ev.Msg.Channel, ev.Msg.Timestamp)
			}

			if ev.SubType == "file_share" {
				if isImageFile(ev.File.Filetype) {
					a.handleFile(ev.File, ev.Channel)
				}
			}

		case *slack.TeamJoinEvent:
			if !ev.User.IsBot {
				a.onBoardUser(ev.User)
			}

		case *slack.GroupLeftEvent:
			log.Println("Bot was removed from private channel: ", ev.Channel)
			a.addBotGroup(ev.Channel)

		case *slack.ChannelLeftEvent:
			log.Println("Bot was removed from public channel: ", ev.Channel)
			a.addBotChannel(ev.Channel)

		case *slack.RTMError:
			fmt.Printf("Error: %s\n", ev.Error())

		case *slack.InvalidAuthEvent:
			fmt.Printf("Invalid credentials")
			return

		default:
			// Ignore other events..
			// fmt.Printf("Unexpected: %v\n", msg.Data)
		}
	}
}
<commit_msg>Use go routines for each image<commit_after>package ares

import (
	"fmt"
	"log"

	"github.com/nlopes/slack"
)

type Ares struct {
	SlackAppToken string
	SlackBotToken string
	SlackAppID    string
	ImgurClientID string
	BotUserID     string
	// Bot won't be added or re-added to these channels
	ManagedChannels []string
	// list of userIDs of admins
	Admins []string
	// maintain a user dict
	Users map[string]string
	// Muted users list
	MutedUsers map[string]bool
}

func (a *Ares) initBot() {
	api := slack.New(a.SlackAppToken)
	a.getBotandAdmin()

	channels, err := api.GetChannels(true)

	if err != nil {
		log.Fatal("Failed to get public channels list: ", err.Error())
	}

	for _, channel := range channels {
		a.addBotChannel(channel.ID)
	}

	groups, err := api.GetGroups(true)

	if err != nil {
		log.Fatal("Failed to get private channels list: ", err.Error())
	}

	for _, group := range groups {
		a.addBotGroup(group.ID)
	}
}

// Fetches the bot user ID and also saves the admin user ids
func (a *Ares) getBotandAdmin() {
	api := slack.New(a.SlackAppToken)
	users, err := api.GetUsers()

	if err != nil {
		log.Fatal("Failed to fetch slack users info", err.Error())
	}

	a.Users = make(map[string]string)
	a.MutedUsers = make(map[string]bool)

	for _, user := range users {

		if user.Profile.ApiAppID == a.SlackAppID {
			a.BotUserID = user.ID
		}

		if user.IsAdmin {
			a.Admins = append(a.Admins, user.ID)
		} else {
			a.Users[user.Name] = user.ID
		}
	}

	if a.BotUserID == "" {
		log.Fatal("Unable to find bot user on the Slack")
	}
}

func (a *Ares) deleteMsg(channel, msgTimestamp string) {
	api := slack.New(a.SlackAppToken)

	if _, _, err := api.DeleteMessage(channel, msgTimestamp); err != nil {
		log.Println("Failed to delete msg: ", err.Error())
	}
}

func (a *Ares) deleteFile(fileId string) {
	api := slack.New(a.SlackAppToken)

	if err := api.DeleteFile(fileId); err != nil {
		fmt.Println("Deletion of the file failed:", err.Error())
	}
}

func (a *Ares) addBotChannel(channelID string) {
	api := slack.New(a.SlackAppToken)
	if _, err := api.InviteUserToChannel(channelID, a.BotUserID); err != nil {
		if err.Error() != "already_in_channel" {
			log.Println(fmt.Sprintf("Failed to add bot to %s: %s", channelID, err.Error()))
		}
	}
}

func (a *Ares) addBotGroup(group string) {
	api := slack.New(a.SlackAppToken)
	if _, _, err := api.InviteUserToGroup(group, a.BotUserID); err != nil {
		log.Println(fmt.Sprintf("Failed to add bot to %s: %s", group, err.Error()))
	}
}

func (a *Ares) notifyUser(user, deleteLink string) {
	api := slack.New(a.SlackBotToken)
	params := slack.PostMessageParameters{}
	attachment := slack.Attachment{
		Color:   "danger",
		Pretext: "Hey there, image uploads to this slack is disabled. Please use external image upload services next time. This incident has been reported :squirrel:",
		Text:    "I have uploaded the image to Imgur for now.",
		Fields: []slack.AttachmentField{
			{
				Title: "If you wish to delete the image uploaded, click on the following link",
				Value: deleteLink,
			},
		},
	}
	params.Attachments = []slack.Attachment{attachment}
	_, _, err := api.PostMessage(user, "", params)
	if err != nil {
		log.Println("Failed to send DM to user", err.Error())
	}
}

func (a *Ares) sendImgToSlack(channel, user, url, commentText string) {
	api := slack.New(a.SlackBotToken)
	params := slack.PostMessageParameters{}
	attachment := slack.Attachment{
		Color:    "#D3D3D3",
		Text:     fmt.Sprintf("image originally uploaded by <@%s>", user),
		ImageURL: url,
		Fields: []slack.AttachmentField{
			{
				Title: "",
				Value: commentText,
			},
		},
	}
	params.Attachments = []slack.Attachment{attachment}
	_, _, err := api.PostMessage(channel, "", params)
	if err != nil {
		log.Println("Failed to send DM to user", err.Error())
	}
}

func (a *Ares) handleFile(file *slack.File, channel string) {
	commentText := file.InitialComment.Comment

	resp := uploadToImgur(file.URLPrivateDownload, a.SlackAppToken, a.ImgurClientID)
	if resp.Status == false {
		log.Println("Failed to download/upload: ", file)
		log.Println("Download url: ", file.URLPrivateDownload)
		return
	}

	a.notifyUser(file.User, resp.DeleteLink)
	a.sendImgToSlack(channel, file.User, resp.Link, commentText)
	a.deleteFile(file.ID)
}

func isImageFile(fileType string) bool {
	for _, fType := range []string{"jpg", "jpeg", "png", "gif"} {
		if fType == fileType {
			return true
		}
	}
	return false
}

func (a *Ares) Run() {
	api := slack.New(a.SlackBotToken)

	a.initBot()
	log.Println("Bot initialized. Starting moderation duty.")

	rtm := api.NewRTM()
	go rtm.ManageConnection()

	for msg := range rtm.IncomingEvents {
		switch ev := msg.Data.(type) {

		case *slack.MessageEvent:

			if a.isAdminUser(ev.Msg.User) {
				a.performMuteAction(ev.Msg.Text)
			} else if a.isMuted(ev.Msg.User) {
				go a.deleteMsg(ev.Msg.Channel, ev.Msg.Timestamp)
			}

			if ev.SubType == "file_share" {
				if isImageFile(ev.File.Filetype) {
					// TODO: Use worker pool
					go a.handleFile(ev.File, ev.Channel)
				}
			}

		case *slack.TeamJoinEvent:
			if !ev.User.IsBot {
				a.onBoardUser(ev.User)
			}

		case *slack.GroupLeftEvent:
			log.Println("Bot was removed from private channel: ", ev.Channel)
			a.addBotGroup(ev.Channel)

		case *slack.ChannelLeftEvent:
			log.Println("Bot was removed from public channel: ", ev.Channel)
			a.addBotChannel(ev.Channel)

		case *slack.RTMError:
			fmt.Printf("Error: %s\n", ev.Error())

		case *slack.InvalidAuthEvent:
			fmt.Printf("Invalid credentials")
			return

		default:
			// Ignore other events..
			// fmt.Printf("Unexpected: %v\n", msg.Data)
		}
	}
}
<|endoftext|>


<commit_before>package core

import (
	"bufio"
	"io/ioutil"
	"os"
	"strings"

	"github.com/tinycedar/lily/common"
	"github.com/tinycedar/lily/conf"
)

const (
	systemHosts = "C:/Windows/System32/drivers/etc/hosts"
)

// var batcher *Batcher

func FireHostsSwitch() bool {
	common.Info("============================== Fire hosts switch ==============================")
	// if batcher != nil {
	// 	batcher.Close()
	// }
	return doProcess()
	// batcher = initSystemHostsWatcher()
	// go startSystemHostsWatcher()
}

// 1. Find collection of same domain names between system hosts and currentHostIndex
// 2. Disconnect the TCP connections(http:80 & https:443) of collection found above
func doProcess() bool {
	success := true
	hostsIPMap := getHostsIpMap()
	overwriteSystemHosts()
	processNameMap := GetProcessNameMap()
	table := getTCPTable()
	for i := uint32(0); i < uint32(table.dwNumEntries); i++ {
		row := table.table[i]
		ip := row.displayIP(row.dwRemoteAddr)
		port := row.displayPort(row.dwRemotePort)
		if row.dwOwningPid <= 0 {
			continue
		}
		if port != 80 && port != 443 {
			continue
		}
		process := strings.ToLower(processNameMap[uint32(row.dwOwningPid)])
		if hostsIPMap[ip] || common.BrowserMap[process] {
			if err := CloseTCPEntry(row); err != nil {
				common.Error("Fail to close TCP connections: Process = %v, Pid = %v, Addr = %v:%v", process, row.dwOwningPid, ip, port)
				success = false
			} else {
				common.Info("Succeed to close TCP connections: Process = %v, Pid = %v, Addr = %v:%v", process, row.dwOwningPid, ip, port)
			}
		}
	}
	return success
}

// find all the ip of system and current hosts
func getHostsIpMap() map[string]bool {
	ipMap := make(map[string]bool)
	for _, v := range readHostConfigMap(systemHosts) {
		ipMap[v] = true
	}
	model := conf.Config.HostConfigModel
	index := conf.Config.CurrentHostIndex
	if length := len(model.Roots); index >= 0 && length > 0 && index < length {
		path := "conf/hosts/" + model.RootAt(index).Text() + ".hosts"
		for _, v := range readHostConfigMap(path) {
			ipMap[v] = true
		}
	}
	return ipMap
}

func overwriteSystemHosts() {
	bytes := ReadCurrentHostConfig()
	if bytes == nil {
		return
	}
	if err := ioutil.WriteFile(systemHosts, bytes, os.ModeExclusive); err != nil {
		common.Error("Error writing to system hosts file: ", err)
	}
}

func ReadCurrentHostConfig() []byte {
	model := conf.Config.HostConfigModel
	index := conf.Config.CurrentHostIndex
	if length := len(model.Roots); index < 0 || length <= 0 || index >= length {
		return nil
	}
	bytes, err := ioutil.ReadFile("conf/hosts/" + model.RootAt(index).Text() + ".hosts")
	if err != nil {
		common.Info("Error reading host config: %v", err)
		return nil
	}
	return bytes
}

func process() {
	// hostConfigMap := readFile()
	browserProcessMap := GetProcessNameMap()
	table := getTCPTable()
	// group by process
	tcpRowByProcessNameMap := make(map[string][]*MIB_TCPROW2)
	for i := uint32(0); i < uint32(table.dwNumEntries); i++ {
		row := table.table[i]
		if row.dwOwningPid <= 0 {
			continue
		}
		// remoteAddr := row.displayIP(row.dwRemoteAddr)
		// if _, ok := hostConfigMap[remoteAddr]; ok {
		// common.Info("====== remoteAddr= %v\tbrowserProcessMap = %v\tpid = %v", remoteAddr, browserProcessMap, row.dwOwningPid)
		if processName, ok := browserProcessMap[uint32(row.dwOwningPid)]; ok {
			pidSlice, ok := tcpRowByProcessNameMap[processName]
			if !ok {
				pidSlice = []*MIB_TCPROW2{}
			}
			pidSlice = append(pidSlice, row)
			tcpRowByProcessNameMap[processName] = pidSlice
		}
		// }
		// common.Info("\t%-6d\t%s:%-16d\t%s:%-16d\t%d\t%d\n", row.dwState, row.displayIP(row.dwLocalAddr), row.displayPort(row.dwLocalPort), row.displayIP(row.dwRemoteAddr), row.displayPort(row.dwRemotePort), row.dwOwningPid, row.dwOffloadState)
	}
	browsers := []string{}
	for k := range tcpRowByProcessNameMap {
		browsers = append(browsers, k)
	}
	common.Info("Browsers: %v", browsers)
	for processName, rowSlice := range tcpRowByProcessNameMap {
		success := true
		for _, row := range rowSlice {
			if err := CloseTCPEntry(row); err != nil {
				success = false
				common.Error("Fail to close TCP connections: %s, %v\n", processName, row.dwOwningPid)
			}
		}
		if success {
			common.Info("Succeed to close TCP connections: %s", processName)
		}
	}
}

func readHostConfigMap(path string) map[string]string {
	hostConfigMap := make(map[string]string)
	file, err := os.Open(path)
	if err != nil {
		common.Info("Fail to open system_hosts: %s", err)
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && !strings.HasPrefix(line, "#") {
			config := strings.Fields(line)
			if len(config) == 2 {
				hostConfigMap[config[1]] = config[0]
			}
		}
	}
	if err := scanner.Err(); err != nil {
		common.Error("Fail to read system_hosts: %s", err)
	}
	return hostConfigMap
}

// func initSystemHostsWatcher() *Batcher {
// 	batcher, err := New(time.Millisecond * 300)
// 	if err != nil {
// 		common.Error("Fail to initialize batcher")
// 	}
// 	if err = batcher.Add(systemHosts); err != nil {
// 		common.Error("Fail to add system hosts: %s", systemHosts)
// 	}
// 	return batcher
// }

// func startSystemHostsWatcher() {
// 	if batcher == nil {
// 		common.Error("Fail to start system hosts watcher, watcher is nil")
// 		return
// 	}
// 	for events := range batcher.Events {
// 		for _, event := range events {
// 			if event.Op&fsnotify.Write == fsnotify.Write {
// 				common.Info("modified file: %v", event)
// 				doProcess()
// 				break
// 			}
// 		}
// 	}
// 	// never return
// }
<commit_msg>Accomplish single process feature<commit_after>package core

import (
	"bufio"
	"io/ioutil"
	"os"
	"strings"

	"github.com/tinycedar/lily/common"
	"github.com/tinycedar/lily/conf"
)

const (
	systemHosts = "C:/Windows/System32/drivers/etc/hosts"
)

// var batcher *Batcher

func FireHostsSwitch() bool {
	common.Info("============================== Fire hosts switch ==============================")
	// if batcher != nil {
	// 	batcher.Close()
	// }
	return doProcess()
	// batcher = initSystemHostsWatcher()
	// go startSystemHostsWatcher()
}

func doProcess() bool {
	success := true
	hostsIPMap := getHostsIpMap()
	overwriteSystemHosts()
	processNameMap := GetProcessNameMap()
	table := getTCPTable()
	for i := uint32(0); i < uint32(table.dwNumEntries); i++ {
		row := table.table[i]
		ip := row.displayIP(row.dwRemoteAddr)
		port := row.displayPort(row.dwRemotePort)
		if row.dwOwningPid <= 0 {
			continue
		}
		if port != 80 && port != 443 {
			continue
		}
		process := strings.ToLower(processNameMap[uint32(row.dwOwningPid)])
		if hostsIPMap[ip] || common.BrowserMap[process] {
			if err := CloseTCPEntry(row); err != nil {
				common.Error("Fail to close TCP connections: Process = %v, Pid = %v, Addr = %v:%v", process, row.dwOwningPid, ip, port)
				success = false
			} else {
				common.Info("Succeed to close TCP connections: Process = %v, Pid = %v, Addr = %v:%v", process, row.dwOwningPid, ip, port)
			}
		}
	}
	return success
}

// find all the ip of system and current hosts
func getHostsIpMap() map[string]bool {
	ipMap := make(map[string]bool)
	for _, v := range readHostConfigMap(systemHosts) {
		ipMap[v] = true
	}
	model := conf.Config.HostConfigModel
	index := conf.Config.CurrentHostIndex
	if length := len(model.Roots); index >= 0 && length > 0 && index < length {
		path := "conf/hosts/" + model.RootAt(index).Text() + ".hosts"
		for _, v := range readHostConfigMap(path) {
			ipMap[v] = true
		}
	}
	return ipMap
}

func overwriteSystemHosts() {
	bytes := ReadCurrentHostConfig()
	if bytes == nil {
		return
	}
	if err := ioutil.WriteFile(systemHosts, bytes, os.ModeExclusive); err != nil {
		common.Error("Error writing to system hosts file: ", err)
	}
}

func ReadCurrentHostConfig() []byte {
	model := conf.Config.HostConfigModel
	index := conf.Config.CurrentHostIndex
	if length := len(model.Roots); index < 0 || length <= 0 || index >= length {
		return nil
	}
	bytes, err := ioutil.ReadFile("conf/hosts/" + model.RootAt(index).Text() + ".hosts")
	if err != nil {
		common.Info("Error reading host config: %v", err)
		return nil
	}
	return bytes
}

func readHostConfigMap(path string) map[string]string {
	hostConfigMap := make(map[string]string)
	file, err := os.Open(path)
	if err != nil {
		common.Info("Fail to open system_hosts: %s", err)
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && !strings.HasPrefix(line, "#") {
			config := strings.Fields(line)
			if len(config) == 2 {
				hostConfigMap[config[1]] = config[0]
			}
		}
	}
	if err := scanner.Err(); err != nil {
		common.Error("Fail to read system_hosts: %s", err)
	}
	return hostConfigMap
}

// func initSystemHostsWatcher() *Batcher {
// 	batcher, err := New(time.Millisecond * 300)
// 	if err != nil {
// 		common.Error("Fail to initialize batcher")
// 	}
// 	if err = batcher.Add(systemHosts); err != nil {
// 		common.Error("Fail to add system hosts: %s", systemHosts)
// 	}
// 	return batcher
// }

// func startSystemHostsWatcher() {
// 	if batcher == nil {
// 		common.Error("Fail to start system hosts watcher, watcher is nil")
// 		return
// 	}
// 	for events := range batcher.Events {
// 		for _, event := range events {
// 			if event.Op&fsnotify.Write == fsnotify.Write {
// 				common.Info("modified file: %v", event)
// 				doProcess()
// 				break
// 			}
// 		}
// 	}
// 	// never return
// }
<|endoftext|>