<commit_before>// Package cache is an interface to multiple storage backends for Shade.  It
// centralizes the implementation of reading and writing to multiple
// drive.Clients.
package cache

import (
	"errors"
	"flag"
	"fmt"
	"log"

	"github.com/asjoyner/shade"
	"github.com/asjoyner/shade/drive"
)

var (
	cacheDebug = flag.Bool("cacheDebug", false, "Print cache debugging traces")
)

func init() {
	drive.RegisterProvider("cache", NewClient)
}

type refreshReq struct {
	sha256sum []byte
	content   []byte
	f         *shade.File
}

// NewClient returns a Drive client which centralizes reading and writing to
// multiple Providers.
func NewClient(c drive.Config) (drive.Client, error) {
	if len(c.Children) == 0 {
		return nil, errors.New("no clients provided")
	}
	d := &Drive{config: c}
	for _, conf := range c.Children {
		child, err := drive.NewClient(conf)
		if err != nil {
			return nil, fmt.Errorf("%s: %s", conf.Provider, err)
		}
		if child.GetConfig().Write {
			d.log(fmt.Sprintf("child %s is writable.", conf.Provider))
			d.config.Write = true
		} else {
			d.log(fmt.Sprintf("child %s is NOT writable.", conf.Provider))
		}
		d.clients = append(d.clients, child)
	}
	d.log(fmt.Sprintf("my final status is: %v", d.config.Write))
	d.files = make(chan refreshReq, 100)
	go func(d *Drive) {
		for r := range d.files {
			d.refreshFile(r.sha256sum, r.content)
		}
	}(d)
	d.chunks = make(chan refreshReq, 100)
	go func(d *Drive) {
		for r := range d.chunks {
			d.refreshChunk(r.sha256sum, r.content, r.f)
		}
	}(d)
	return d, nil
}

// Drive implements the drive.Client interface by reading and writing to the
// slice of drive.Client interfaces it was provided.  It can return a config
// which describes only its name.
//
// If any of its clients are not Local(), it reports itself as not Local() by
// returning false.  If any of its clients are Persistent(), it requires writes
// to at least one of those backends to succeed, and reports itself as
// Persistent().
type Drive struct {
	config  drive.Config
	clients []drive.Client
	chunks  chan refreshReq
	files   chan refreshReq
	debug   bool
}

// ListFiles retrieves all of the File objects known to all of the provided
// clients.  The return is a list of sha256sums of the file object.  The keys
// may be passed to GetChunk() to retrieve the corresponding shade.File.
func (s *Drive) ListFiles() ([][]byte, error) {
	c := make(chan [][]byte, len(s.clients))
	for _, client := range s.clients {
		go func(client drive.Client) {
			f, err := client.ListFiles()
			if err != nil {
				s.log(fmt.Sprintf("Error reading from %q: %s", client.GetConfig().Provider, err))
			}
			c <- f
		}(client)
	}

	var resp [][]byte
	for i := 0; i < len(s.clients); i++ {
		resp = append(resp, <-c...)
	}

	return resp, nil
}

// GetFile retrieves a file with a given SHA-256 sum.  It will be returned
// from the first client in the slice of structs that returns the chunk.
func (s *Drive) GetFile(sha256sum []byte) ([]byte, error) {
	for _, client := range s.clients {
		file, err := client.GetFile(sha256sum)
		if err != nil {
			s.log(fmt.Sprintf("File %x not found in %q: %s", sha256sum, client.GetConfig().Provider, err))
			continue
		}
		s.files <- refreshReq{sha256sum: sha256sum, content: file, f: nil}
		return file, nil
	}
	return nil, errors.New("file not found")
}

// PutFile writes the metadata describing a new file.  It will be written to
// all shade backends configured to Write.  If any backends are Persistent, it
// returns an error if all Persistent backends fail to write.
// f should be marshalled JSON, and may be encrypted.
func (s *Drive) PutFile(sha256sum, f []byte) error {
	if s.config.Write == false {
		return errors.New("no clients configured to write")
	}

	persisted := make(chan struct{}, len(s.clients))
	done := make(chan struct{}, len(s.clients))
	for _, client := range s.clients {
		go func(client drive.Client) {
			if err := client.PutFile(sha256sum, f); err != nil {
				s.log(fmt.Sprintf("%s.PutFile(%x) failed: %s", client.GetConfig().Provider, sha256sum, err))
				done <- struct{}{}
				return
			}
			if !s.Persistent() || client.Persistent() {
				persisted <- struct{}{}
				return
			}
			done <- struct{}{}
		}(client)
	}
	for range s.clients {
		select {
		case <-persisted:
			return nil
		case <-done:
		}
	}
	return fmt.Errorf("persistent storage configured, but all writes failed: %x", sha256sum)
}

// GetChunk retrieves a chunk with a given SHA-256 sum.  It will be returned
// from the first client in the slice of structs that returns the chunk.
func (s *Drive) GetChunk(sha256sum []byte, f *shade.File) ([]byte, error) {
	// TODO(asjoyner): consider adding the ability to cancel GetChunk, then
	// paralellize this with a slight delay between launching each request.
	for _, client := range s.clients {
		chunk, err := client.GetChunk(sha256sum, f)
		if err != nil {
			s.log(fmt.Sprintf("Chunk %x not found in %q: %s", sha256sum, client.GetConfig().Provider, err))
			continue
		}
		s.files <- refreshReq{sha256sum: sha256sum, content: chunk, f: f}
		return chunk, nil
	}
	return nil, errors.New("chunk not found")
}

// PutChunk writes a chunk associated with a SHA-256 sum.  It will attempt to write to
// all shade backends configured to Write.  If any backends are Persistent, it
// returns an error if all Persistent backends fail to write.
func (s *Drive) PutChunk(sha256sum []byte, chunk []byte, f *shade.File) error {
	if s.config.Write == false {
		return errors.New("no clients configured to write")
	}

	persisted := make(chan struct{}, len(s.clients))
	done := make(chan struct{}, len(s.clients))
	for _, client := range s.clients {
		go func(client drive.Client) {
			if err := client.PutChunk(sha256sum, chunk, f); err != nil {
				s.log(fmt.Sprintf("%s.PutChunk(%x) failed: %s", client.GetConfig().Provider, sha256sum, err))
				done <- struct{}{}
				return
			}
			if !s.Persistent() || client.Persistent() {
				persisted <- struct{}{}
				return
			}
			done <- struct{}{}
		}(client)
	}
	for range s.clients {
		select {
		case <-persisted:
			return nil
		case <-done:
		}
	}
	return fmt.Errorf("persistent storage configured, but all writes failed: %x", sha256sum)
}

// GetConfig returns the config used to initialize this client.
func (s *Drive) GetConfig() drive.Config {
	return s.config
}

// Local returns true only if all configured storage backends are local to this
// machine.
func (s *Drive) Local() bool {
	for _, c := range s.clients {
		if !c.Local() {
			return false
		}
	}
	return true
}

// Persistent returns true if at least one configured storage backend is
// Persistent().
func (s *Drive) Persistent() bool {
	for _, c := range s.clients {
		if c.Persistent() {
			return true
		}
	}
	return false
}

// Debug enables debug statements to STDERR for non-critical failures to read or
// write from clients.
func (s *Drive) Debug() {
	flag.Set("cacheDebug", "true")
}

func (s *Drive) log(output string) {
	if *cacheDebug {
		log.Printf("drive.Cache: %s\n", output)
	}
}

func (s *Drive) refreshWorker() {
	select {
	case r := <-s.files:
		s.refreshFile(r.sha256sum, r.content)
	case r := <-s.chunks:
		s.refreshChunk(r.sha256sum, r.content, r.f)
	}
}

// refreshFile calls PutFile on each client which is Local()
// This populates eg. memory and disk clients with files that are
// fetched from remote clients.  Errors are logged, but not returned.
func (s *Drive) refreshFile(sha256sum, file []byte) {
	for _, client := range s.clients {
		if client.Local() {
			client.PutFile(sha256sum, file)
		}
	}
}

// refreshChunk calls PutChunk on each client which is Local()
// This populates eg. memory and disk clients with chunks that are
// fetched from remote clients.  Errors are logged, but not returned.
func (s *Drive) refreshChunk(sha256sum, chunk []byte, f *shade.File) {
	for _, client := range s.clients {
		if client.Local() {
			client.PutChunk(sha256sum, chunk, f)
		}
	}
}
<commit_msg>Write chunk refresh requests to the chunks chan :)<commit_after>// Package cache is an interface to multiple storage backends for Shade.  It
// centralizes the implementation of reading and writing to multiple
// drive.Clients.
package cache

import (
	"errors"
	"flag"
	"fmt"
	"log"

	"github.com/asjoyner/shade"
	"github.com/asjoyner/shade/drive"
)

var (
	cacheDebug = flag.Bool("cacheDebug", false, "Print cache debugging traces")
)

func init() {
	drive.RegisterProvider("cache", NewClient)
}

type refreshReq struct {
	sha256sum []byte
	content   []byte
	f         *shade.File
}

// NewClient returns a Drive client which centralizes reading and writing to
// multiple Providers.
func NewClient(c drive.Config) (drive.Client, error) {
	if len(c.Children) == 0 {
		return nil, errors.New("no clients provided")
	}
	d := &Drive{config: c}
	for _, conf := range c.Children {
		child, err := drive.NewClient(conf)
		if err != nil {
			return nil, fmt.Errorf("%s: %s", conf.Provider, err)
		}
		if child.GetConfig().Write {
			d.log(fmt.Sprintf("child %s is writable.", conf.Provider))
			d.config.Write = true
		} else {
			d.log(fmt.Sprintf("child %s is NOT writable.", conf.Provider))
		}
		d.clients = append(d.clients, child)
	}
	d.log(fmt.Sprintf("my final status is: %v", d.config.Write))
	d.files = make(chan refreshReq, 100)
	go func(d *Drive) {
		for r := range d.files {
			d.refreshFile(r.sha256sum, r.content)
		}
	}(d)
	d.chunks = make(chan refreshReq, 100)
	go func(d *Drive) {
		for r := range d.chunks {
			d.refreshChunk(r.sha256sum, r.content, r.f)
		}
	}(d)
	return d, nil
}

// Drive implements the drive.Client interface by reading and writing to the
// slice of drive.Client interfaces it was provided.  It can return a config
// which describes only its name.
//
// If any of its clients are not Local(), it reports itself as not Local() by
// returning false.  If any of its clients are Persistent(), it requires writes
// to at least one of those backends to succeed, and reports itself as
// Persistent().
type Drive struct {
	config  drive.Config
	clients []drive.Client
	chunks  chan refreshReq
	files   chan refreshReq
	debug   bool
}

// ListFiles retrieves all of the File objects known to all of the provided
// clients.  The return is a list of sha256sums of the file object.  The keys
// may be passed to GetChunk() to retrieve the corresponding shade.File.
func (s *Drive) ListFiles() ([][]byte, error) {
	c := make(chan [][]byte, len(s.clients))
	for _, client := range s.clients {
		go func(client drive.Client) {
			f, err := client.ListFiles()
			if err != nil {
				s.log(fmt.Sprintf("Error reading from %q: %s", client.GetConfig().Provider, err))
			}
			c <- f
		}(client)
	}

	var resp [][]byte
	for i := 0; i < len(s.clients); i++ {
		resp = append(resp, <-c...)
	}

	return resp, nil
}

// GetFile retrieves a file with a given SHA-256 sum.  It will be returned
// from the first client in the slice of structs that returns the chunk.
func (s *Drive) GetFile(sha256sum []byte) ([]byte, error) {
	for _, client := range s.clients {
		file, err := client.GetFile(sha256sum)
		if err != nil {
			s.log(fmt.Sprintf("File %x not found in %q: %s", sha256sum, client.GetConfig().Provider, err))
			continue
		}
		s.files <- refreshReq{sha256sum: sha256sum, content: file, f: nil}
		return file, nil
	}
	return nil, errors.New("file not found")
}

// PutFile writes the metadata describing a new file.  It will be written to
// all shade backends configured to Write.  If any backends are Persistent, it
// returns an error if all Persistent backends fail to write.
// f should be marshalled JSON, and may be encrypted.
func (s *Drive) PutFile(sha256sum, f []byte) error {
	if s.config.Write == false {
		return errors.New("no clients configured to write")
	}

	persisted := make(chan struct{}, len(s.clients))
	done := make(chan struct{}, len(s.clients))
	for _, client := range s.clients {
		go func(client drive.Client) {
			if err := client.PutFile(sha256sum, f); err != nil {
				s.log(fmt.Sprintf("%s.PutFile(%x) failed: %s", client.GetConfig().Provider, sha256sum, err))
				done <- struct{}{}
				return
			}
			if !s.Persistent() || client.Persistent() {
				persisted <- struct{}{}
				return
			}
			done <- struct{}{}
		}(client)
	}
	for range s.clients {
		select {
		case <-persisted:
			return nil
		case <-done:
		}
	}
	return fmt.Errorf("persistent storage configured, but all writes failed: %x", sha256sum)
}

// GetChunk retrieves a chunk with a given SHA-256 sum.  It will be returned
// from the first client in the slice of structs that returns the chunk.
func (s *Drive) GetChunk(sha256sum []byte, f *shade.File) ([]byte, error) {
	// TODO(asjoyner): consider adding the ability to cancel GetChunk, then
	// paralellize this with a slight delay between launching each request.
	for _, client := range s.clients {
		chunk, err := client.GetChunk(sha256sum, f)
		if err != nil {
			s.log(fmt.Sprintf("Chunk %x not found in %q: %s", sha256sum, client.GetConfig().Provider, err))
			continue
		}
		s.chunks <- refreshReq{sha256sum: sha256sum, content: chunk, f: f}
		return chunk, nil
	}
	return nil, errors.New("chunk not found")
}

// PutChunk writes a chunk associated with a SHA-256 sum.  It will attempt to write to
// all shade backends configured to Write.  If any backends are Persistent, it
// returns an error if all Persistent backends fail to write.
func (s *Drive) PutChunk(sha256sum []byte, chunk []byte, f *shade.File) error {
	if s.config.Write == false {
		return errors.New("no clients configured to write")
	}

	persisted := make(chan struct{}, len(s.clients))
	done := make(chan struct{}, len(s.clients))
	for _, client := range s.clients {
		go func(client drive.Client) {
			if err := client.PutChunk(sha256sum, chunk, f); err != nil {
				s.log(fmt.Sprintf("%s.PutChunk(%x) failed: %s", client.GetConfig().Provider, sha256sum, err))
				done <- struct{}{}
				return
			}
			if !s.Persistent() || client.Persistent() {
				persisted <- struct{}{}
				return
			}
			done <- struct{}{}
		}(client)
	}
	for range s.clients {
		select {
		case <-persisted:
			return nil
		case <-done:
		}
	}
	return fmt.Errorf("persistent storage configured, but all writes failed: %x", sha256sum)
}

// GetConfig returns the config used to initialize this client.
func (s *Drive) GetConfig() drive.Config {
	return s.config
}

// Local returns true only if all configured storage backends are local to this
// machine.
func (s *Drive) Local() bool {
	for _, c := range s.clients {
		if !c.Local() {
			return false
		}
	}
	return true
}

// Persistent returns true if at least one configured storage backend is
// Persistent().
func (s *Drive) Persistent() bool {
	for _, c := range s.clients {
		if c.Persistent() {
			return true
		}
	}
	return false
}

// Debug enables debug statements to STDERR for non-critical failures to read or
// write from clients.
func (s *Drive) Debug() {
	flag.Set("cacheDebug", "true")
}

func (s *Drive) log(output string) {
	if *cacheDebug {
		log.Printf("drive.Cache: %s\n", output)
	}
}

func (s *Drive) refreshWorker() {
	select {
	case r := <-s.files:
		s.refreshFile(r.sha256sum, r.content)
	case r := <-s.chunks:
		s.refreshChunk(r.sha256sum, r.content, r.f)
	}
}

// refreshFile calls PutFile on each client which is Local()
// This populates eg. memory and disk clients with files that are
// fetched from remote clients.  Errors are logged, but not returned.
func (s *Drive) refreshFile(sha256sum, file []byte) {
	for _, client := range s.clients {
		if client.Local() {
			client.PutFile(sha256sum, file)
		}
	}
}

// refreshChunk calls PutChunk on each client which is Local()
// This populates eg. memory and disk clients with chunks that are
// fetched from remote clients.  Errors are logged, but not returned.
func (s *Drive) refreshChunk(sha256sum, chunk []byte, f *shade.File) {
	for _, client := range s.clients {
		if client.Local() {
			client.PutChunk(sha256sum, chunk, f)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"

	"github.com/ClusterHQ/dvol/pkg/api"
)

const PLUGINS_DIR = "/run/docker/plugins"
const DVOL_SOCKET = PLUGINS_DIR + "/dvol.sock"
const VOL_DIR = "/var/lib/dvol/volumes"

type ResponseImplements struct {
	// A response to the Plugin.Activate request
	Implements []string
}

type RequestCreate struct {
	// A request to create a volume for Docker
	Name string
	Opts map[string]string
}

type RequestMount struct {
	// A request to mount a volume for Docker
	Name string
}

type RequestRemove struct {
	// A request to remove a volume for Docker
	Name string
}

type ResponseSimple struct {
	// A response which only indicates if there was an error or not
	Err string
}

type ResponseMount struct {
	// A response to the VolumeDriver.Mount request
	Mountpoint string
	Err        string
}

func main() {
	if _, err := os.Stat(PLUGINS_DIR); err != nil {
		if err := os.MkdirAll(PLUGINS_DIR, 0700); err != nil {
			log.Fatalf("Could not make plugin directory %s: %v", PLUGINS_DIR, err)
		}
	}
	if _, err := os.Stat(DVOL_SOCKET); err == nil {
		if err = os.Remove(DVOL_SOCKET); err != nil {
			log.Fatalf("Could not clean up existing socket at %s: %v", DVOL_SOCKET, err)
		}
	}
	if _, err := os.Stat(VOL_DIR); err != nil {
		if err := os.MkdirAll(VOL_DIR, 0700); err != nil {
			log.Fatalf("Could not make volumes directory %s: %v", VOL_DIR, err)
		}
	}

	listener, err := net.Listen("unix", DVOL_SOCKET)

	if err != nil {
		log.Fatalf("Could not listen on %s: %v", DVOL_SOCKET, err)
	}

	http.HandleFunc("/Plugin.Activate", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /Plugin.Activate")
		responseJSON, _ := json.Marshal(&ResponseImplements{
			Implements: []string{"VolumeDriver"},
		})
		w.Write(responseJSON)
	})

	http.HandleFunc("/VolumeDriver.Create", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /VolumeDriver.Create")
		requestJSON, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Fatalf("Unable to read response body %s", err)
		}
		request := new(RequestCreate)
		json.Unmarshal(requestJSON, request)
		name := request.Name
		dvol := api.NewDvolAPI(VOL_DIR)
		if !dvol.VolumeExists(name) {
			log.Print("Creating volume", name, " which doesn't exist")
			err := dvol.CreateVolume(name)
			if err != nil {
				WriteResponseErr(err, w)
				return
			}
		}
		WriteResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.Remove", func(w http.ResponseWriter, r *http.Request) {
		WriteResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.Path", func(w http.ResponseWriter, r *http.Request) {
	})

	http.HandleFunc("/VolumeDriver.Mount", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /VolumeDriver.Mount")
		requestJSON, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Fatalf("Unable to read response body %s", err)
		}
		request := new(RequestMount)
		json.Unmarshal(requestJSON, request)
		name := request.Name

		dvol := api.NewDvolAPI(VOL_DIR)

		if dvol.VolumeExists(name) {
			err := dvol.SwitchVolume(name)
			if err != nil {
				WriteResponseErr(err, w)
			}
			_, err = dvol.ActiveVolume()
			if err != nil {
				WriteResponseErr(err, w)
			}
			// mountpoint should be:
			// /var/lib/docker/volumes/<volumename>/running_point
			responseJSON, _ := json.Marshal(&ResponseMount{
				Mountpoint: "/tmp", // TODO: Get the real path
				Err:        "",
			})
			w.Write(responseJSON)
		} else {
			WriteResponseErr(errors.New("Requested to mount unknown volume "+name), w)
		}
	})

	http.HandleFunc("/VolumeDriver.Unmount", func(w http.ResponseWriter, r *http.Request) {
		WriteResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.List", func(w http.ResponseWriter, r *http.Request) {
	})

	http.Serve(listener, nil)
}

func WriteResponseOK(w http.ResponseWriter) {
	// A shortcut to writing a ResponseOK to w
	responseJSON, _ := json.Marshal(&ResponseSimple{Err: ""})
	w.Write(responseJSON)
}

func WriteResponseErr(err error, w http.ResponseWriter) {
	// A shortcut to responding with an error, and then log the error
	errString := fmt.Sprintln(err)
	log.Printf("Error: %v", err)
	responseJSON, _ := json.Marshal(&ResponseSimple{Err: errString})
	w.Write(responseJSON)
}
<commit_msg>Address review comment: move this closer to where it gets used.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"

	"github.com/ClusterHQ/dvol/pkg/api"
)

const PLUGINS_DIR = "/run/docker/plugins"
const DVOL_SOCKET = PLUGINS_DIR + "/dvol.sock"
const VOL_DIR = "/var/lib/dvol/volumes"

type ResponseImplements struct {
	// A response to the Plugin.Activate request
	Implements []string
}

type RequestCreate struct {
	// A request to create a volume for Docker
	Name string
	Opts map[string]string
}

type RequestMount struct {
	// A request to mount a volume for Docker
	Name string
}

type RequestRemove struct {
	// A request to remove a volume for Docker
	Name string
}

type ResponseSimple struct {
	// A response which only indicates if there was an error or not
	Err string
}

type ResponseMount struct {
	// A response to the VolumeDriver.Mount request
	Mountpoint string
	Err        string
}

func main() {
	if _, err := os.Stat(PLUGINS_DIR); err != nil {
		if err := os.MkdirAll(PLUGINS_DIR, 0700); err != nil {
			log.Fatalf("Could not make plugin directory %s: %v", PLUGINS_DIR, err)
		}
	}
	if _, err := os.Stat(DVOL_SOCKET); err == nil {
		if err = os.Remove(DVOL_SOCKET); err != nil {
			log.Fatalf("Could not clean up existing socket at %s: %v", DVOL_SOCKET, err)
		}
	}
	if _, err := os.Stat(VOL_DIR); err != nil {
		if err := os.MkdirAll(VOL_DIR, 0700); err != nil {
			log.Fatalf("Could not make volumes directory %s: %v", VOL_DIR, err)
		}
	}

	if err != nil {
		log.Fatalf("Could not listen on %s: %v", DVOL_SOCKET, err)
	}

	http.HandleFunc("/Plugin.Activate", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /Plugin.Activate")
		responseJSON, _ := json.Marshal(&ResponseImplements{
			Implements: []string{"VolumeDriver"},
		})
		w.Write(responseJSON)
	})

	http.HandleFunc("/VolumeDriver.Create", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /VolumeDriver.Create")
		requestJSON, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Fatalf("Unable to read response body %s", err)
		}
		request := new(RequestCreate)
		json.Unmarshal(requestJSON, request)
		name := request.Name
		dvol := api.NewDvolAPI(VOL_DIR)
		if !dvol.VolumeExists(name) {
			log.Print("Creating volume", name, " which doesn't exist")
			err := dvol.CreateVolume(name)
			if err != nil {
				WriteResponseErr(err, w)
				return
			}
		}
		WriteResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.Remove", func(w http.ResponseWriter, r *http.Request) {
		WriteResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.Path", func(w http.ResponseWriter, r *http.Request) {
	})

	http.HandleFunc("/VolumeDriver.Mount", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /VolumeDriver.Mount")
		requestJSON, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Fatalf("Unable to read response body %s", err)
		}
		request := new(RequestMount)
		json.Unmarshal(requestJSON, request)
		name := request.Name

		dvol := api.NewDvolAPI(VOL_DIR)

		if dvol.VolumeExists(name) {
			err := dvol.SwitchVolume(name)
			if err != nil {
				WriteResponseErr(err, w)
			}
			_, err = dvol.ActiveVolume()
			if err != nil {
				WriteResponseErr(err, w)
			}
			// mountpoint should be:
			// /var/lib/docker/volumes/<volumename>/running_point
			responseJSON, _ := json.Marshal(&ResponseMount{
				Mountpoint: "/tmp", // TODO: Get the real path
				Err:        "",
			})
			w.Write(responseJSON)
		} else {
			WriteResponseErr(errors.New("Requested to mount unknown volume "+name), w)
		}
	})

	http.HandleFunc("/VolumeDriver.Unmount", func(w http.ResponseWriter, r *http.Request) {
		WriteResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.List", func(w http.ResponseWriter, r *http.Request) {
	})

	listener, err := net.Listen("unix", DVOL_SOCKET)
	http.Serve(listener, nil)
}

func WriteResponseOK(w http.ResponseWriter) {
	// A shortcut to writing a ResponseOK to w
	responseJSON, _ := json.Marshal(&ResponseSimple{Err: ""})
	w.Write(responseJSON)
}

func WriteResponseErr(err error, w http.ResponseWriter) {
	// A shortcut to responding with an error, and then log the error
	errString := fmt.Sprintln(err)
	log.Printf("Error: %v", err)
	responseJSON, _ := json.Marshal(&ResponseSimple{Err: errString})
	w.Write(responseJSON)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"

	"github.com/ClusterHQ/dvol/pkg/api"
)

const PLUGINS_DIR = "/run/docker/plugins"
const DVOL_SOCKET = PLUGINS_DIR + "/dvol.sock"
const VOL_DIR = "/var/lib/dvol/volumes"

type DockerVolumePluginType string
type VolumeName string
type OptKey string
type OptValue string
type ErrResponse string

type ResponseImplements struct {
	// A response to the Plugin.Activate request
	Implements []DockerVolumePluginType
}

type RequestCreate struct {
	// A request to create a volume for Docker
	Name VolumeName
	Opts map[OptKey]OptValue
}

type RequestMount struct {
	// A request to mount a volume for Docker
	Name VolumeName
}

type RequestRemove struct {
	// A request to remove a volume for Docker
	Name VolumeName
}

type ResponseSimple struct {
	// A response which only indicates if there was an error or not
	Err ErrResponse
}

type ResponseMount struct {
	// A response to the VolumeDriver.Mount request
	Mountpoint MountPath
	Err        ErrResponse
}

func main() {
	log.Print("Starting dvol plugin")

	if _, err := os.Stat(PLUGINS_DIR); err != nil {
		if err := os.MkdirAll(PLUGINS_DIR, 0700); err != nil {
			log.Fatalf("Could not make plugin directory %s: %v", PLUGINS_DIR, err)
		}
	}
	if _, err := os.Stat(DVOL_SOCKET); err == nil {
		if err = os.Remove(DVOL_SOCKET); err != nil {
			log.Fatalf("Could not clean up existing socket at %s: %v", DVOL_SOCKET, err)
		}
	}
	if _, err := os.Stat(VOL_DIR); err != nil {
		if err := os.MkdirAll(VOL_DIR, 0700); err != nil {
			log.Fatalf("Could not make volumes directory %s: %v", VOL_DIR, err)
		}
	}

	http.HandleFunc("/Plugin.Activate", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /Plugin.Activate")
		responseJSON, _ := json.Marshal(&ResponseImplements{
			Implements: []string{"VolumeDriver"},
		})
		w.Write(responseJSON)
	})

	http.HandleFunc("/VolumeDriver.Create", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /VolumeDriver.Create")
		requestJSON, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Fatalf("Unable to read response body %s", err)
		}
		request := new(RequestCreate)
		json.Unmarshal(requestJSON, request)
		name := request.Name
		dvol := api.NewDvolAPI(VOL_DIR)
		if dvol.VolumeExists(name) {
			log.Print("Volume already exists: %s", name)
		} else {
			err := dvol.CreateVolume(name)
			if err != nil {
				writeResponseErr(fmt.Errorf("Could not create volume %s: %v", name, err), w)
				return
			}
		}
		writeResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.Remove", func(w http.ResponseWriter, r *http.Request) {
		writeResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.Path", func(w http.ResponseWriter, r *http.Request) {
	})

	http.HandleFunc("/VolumeDriver.Mount", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /VolumeDriver.Mount")
		requestJSON, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Fatalf("Unable to read response body %s", err)
		}
		request := new(RequestMount)
		json.Unmarshal(requestJSON, request)
		name := request.Name

		dvol := api.NewDvolAPI(VOL_DIR)

		if dvol.VolumeExists(name) {
			err := dvol.SwitchVolume(name)
			if err != nil {
				writeResponseErr(err, w)
			}
			_, err = dvol.ActiveVolume()
			if err != nil {
				writeResponseErr(err, w)
			}
			// mountpoint should be:
			// /var/lib/docker/volumes/<volumename>/running_point
			responseJSON, _ := json.Marshal(&ResponseMount{
				Mountpoint: "/tmp", // TODO: Get the real path
				Err:        "",
			})
			w.Write(responseJSON)
		} else {
			writeResponseErr(fmt.Errorf("Requested to mount unknown volume %s", name), w)
		}
	})

	http.HandleFunc("/VolumeDriver.Unmount", func(w http.ResponseWriter, r *http.Request) {
		writeResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.List", func(w http.ResponseWriter, r *http.Request) {
	})

	listener, err := net.Listen("unix", DVOL_SOCKET)
	if err != nil {
		log.Fatalf("Could not listen on %s: %v", DVOL_SOCKET, err)
	}

	http.Serve(listener, nil)
}

func writeResponseOK(w http.ResponseWriter) {
	// A shortcut to writing a ResponseOK to w
	responseJSON, _ := json.Marshal(&ResponseSimple{Err: ""})
	w.Write(responseJSON)
}

func writeResponseErr(err error, w http.ResponseWriter) {
	// A shortcut to responding with an error, and then log the error
	errString := fmt.Sprintln(err)
	log.Printf("Error: %v", err)
	responseJSON, _ := json.Marshal(&ResponseSimple{Err: errString})
	w.Write(responseJSON)
}
<commit_msg>Undo typing all the strings for now.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"

	"github.com/ClusterHQ/dvol/pkg/api"
)

const PLUGINS_DIR = "/run/docker/plugins"
const DVOL_SOCKET = PLUGINS_DIR + "/dvol.sock"
const VOL_DIR = "/var/lib/dvol/volumes"

type ResponseImplements struct {
	// A response to the Plugin.Activate request
	Implements []string
}

type RequestCreate struct {
	// A request to create a volume for Docker
	Name string
	Opts map[string]string
}

type RequestMount struct {
	// A request to mount a volume for Docker
	Name string
}

type RequestRemove struct {
	// A request to remove a volume for Docker
	Name string
}

type ResponseSimple struct {
	// A response which only indicates if there was an error or not
	Err string
}

type ResponseMount struct {
	// A response to the VolumeDriver.Mount request
	Mountpoint string
	Err        string
}

func main() {
	log.Print("Starting dvol plugin")

	if _, err := os.Stat(PLUGINS_DIR); err != nil {
		if err := os.MkdirAll(PLUGINS_DIR, 0700); err != nil {
			log.Fatalf("Could not make plugin directory %s: %v", PLUGINS_DIR, err)
		}
	}
	if _, err := os.Stat(DVOL_SOCKET); err == nil {
		if err = os.Remove(DVOL_SOCKET); err != nil {
			log.Fatalf("Could not clean up existing socket at %s: %v", DVOL_SOCKET, err)
		}
	}
	if _, err := os.Stat(VOL_DIR); err != nil {
		if err := os.MkdirAll(VOL_DIR, 0700); err != nil {
			log.Fatalf("Could not make volumes directory %s: %v", VOL_DIR, err)
		}
	}

	http.HandleFunc("/Plugin.Activate", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /Plugin.Activate")
		responseJSON, _ := json.Marshal(&ResponseImplements{
			Implements: []string{"VolumeDriver"},
		})
		w.Write(responseJSON)
	})

	http.HandleFunc("/VolumeDriver.Create", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /VolumeDriver.Create")
		requestJSON, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Fatalf("Unable to read response body %s", err)
		}
		request := new(RequestCreate)
		json.Unmarshal(requestJSON, request)
		name := request.Name
		dvol := api.NewDvolAPI(VOL_DIR)
		if dvol.VolumeExists(name) {
			log.Print("Volume already exists: %s", name)
		} else {
			err := dvol.CreateVolume(name)
			if err != nil {
				writeResponseErr(fmt.Errorf("Could not create volume %s: %v", name, err), w)
				return
			}
		}
		writeResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.Remove", func(w http.ResponseWriter, r *http.Request) {
		writeResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.Path", func(w http.ResponseWriter, r *http.Request) {
	})

	http.HandleFunc("/VolumeDriver.Mount", func(w http.ResponseWriter, r *http.Request) {
		log.Print("<= /VolumeDriver.Mount")
		requestJSON, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Fatalf("Unable to read response body %s", err)
		}
		request := new(RequestMount)
		json.Unmarshal(requestJSON, request)
		name := request.Name

		dvol := api.NewDvolAPI(VOL_DIR)

		if dvol.VolumeExists(name) {
			err := dvol.SwitchVolume(name)
			if err != nil {
				writeResponseErr(err, w)
			}
			_, err = dvol.ActiveVolume()
			if err != nil {
				writeResponseErr(err, w)
			}
			// mountpoint should be:
			// /var/lib/docker/volumes/<volumename>/running_point
			responseJSON, _ := json.Marshal(&ResponseMount{
				Mountpoint: "/tmp", // TODO: Get the real path
				Err:        "",
			})
			w.Write(responseJSON)
		} else {
			writeResponseErr(fmt.Errorf("Requested to mount unknown volume %s", name), w)
		}
	})

	http.HandleFunc("/VolumeDriver.Unmount", func(w http.ResponseWriter, r *http.Request) {
		writeResponseOK(w)
	})

	http.HandleFunc("/VolumeDriver.List", func(w http.ResponseWriter, r *http.Request) {
	})

	listener, err := net.Listen("unix", DVOL_SOCKET)
	if err != nil {
		log.Fatalf("Could not listen on %s: %v", DVOL_SOCKET, err)
	}

	http.Serve(listener, nil)
}

func writeResponseOK(w http.ResponseWriter) {
	// A shortcut to writing a ResponseOK to w
	responseJSON, _ := json.Marshal(&ResponseSimple{Err: ""})
	w.Write(responseJSON)
}

func writeResponseErr(err error, w http.ResponseWriter) {
	// A shortcut to responding with an error, and then log the error
	errString := fmt.Sprintln(err)
	log.Printf("Error: %v", err)
	responseJSON, _ := json.Marshal(&ResponseSimple{Err: errString})
	w.Write(responseJSON)
}
<|endoftext|>


<commit_before>package gofast

import (
	"fmt"
	"net"
	"net/http"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/tools/godoc/vfs"
	"golang.org/x/tools/godoc/vfs/httpfs"
)

// SessionHandler handles the gofast *Reqeust with the provided given Client.
// The Client should properly handle the transport to the fastcgi application.
// Should do proper routing or other parameter mapping here.
type SessionHandler func(client Client, req *Request) (resp *ResponsePipe, err error)

// Middleware transform a SessionHandler as another SessionHandler. The
// middlewares provided by this library helps to map fastcgi parameters
// according to the need of different application.
//
// You may also implement your own Middleware can be provided to modify
// the *Request, add extra business logic in between, rewrite the response
// stream from *ResponsePipe. or better handle errors
//
// Ordinary fastcgi parameters on nginx (for PHP at least):
//
//  fastcgi_split_path_info ^(.+\.php)(/?.+)$;
//  fastcgi_param  SCRIPT_FILENAME    $document_root$fastcgi_script_name;
//  fastcgi_param  PATH_INFO          $fastcgi_path_info;
//  fastcgi_param  PATH_TRANSLATED    $document_root$fastcgi_path_info;
//  fastcgi_param  QUERY_STRING       $query_string;
//  fastcgi_param  REQUEST_METHOD     $request_method;
//  fastcgi_param  CONTENT_TYPE       $content_type;
//  fastcgi_param  CONTENT_LENGTH     $content_length;
//  fastcgi_param  SCRIPT_NAME        $fastcgi_script_name;
//  fastcgi_param  REQUEST_URI        $request_uri;
//  fastcgi_param  DOCUMENT_URI       $document_uri;
//  fastcgi_param  DOCUMENT_ROOT      $document_root;
//  fastcgi_param  SERVER_PROTOCOL    $server_protocol;
//  fastcgi_param  HTTPS              $https if_not_empty;
//  fastcgi_param  GATEWAY_INTERFACE  CGI/1.1;
//  fastcgi_param  SERVER_SOFTWARE    nginx/$nginx_version;
//  fastcgi_param  REMOTE_ADDR        $remote_addr;
//  fastcgi_param  REMOTE_PORT        $remote_port;
//  fastcgi_param  SERVER_ADDR        $server_addr;
//  fastcgi_param  SERVER_PORT        $server_port;
//  fastcgi_param  SERVER_NAME        $server_name;
//  # PHP only, required if PHP was built with --enable-force-cgi-redirect
//  fastcgi_param  REDIRECT_STATUS    200;
//
type Middleware func(SessionHandler) SessionHandler

// Chain chains middlewares into a single middleware
func Chain(middlewares ...Middleware) Middleware {
	if len(middlewares) == 0 {
		return nil
	}
	return func(inner SessionHandler) (out SessionHandler) {
		out = inner
		for i := len(middlewares) - 1; i >= 0; i-- {
			out = middlewares[i](out)
		}
		return
	}
}

// BasicSession is the default SessionHandler used in the default Handler
func BasicSession(client Client, req *Request) (*ResponsePipe, error) {
	return client.Do(req)
}

// BasicParamsMap implements Middleware. It maps basic parameters to the
// req.Params.
//
// Parameters included:
//  CONTENT_TYPE
//  CONTENT_LENGTH
//  HTTPS
//  GATEWAY_INTERFACE
//  REMOTE_ADDR
//  REMOTE_PORT
//  SERVER_PORT
//  SERVER_NAME
//  SERVER_PROTOCOL
//  SERVER_SOFTWARE
//  REDIRECT_STATUS
//  REQUEST_METHOD
//  REQUEST_URI
//  QUERY_STRING
//
func BasicParamsMap(inner SessionHandler) SessionHandler {
	return func(client Client, req *Request) (*ResponsePipe, error) {

		r := req.Raw

		isHTTPS := r.TLS != nil
		if isHTTPS {
			req.Params["HTTPS"] = "on"
		}

		remoteAddr, remotePort, _ := net.SplitHostPort(r.RemoteAddr)
		host, serverPort, err := net.SplitHostPort(r.Host)
		if err != nil {
			if isHTTPS {
				serverPort = "443"
			} else {
				serverPort = "80"
			}
		}

		// the basic information here
		req.Params["CONTENT_TYPE"] = r.Header.Get("Content-Type")
		req.Params["CONTENT_LENGTH"] = r.Header.Get("Content-Length")
		req.Params["GATEWAY_INTERFACE"] = "CGI/1.1"
		req.Params["REMOTE_ADDR"] = remoteAddr
		req.Params["REMOTE_PORT"] = remotePort
		req.Params["SERVER_PORT"] = serverPort
		req.Params["SERVER_NAME"] = r.Host
		req.Params["SERVER_PROTOCOL"] = r.Proto
		req.Params["SERVER_SOFTWARE"] = "gofast"
		req.Params["REDIRECT_STATUS"] = "200"
		req.Params["REQUEST_METHOD"] = r.Method
		req.Params["REQUEST_URI"] = r.RequestURI
		req.Params["QUERY_STRING"] = r.URL.RawQuery

		return inner(client, req)
	}
}

// FilterAuthReqParams filter out FCGI_PARAMS key-value that is explicitly
// forbidden to passed on in factcgi specification, include:
//  CONTENT_LENGTH;
//  PATH_INFO;
//  PATH_TRANSLATED; and
//  SCRIPT_NAME
func FilterAuthReqParams(inner SessionHandler) SessionHandler {
	return func(client Client, req *Request) (*ResponsePipe, error) {
		if _, ok := req.Params["CONTENT_LENGTH"]; ok {
			delete(req.Params, "CONTENT_LENGTH")
		}
		if _, ok := req.Params["PATH_INFO"]; ok {
			delete(req.Params, "PATH_INFO")
		}
		if _, ok := req.Params["PATH_TRANSLATED"]; ok {
			delete(req.Params, "PATH_TRANSLATED")
		}
		if _, ok := req.Params["SCRIPT_NAME"]; ok {
			delete(req.Params, "SCRIPT_NAME")
		}

		return inner(client, req)
	}
}

// FileSystemRouter helps to produce Middleware implementation for
// mapping path related fastcgi parameters. See method Router for usage.
type FileSystemRouter struct {

	// DocRoot stores the ordinary Apache DocumentRoot parameter
	DocRoot string

	// Exts stores accepted extensions
	Exts []string

	// DirIndex stores ordinary Apache DirectoryIndex parameter
	// for to identify file to show in directory
	DirIndex []string
}

// Router returns a Middleware that prepare session parameters that are
// path related. With information provided in the FileSystemRouter, it will
// route request to script files which path matches the http request path.
//
// i.e. classic PHP hosting environment like Apache + mod_php
//
// Parameters included:
//  PATH_INFO
//  PATH_TRANSLATED
//  SCRIPT_NAME
//  SCRIPT_FILENAME
//  DOCUMENT_URI
//  DOCUMENT_ROOT
//
func (fs *FileSystemRouter) Router() Middleware {
	return func(inner SessionHandler) SessionHandler {
		return func(client Client, req *Request) (*ResponsePipe, error) {

			// define some required cgi parameters
			// with the given http request
			r := req.Raw
			fastcgiScriptName := r.URL.Path

			var fastcgiPathInfo string
			pathinfoRe := regexp.MustCompile(`^(.+\.php)(/?.+)$`)
			if matches := pathinfoRe.FindStringSubmatch(fastcgiScriptName); len(matches) > 0 {
				fastcgiScriptName, fastcgiPathInfo = matches[1], matches[2]
			}

			req.Params["PATH_INFO"] = fastcgiPathInfo
			req.Params["PATH_TRANSLATED"] = filepath.Join(fs.DocRoot, fastcgiPathInfo)
			req.Params["SCRIPT_NAME"] = fastcgiScriptName
			req.Params["SCRIPT_FILENAME"] = filepath.Join(fs.DocRoot, fastcgiScriptName)
			req.Params["DOCUMENT_URI"] = r.URL.Path
			req.Params["DOCUMENT_ROOT"] = fs.DocRoot

			// handle directory index
			urlPath := r.URL.Path
			if strings.HasSuffix(urlPath, "/") {
				urlPath = path.Join(urlPath, "index.php")
			}
			req.Params["SCRIPT_FILENAME"] = path.Join(fs.DocRoot, urlPath)

			return inner(client, req)
		}
	}
}

// MapHeader implement Middleware to map header field HTTP_*
//
// It is a convention to map header field SomeRandomField to
// HTTP_SOME_RANDOM_FIELD. For example, if a header field "X-Hello-World" is in
// the header, it will be mapped as "X_HELLO_WORLD" in the fastcgi parameter
// field.
//
// Note: HTTP_CONTENT_TYPE and HTTP_CONTENT_LENGTH cannot be overridden.
//
func MapHeader(inner SessionHandler) SessionHandler {
	return func(client Client, req *Request) (*ResponsePipe, error) {
		r := req.Raw

		// http header
		for k, v := range r.Header {
			formattedKey := strings.Replace(strings.ToUpper(k), "-", "_", -1)
			if formattedKey == "CONTENT_TYPE" || formattedKey == "CONTENT_LENGTH" {
				continue
			}

			key := "HTTP_" + formattedKey
			var value string
			if len(v) > 0 {
				//   refer to https://tools.ietf.org/html/rfc7230#section-3.2.2
				//
				//   A recipient MAY combine multiple header fields with the same field
				//   name into one "field-name: field-value" pair, without changing the
				//   semantics of the message, by appending each subsequent field value to
				//   the combined field value in order, separated by a comma.  The order
				//   in which header fields with the same field name are received is
				//   therefore significant to the interpretation of the combined field
				//   value; a proxy MUST NOT change the order of these field values when
				//   forwarding a message.
				value = strings.Join(v, ",")
			}
			req.Params[key] = value
		}

		return inner(client, req)
	}
}

// MapEndpoint returns a Middleware implementation that prepare session for
// application with only 1 file as endpoint (i.e. it will handle script routing
// on its own). Suitable for web.py based application.
//
// Parameters included:
//  PATH_INFO
//  PATH_TRANSLATED
//  SCRIPT_NAME
//  SCRIPT_FILENAME
//  DOCUMENT_URI
//  DOCUMENT_ROOT
//
func MapEndpoint(endpointFile string) Middleware {
	dir, webpath := filepath.Dir(endpointFile), "/"+filepath.Base(endpointFile)
	return func(inner SessionHandler) SessionHandler {
		return func(client Client, req *Request) (*ResponsePipe, error) {
			r := req.Raw
			req.Params["REQUEST_URI"] = webpath + r.URL.RequestURI()
			req.Params["SCRIPT_NAME"] = webpath
			req.Params["SCRIPT_FILENAME"] = endpointFile
			req.Params["DOCUMENT_URI"] = r.URL.Path
			req.Params["DOCUMENT_ROOT"] = dir
			return inner(client, req)
		}
	}
}

// MapFilterRequest changes the request role to RoleFilter and add the
// Data stream from the given file system, if file exists. Also
// set the required params to request.
//
// If the file do not exists or cannot be opened, the middleware
// will return empty response pipe and the error.
func MapFilterRequest(fs http.FileSystem) Middleware {
	return func(inner SessionHandler) SessionHandler {
		return func(client Client, req *Request) (*ResponsePipe, error) {

			// force role to be RoleFilter
			req.Role = RoleFilter

			// define some required cgi parameters
			// with the given http request
			r := req.Raw
			fastcgiScriptName := r.URL.Path

			var fastcgiPathInfo string
			pathinfoRe := regexp.MustCompile(`^(.+\.php)(/?.+)$`)
			if matches := pathinfoRe.FindStringSubmatch(fastcgiScriptName); len(matches) > 0 {
				fastcgiScriptName, fastcgiPathInfo = matches[1], matches[2]
			}

			req.Params["PATH_INFO"] = fastcgiPathInfo
			req.Params["SCRIPT_NAME"] = fastcgiScriptName
			req.Params["DOCUMENT_URI"] = r.URL.Path

			// handle directory index
			urlPath := r.URL.Path
			if strings.HasSuffix(urlPath, "/") {
				urlPath = path.Join(urlPath, "index.php")
			}

			// find the file
			f, err := fs.Open(urlPath)
			if err != nil {
				err = fmt.Errorf("cannot open file: %s", err)
				return nil, err
			}

			// map fcgi params for filtering
			s, err := f.Stat()
			if err != nil {
				err = fmt.Errorf("cannot stat file: %s", err)
				return nil, err
			}
			req.Params["FCGI_DATA_LAST_MOD"] = fmt.Sprintf("%d", s.ModTime().Unix())
			req.Params["FCGI_DATA_LENGTH"] = fmt.Sprintf("%d", s.Size())

			// use the file as FCGI_DATA in request
			req.Data = f
			return inner(client, req)
		}
	}
}

// NewFilterLocalFS is a shortcut to use NewFilterFS with
// a http.FileSystem created for the given local folder.
func NewFilterLocalFS(root string) Middleware {
	fs := httpfs.New(vfs.OS(root))
	return NewFilterFS(fs)
}

// NewFilterFS chains BasicParamsMap, MapHeader and MapFilterRequest
// to implement Middleware that prepares a fastcgi Filter session
// environment.
func NewFilterFS(fs http.FileSystem) Middleware {
	return Chain(
		BasicParamsMap,
		MapHeader,
		MapFilterRequest(fs),
	)
}

// NewPHPFS chains BasicParamsMap, MapHeader and FileSystemRouter to implement
// Middleware that prepares an ordinary PHP hosting session environment.
func NewPHPFS(root string) Middleware {
	fs := &FileSystemRouter{
		DocRoot:  root,
		Exts:     []string{"php"},
		DirIndex: []string{"index.php"},
	}
	return Chain(
		BasicParamsMap,
		MapHeader,
		fs.Router(),
	)
}

// NewFileEndpoint chains BasicParamsMap, MapHeader and MapEndpoint to implement
// Middleware that prepares an ordinary web.py hosting session environment
func NewFileEndpoint(endpointFile string) Middleware {
	return Chain(
		BasicParamsMap,
		MapHeader,
		MapEndpoint(endpointFile),
	)
}

// NewAuthPrepare chains BasicParamsMap, MapHeader, FilterAuthReqParams to
// implement Middleware that prepares a authorizer request
func NewAuthPrepare() Middleware {
	return Chain(
		BasicParamsMap,
		MapHeader,
		FilterAuthReqParams,
	)
}
<commit_msg>bugfix: can not get schema from http.Request.URL <commit_after>package gofast

import (
	"fmt"
	"net"
	"net/http"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/tools/godoc/vfs"
	"golang.org/x/tools/godoc/vfs/httpfs"
)

// SessionHandler handles the gofast *Reqeust with the provided given Client.
// The Client should properly handle the transport to the fastcgi application.
// Should do proper routing or other parameter mapping here.
type SessionHandler func(client Client, req *Request) (resp *ResponsePipe, err error)

// Middleware transform a SessionHandler as another SessionHandler. The
// middlewares provided by this library helps to map fastcgi parameters
// according to the need of different application.
//
// You may also implement your own Middleware can be provided to modify
// the *Request, add extra business logic in between, rewrite the response
// stream from *ResponsePipe. or better handle errors
//
// Ordinary fastcgi parameters on nginx (for PHP at least):
//
//  fastcgi_split_path_info ^(.+\.php)(/?.+)$;
//  fastcgi_param  SCRIPT_FILENAME    $document_root$fastcgi_script_name;
//  fastcgi_param  PATH_INFO          $fastcgi_path_info;
//  fastcgi_param  PATH_TRANSLATED    $document_root$fastcgi_path_info;
//  fastcgi_param  QUERY_STRING       $query_string;
//  fastcgi_param  REQUEST_METHOD     $request_method;
//  fastcgi_param  CONTENT_TYPE       $content_type;
//  fastcgi_param  CONTENT_LENGTH     $content_length;
//  fastcgi_param  SCRIPT_NAME        $fastcgi_script_name;
//  fastcgi_param  REQUEST_URI        $request_uri;
//  fastcgi_param  DOCUMENT_URI       $document_uri;
//  fastcgi_param  DOCUMENT_ROOT      $document_root;
//  fastcgi_param  SERVER_PROTOCOL    $server_protocol;
//  fastcgi_param  HTTPS              $https if_not_empty;
//  fastcgi_param  GATEWAY_INTERFACE  CGI/1.1;
//  fastcgi_param  SERVER_SOFTWARE    nginx/$nginx_version;
//  fastcgi_param  REMOTE_ADDR        $remote_addr;
//  fastcgi_param  REMOTE_PORT        $remote_port;
//  fastcgi_param  SERVER_ADDR        $server_addr;
//  fastcgi_param  SERVER_PORT        $server_port;
//  fastcgi_param  SERVER_NAME        $server_name;
//  # PHP only, required if PHP was built with --enable-force-cgi-redirect
//  fastcgi_param  REDIRECT_STATUS    200;
//
type Middleware func(SessionHandler) SessionHandler

// Chain chains middlewares into a single middleware
func Chain(middlewares ...Middleware) Middleware {
	if len(middlewares) == 0 {
		return nil
	}
	return func(inner SessionHandler) (out SessionHandler) {
		out = inner
		for i := len(middlewares) - 1; i >= 0; i-- {
			out = middlewares[i](out)
		}
		return
	}
}

// BasicSession is the default SessionHandler used in the default Handler
func BasicSession(client Client, req *Request) (*ResponsePipe, error) {
	return client.Do(req)
}

// BasicParamsMap implements Middleware. It maps basic parameters to the
// req.Params.
//
// Parameters included:
//  CONTENT_TYPE
//  CONTENT_LENGTH
//  HTTPS
//  GATEWAY_INTERFACE
//  REMOTE_ADDR
//  REMOTE_PORT
//  SERVER_PORT
//  SERVER_NAME
//  SERVER_PROTOCOL
//  SERVER_SOFTWARE
//  REDIRECT_STATUS
//  REQUEST_METHOD
//  REQUEST_URI
//  QUERY_STRING
//
func BasicParamsMap(inner SessionHandler) SessionHandler {
	return func(client Client, req *Request) (*ResponsePipe, error) {

		r := req.Raw

		isHTTPS := r.TLS != nil
		if isHTTPS {
			req.Params["HTTPS"] = "on"
		}

		remoteAddr, remotePort, _ := net.SplitHostPort(r.RemoteAddr)
		host, serverPort, err := net.SplitHostPort(r.Host)
		if err != nil {
			if isHTTPS {
				serverPort = "443"
			} else {
				serverPort = "80"
			}
		}

		// the basic information here
		req.Params["CONTENT_TYPE"] = r.Header.Get("Content-Type")
		req.Params["CONTENT_LENGTH"] = r.Header.Get("Content-Length")
		req.Params["GATEWAY_INTERFACE"] = "CGI/1.1"
		req.Params["REMOTE_ADDR"] = remoteAddr
		req.Params["REMOTE_PORT"] = remotePort
		req.Params["SERVER_PORT"] = serverPort
		req.Params["SERVER_NAME"] = host
		req.Params["SERVER_PROTOCOL"] = r.Proto
		req.Params["SERVER_SOFTWARE"] = "gofast"
		req.Params["REDIRECT_STATUS"] = "200"
		req.Params["REQUEST_METHOD"] = r.Method
		req.Params["REQUEST_URI"] = r.RequestURI
		req.Params["QUERY_STRING"] = r.URL.RawQuery

		return inner(client, req)
	}
}

// FilterAuthReqParams filter out FCGI_PARAMS key-value that is explicitly
// forbidden to passed on in factcgi specification, include:
//  CONTENT_LENGTH;
//  PATH_INFO;
//  PATH_TRANSLATED; and
//  SCRIPT_NAME
func FilterAuthReqParams(inner SessionHandler) SessionHandler {
	return func(client Client, req *Request) (*ResponsePipe, error) {
		if _, ok := req.Params["CONTENT_LENGTH"]; ok {
			delete(req.Params, "CONTENT_LENGTH")
		}
		if _, ok := req.Params["PATH_INFO"]; ok {
			delete(req.Params, "PATH_INFO")
		}
		if _, ok := req.Params["PATH_TRANSLATED"]; ok {
			delete(req.Params, "PATH_TRANSLATED")
		}
		if _, ok := req.Params["SCRIPT_NAME"]; ok {
			delete(req.Params, "SCRIPT_NAME")
		}

		return inner(client, req)
	}
}

// FileSystemRouter helps to produce Middleware implementation for
// mapping path related fastcgi parameters. See method Router for usage.
type FileSystemRouter struct {

	// DocRoot stores the ordinary Apache DocumentRoot parameter
	DocRoot string

	// Exts stores accepted extensions
	Exts []string

	// DirIndex stores ordinary Apache DirectoryIndex parameter
	// for to identify file to show in directory
	DirIndex []string
}

// Router returns a Middleware that prepare session parameters that are
// path related. With information provided in the FileSystemRouter, it will
// route request to script files which path matches the http request path.
//
// i.e. classic PHP hosting environment like Apache + mod_php
//
// Parameters included:
//  PATH_INFO
//  PATH_TRANSLATED
//  SCRIPT_NAME
//  SCRIPT_FILENAME
//  DOCUMENT_URI
//  DOCUMENT_ROOT
//
func (fs *FileSystemRouter) Router() Middleware {
	return func(inner SessionHandler) SessionHandler {
		return func(client Client, req *Request) (*ResponsePipe, error) {

			// define some required cgi parameters
			// with the given http request
			r := req.Raw
			fastcgiScriptName := r.URL.Path

			var fastcgiPathInfo string
			pathinfoRe := regexp.MustCompile(`^(.+\.php)(/?.+)$`)
			if matches := pathinfoRe.FindStringSubmatch(fastcgiScriptName); len(matches) > 0 {
				fastcgiScriptName, fastcgiPathInfo = matches[1], matches[2]
			}

			req.Params["PATH_INFO"] = fastcgiPathInfo
			req.Params["PATH_TRANSLATED"] = filepath.Join(fs.DocRoot, fastcgiPathInfo)
			req.Params["SCRIPT_NAME"] = fastcgiScriptName
			req.Params["SCRIPT_FILENAME"] = filepath.Join(fs.DocRoot, fastcgiScriptName)
			req.Params["DOCUMENT_URI"] = r.URL.Path
			req.Params["DOCUMENT_ROOT"] = fs.DocRoot

			// handle directory index
			urlPath := r.URL.Path
			if strings.HasSuffix(urlPath, "/") {
				urlPath = path.Join(urlPath, "index.php")
			}
			req.Params["SCRIPT_FILENAME"] = path.Join(fs.DocRoot, urlPath)

			return inner(client, req)
		}
	}
}

// MapHeader implement Middleware to map header field HTTP_*
//
// It is a convention to map header field SomeRandomField to
// HTTP_SOME_RANDOM_FIELD. For example, if a header field "X-Hello-World" is in
// the header, it will be mapped as "X_HELLO_WORLD" in the fastcgi parameter
// field.
//
// Note: HTTP_CONTENT_TYPE and HTTP_CONTENT_LENGTH cannot be overridden.
//
func MapHeader(inner SessionHandler) SessionHandler {
	return func(client Client, req *Request) (*ResponsePipe, error) {
		r := req.Raw

		// http header
		for k, v := range r.Header {
			formattedKey := strings.Replace(strings.ToUpper(k), "-", "_", -1)
			if formattedKey == "CONTENT_TYPE" || formattedKey == "CONTENT_LENGTH" {
				continue
			}

			key := "HTTP_" + formattedKey
			var value string
			if len(v) > 0 {
				//   refer to https://tools.ietf.org/html/rfc7230#section-3.2.2
				//
				//   A recipient MAY combine multiple header fields with the same field
				//   name into one "field-name: field-value" pair, without changing the
				//   semantics of the message, by appending each subsequent field value to
				//   the combined field value in order, separated by a comma.  The order
				//   in which header fields with the same field name are received is
				//   therefore significant to the interpretation of the combined field
				//   value; a proxy MUST NOT change the order of these field values when
				//   forwarding a message.
				value = strings.Join(v, ",")
			}
			req.Params[key] = value
		}

		return inner(client, req)
	}
}

// MapEndpoint returns a Middleware implementation that prepare session for
// application with only 1 file as endpoint (i.e. it will handle script routing
// on its own). Suitable for web.py based application.
//
// Parameters included:
//  PATH_INFO
//  PATH_TRANSLATED
//  SCRIPT_NAME
//  SCRIPT_FILENAME
//  DOCUMENT_URI
//  DOCUMENT_ROOT
//
func MapEndpoint(endpointFile string) Middleware {
	dir, webpath := filepath.Dir(endpointFile), "/"+filepath.Base(endpointFile)
	return func(inner SessionHandler) SessionHandler {
		return func(client Client, req *Request) (*ResponsePipe, error) {
			r := req.Raw
			req.Params["REQUEST_URI"] = webpath + r.URL.RequestURI()
			req.Params["SCRIPT_NAME"] = webpath
			req.Params["SCRIPT_FILENAME"] = endpointFile
			req.Params["DOCUMENT_URI"] = r.URL.Path
			req.Params["DOCUMENT_ROOT"] = dir
			return inner(client, req)
		}
	}
}

// MapFilterRequest changes the request role to RoleFilter and add the
// Data stream from the given file system, if file exists. Also
// set the required params to request.
//
// If the file do not exists or cannot be opened, the middleware
// will return empty response pipe and the error.
func MapFilterRequest(fs http.FileSystem) Middleware {
	return func(inner SessionHandler) SessionHandler {
		return func(client Client, req *Request) (*ResponsePipe, error) {

			// force role to be RoleFilter
			req.Role = RoleFilter

			// define some required cgi parameters
			// with the given http request
			r := req.Raw
			fastcgiScriptName := r.URL.Path

			var fastcgiPathInfo string
			pathinfoRe := regexp.MustCompile(`^(.+\.php)(/?.+)$`)
			if matches := pathinfoRe.FindStringSubmatch(fastcgiScriptName); len(matches) > 0 {
				fastcgiScriptName, fastcgiPathInfo = matches[1], matches[2]
			}

			req.Params["PATH_INFO"] = fastcgiPathInfo
			req.Params["SCRIPT_NAME"] = fastcgiScriptName
			req.Params["DOCUMENT_URI"] = r.URL.Path

			// handle directory index
			urlPath := r.URL.Path
			if strings.HasSuffix(urlPath, "/") {
				urlPath = path.Join(urlPath, "index.php")
			}

			// find the file
			f, err := fs.Open(urlPath)
			if err != nil {
				err = fmt.Errorf("cannot open file: %s", err)
				return nil, err
			}

			// map fcgi params for filtering
			s, err := f.Stat()
			if err != nil {
				err = fmt.Errorf("cannot stat file: %s", err)
				return nil, err
			}
			req.Params["FCGI_DATA_LAST_MOD"] = fmt.Sprintf("%d", s.ModTime().Unix())
			req.Params["FCGI_DATA_LENGTH"] = fmt.Sprintf("%d", s.Size())

			// use the file as FCGI_DATA in request
			req.Data = f
			return inner(client, req)
		}
	}
}

// NewFilterLocalFS is a shortcut to use NewFilterFS with
// a http.FileSystem created for the given local folder.
func NewFilterLocalFS(root string) Middleware {
	fs := httpfs.New(vfs.OS(root))
	return NewFilterFS(fs)
}

// NewFilterFS chains BasicParamsMap, MapHeader and MapFilterRequest
// to implement Middleware that prepares a fastcgi Filter session
// environment.
func NewFilterFS(fs http.FileSystem) Middleware {
	return Chain(
		BasicParamsMap,
		MapHeader,
		MapFilterRequest(fs),
	)
}

// NewPHPFS chains BasicParamsMap, MapHeader and FileSystemRouter to implement
// Middleware that prepares an ordinary PHP hosting session environment.
func NewPHPFS(root string) Middleware {
	fs := &FileSystemRouter{
		DocRoot:  root,
		Exts:     []string{"php"},
		DirIndex: []string{"index.php"},
	}
	return Chain(
		BasicParamsMap,
		MapHeader,
		fs.Router(),
	)
}

// NewFileEndpoint chains BasicParamsMap, MapHeader and MapEndpoint to implement
// Middleware that prepares an ordinary web.py hosting session environment
func NewFileEndpoint(endpointFile string) Middleware {
	return Chain(
		BasicParamsMap,
		MapHeader,
		MapEndpoint(endpointFile),
	)
}

// NewAuthPrepare chains BasicParamsMap, MapHeader, FilterAuthReqParams to
// implement Middleware that prepares a authorizer request
func NewAuthPrepare() Middleware {
	return Chain(
		BasicParamsMap,
		MapHeader,
		FilterAuthReqParams,
	)
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package utils

import (
	"os"
	"os/exec"

        "launchpad.net/loggo"
)

var aptLogger = loggo.GetLogger("juju.utils.apt")

// Some helpful functions for running apt in a sane way

// TODO: When we have a unit-level lock to avoid multiple unit agents running
//       apt concurrently, we should use that same locking here

// osRunCommand calls cmd.Run, this is used as an overloading point so we can
// test what *would* be run without actually executing another program
func osRunCommand(cmd *exec.Cmd) error {
	return cmd.Run()
}

var runCommand = osRunCommand

// This is the default apt-get command used in cloud-init, the various settings
// mean that apt won't actually block waiting for a prompt from the user.
var aptGetCommand = []string{
	"apt-get", "--option=Dpkg::Options::=--force-confold",
	"--option=Dpkg::options::=--force-unsafe-io", "--assume-yes", "--quiet",
}

// aptEnvOptions are options we need to pass to apt-get to not have it prompt
// the user
var aptGetEnvOptions = []string{"DEBIAN_FRONTEND=noninteractive"}

// AptGetInstall runs 'apt-get install packages' for the packages listed here
func AptGetInstall(packages ...string) error {
	cmdArgs := append([]string(nil), aptGetCommand...)
	cmdArgs = append(cmdArgs, "install")
	cmdArgs = append(cmdArgs, packages...)
        aptLogger.Infof("Running: %s", cmdArgs)
	cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
	cmd.Env = append(os.Environ(), aptGetEnvOptions...)
	return runCommand(cmd)
}
<commit_msg>go fmt<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package utils

import (
	"os"
	"os/exec"

	"launchpad.net/loggo"
)

var aptLogger = loggo.GetLogger("juju.utils.apt")

// Some helpful functions for running apt in a sane way

// TODO: When we have a unit-level lock to avoid multiple unit agents running
//       apt concurrently, we should use that same locking here

// osRunCommand calls cmd.Run, this is used as an overloading point so we can
// test what *would* be run without actually executing another program
func osRunCommand(cmd *exec.Cmd) error {
	return cmd.Run()
}

var runCommand = osRunCommand

// This is the default apt-get command used in cloud-init, the various settings
// mean that apt won't actually block waiting for a prompt from the user.
var aptGetCommand = []string{
	"apt-get", "--option=Dpkg::Options::=--force-confold",
	"--option=Dpkg::options::=--force-unsafe-io", "--assume-yes", "--quiet",
}

// aptEnvOptions are options we need to pass to apt-get to not have it prompt
// the user
var aptGetEnvOptions = []string{"DEBIAN_FRONTEND=noninteractive"}

// AptGetInstall runs 'apt-get install packages' for the packages listed here
func AptGetInstall(packages ...string) error {
	cmdArgs := append([]string(nil), aptGetCommand...)
	cmdArgs = append(cmdArgs, "install")
	cmdArgs = append(cmdArgs, packages...)
	aptLogger.Infof("Running: %s", cmdArgs)
	cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
	cmd.Env = append(os.Environ(), aptGetEnvOptions...)
	return runCommand(cmd)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/go-ini/ini"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

type Profile struct {
	RoleArn            string
	SourceProfile      string
	MfaSerial          string
	AwsAccessKeyId     string
	AwsSecretAccessKey string
	Region             string
	Token              string
	Name               string
}

var ErrNoAccessKeyGiven = errors.New("no access key given")
var ErrUnknownRegion = errors.New("unknown region given")

func getProfile(profiles []string, iniFile ini.File, hasPrefix bool) (profile Profile, err error) {
	for _, p := range profiles {
		n := p
		if hasPrefix {
			n = "profile " + n
		}
		var section, err = iniFile.GetSection(n)
		if section != nil && err == nil {
			if section.HasKey("mfa_serial") {
				profile.MfaSerial = section.Key("mfa_serial").String()
			}
			if section.HasKey("source_profile") {
				profile.SourceProfile = section.Key("source_profile").String()
			}
			if section.HasKey("region") {
				profile.Region = section.Key("region").String()
			}
			if section.HasKey("role_arn") {
				profile.RoleArn = section.Key("role_arn").String()
			}
			if section.HasKey("aws_access_key_id") {
				profile.AwsAccessKeyId = section.Key("aws_access_key_id").String()
			}
			if section.HasKey("aws_secret_access_key") {
				profile.AwsSecretAccessKey = section.Key("aws_secret_access_key").String()
			}
			break
		}
	}
	return
}

func getProfileKeys(profileName string) (profiles []string) {
	profiles = append(profiles, profileName)
	lowerProjectName := strings.ToLower(profileName)
	profiles = append(profiles, strings.Replace(lowerProjectName, " ", "_", -1))
	profiles = append(profiles, strings.Replace(lowerProjectName, " ", "-", -1))
	return
}

func getAWSConf(projectName string) (sess *session.Session, err error) {
	var creds *credentials.Credentials
	hasPrefix := false
	confFn := os.Getenv("AWS_CONFIG_FILE")
	if confFn == "" {
		confFn = os.Getenv("HOME") + "/.aws/credentials"
		if _, err = os.Stat(confFn); os.IsNotExist(err) {
			confFn = os.Getenv("HOME") + "/.aws/config"
			hasPrefix = true
		}
	}
	if os.Getenv("AWS_ACCESS_KEY_ID") != "" && os.Getenv("AWS_SECRET_ACCESS_KEY") != "" && (os.Getenv("AWS_DEFAULT_REGION") != "" || os.Getenv("AWS_REGION") != "") {
		creds = credentials.NewEnvCredentials()
		region := os.Getenv("AWS_DEFAULT_REGION")
		sess = session.New(&aws.Config{Credentials: creds, Region: &region})
	} else {
		var iniFile *ini.File
		iniFile, err = ini.Load(confFn)
		if err != nil {
			log.Fatalf("Failed to load AWS credentials file  %s", confFn)
		}
		profileKeys := getProfileKeys(projectName)
		profile, _ := getProfile(profileKeys, *iniFile, hasPrefix)
		profile.Name = projectName
		if profile.SourceProfile != "" {
			source_profile, err := getProfile([]string{profile.SourceProfile}, *iniFile, hasPrefix)
			if err == nil {
				profile.AwsAccessKeyId = source_profile.AwsAccessKeyId
				profile.AwsSecretAccessKey = source_profile.AwsSecretAccessKey
				profile.Region = source_profile.Region
			} else {
				log.Fatalf("Failed to load source profile %s", profile.SourceProfile)
			}
		}
		creds = loadCachedCreds(profile)
		if creds == nil {
			if profile.RoleArn != "" {
				if profile.MfaSerial != "" {
					profile.Token = readToken()
				}
				creds = getStsCredentials(profile)
			} else {
				creds = credentials.NewStaticCredentials(profile.AwsAccessKeyId, profile.AwsSecretAccessKey, "")
				creds.Get()
			}
		}
		sess = session.New(&aws.Config{Credentials: creds, Region: &profile.Region})
	}

	return
}

func getStsCredentials(profile Profile) (creds *credentials.Credentials) {
	staticCreds := credentials.NewStaticCredentials(profile.AwsAccessKeyId, profile.AwsSecretAccessKey, "")
	staticCreds.Get()
	client := sts.New(session.New(&aws.Config{Credentials: staticCreds, Region: &profile.Region}))

	sessionName := "AWS-Profile-session-" + strconv.Itoa(int(time.Now().Unix()))
	input := sts.AssumeRoleInput{
		RoleArn:         &profile.RoleArn,
		SerialNumber:    &profile.MfaSerial,
		RoleSessionName: &sessionName,
		TokenCode:       &profile.Token,
	}

	output, err := client.AssumeRole(&input)
	if err != nil {
		log.Fatal(err)
	}
	saveCachedCreds(profile, output)

	creds = credentials.NewStaticCredentials(*output.Credentials.AccessKeyId, *output.Credentials.SecretAccessKey, *output.Credentials.SessionToken)
	creds.Get()
	return
}

func readToken() (token string) {
	var err error
	for {
		fmt.Print("Enter MFA code: ")
		fmt.Scanln(&token)
		if len(token) != 6 {
			fmt.Println("Please make sure your token length is 6")
			continue
		}
		_, err = strconv.Atoi(token)
		if err != nil {
			fmt.Println("Please make sure your token is an integer")
			continue
		}
		return
	}
}

func getCachePath(profile Profile) (path string) {
	path = strings.Replace(profile.RoleArn, ":", "_", -1)
	path = strings.Replace(path, "/", "-", -1)
	path = profile.Name + "--" + path + ".json"
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	path = filepath.Join(usr.HomeDir, ".aws/cli/cache/", path)
	return
}

func loadCachedCreds(profile Profile) (creds *credentials.Credentials) {
	path := getCachePath(profile)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return
	}

	b, err := ioutil.ReadFile(path)

	if err != nil {
		log.Fatalf("Failed to read cache path %s", path)
	}
	assumeRole := new(sts.AssumeRoleOutput)
	err = json.Unmarshal(b, &assumeRole)
	if err != nil {
		log.Fatal(err)
	}
	now := time.Now()
	if now.Unix() > assumeRole.Credentials.Expiration.Unix() {
		return
	}

	creds = credentials.NewStaticCredentials(*assumeRole.Credentials.AccessKeyId, *assumeRole.Credentials.SecretAccessKey, *assumeRole.Credentials.SessionToken)
	creds.Get()
	return
}

func saveCachedCreds(profile Profile, assumeRoleOutput *sts.AssumeRoleOutput) {
	path := getCachePath(profile)
	b, err := json.Marshal(assumeRoleOutput)
	if err != nil {
		log.Fatal(err)
	}
	err = ioutil.WriteFile(path, b, 0600)
	if err != nil {
		log.Fatalf("Failed to write to cache path %s", path)
	}
}
<commit_msg>Catch potential error reading from stdin<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/go-ini/ini"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

type Profile struct {
	RoleArn            string
	SourceProfile      string
	MfaSerial          string
	AwsAccessKeyId     string
	AwsSecretAccessKey string
	Region             string
	Token              string
	Name               string
}

var ErrNoAccessKeyGiven = errors.New("no access key given")
var ErrUnknownRegion = errors.New("unknown region given")

func getProfile(profiles []string, iniFile ini.File, hasPrefix bool) (profile Profile, err error) {
	for _, p := range profiles {
		n := p
		if hasPrefix {
			n = "profile " + n
		}
		var section, err = iniFile.GetSection(n)
		if section != nil && err == nil {
			if section.HasKey("mfa_serial") {
				profile.MfaSerial = section.Key("mfa_serial").String()
			}
			if section.HasKey("source_profile") {
				profile.SourceProfile = section.Key("source_profile").String()
			}
			if section.HasKey("region") {
				profile.Region = section.Key("region").String()
			}
			if section.HasKey("role_arn") {
				profile.RoleArn = section.Key("role_arn").String()
			}
			if section.HasKey("aws_access_key_id") {
				profile.AwsAccessKeyId = section.Key("aws_access_key_id").String()
			}
			if section.HasKey("aws_secret_access_key") {
				profile.AwsSecretAccessKey = section.Key("aws_secret_access_key").String()
			}
			break
		}
	}
	return
}

func getProfileKeys(profileName string) (profiles []string) {
	profiles = append(profiles, profileName)
	lowerProjectName := strings.ToLower(profileName)
	profiles = append(profiles, strings.Replace(lowerProjectName, " ", "_", -1))
	profiles = append(profiles, strings.Replace(lowerProjectName, " ", "-", -1))
	return
}

func getAWSConf(projectName string) (sess *session.Session, err error) {
	var creds *credentials.Credentials
	hasPrefix := false
	confFn := os.Getenv("AWS_CONFIG_FILE")
	if confFn == "" {
		confFn = os.Getenv("HOME") + "/.aws/credentials"
		if _, err = os.Stat(confFn); os.IsNotExist(err) {
			confFn = os.Getenv("HOME") + "/.aws/config"
			hasPrefix = true
		}
	}
	if os.Getenv("AWS_ACCESS_KEY_ID") != "" && os.Getenv("AWS_SECRET_ACCESS_KEY") != "" && (os.Getenv("AWS_DEFAULT_REGION") != "" || os.Getenv("AWS_REGION") != "") {
		creds = credentials.NewEnvCredentials()
		region := os.Getenv("AWS_DEFAULT_REGION")
		sess = session.New(&aws.Config{Credentials: creds, Region: &region})
	} else {
		var iniFile *ini.File
		iniFile, err = ini.Load(confFn)
		if err != nil {
			log.Fatalf("Failed to load AWS credentials file  %s", confFn)
		}
		profileKeys := getProfileKeys(projectName)
		profile, _ := getProfile(profileKeys, *iniFile, hasPrefix)
		profile.Name = projectName
		if profile.SourceProfile != "" {
			source_profile, err := getProfile([]string{profile.SourceProfile}, *iniFile, hasPrefix)
			if err == nil {
				profile.AwsAccessKeyId = source_profile.AwsAccessKeyId
				profile.AwsSecretAccessKey = source_profile.AwsSecretAccessKey
				profile.Region = source_profile.Region
			} else {
				log.Fatalf("Failed to load source profile %s", profile.SourceProfile)
			}
		}
		creds = loadCachedCreds(profile)
		if creds == nil {
			if profile.RoleArn != "" {
				if profile.MfaSerial != "" {
					profile.Token = readToken()
				}
				creds = getStsCredentials(profile)
			} else {
				creds = credentials.NewStaticCredentials(profile.AwsAccessKeyId, profile.AwsSecretAccessKey, "")
				creds.Get()
			}
		}
		sess = session.New(&aws.Config{Credentials: creds, Region: &profile.Region})
	}

	return
}

func getStsCredentials(profile Profile) (creds *credentials.Credentials) {
	staticCreds := credentials.NewStaticCredentials(profile.AwsAccessKeyId, profile.AwsSecretAccessKey, "")
	staticCreds.Get()
	client := sts.New(session.New(&aws.Config{Credentials: staticCreds, Region: &profile.Region}))

	sessionName := "AWS-Profile-session-" + strconv.Itoa(int(time.Now().Unix()))
	input := sts.AssumeRoleInput{
		RoleArn:         &profile.RoleArn,
		SerialNumber:    &profile.MfaSerial,
		RoleSessionName: &sessionName,
		TokenCode:       &profile.Token,
	}

	output, err := client.AssumeRole(&input)
	if err != nil {
		log.Fatal(err)
	}
	saveCachedCreds(profile, output)

	creds = credentials.NewStaticCredentials(*output.Credentials.AccessKeyId, *output.Credentials.SecretAccessKey, *output.Credentials.SessionToken)
	creds.Get()
	return
}

func readToken() (token string) {
	var err error
	for {
		fmt.Print("Enter MFA code: ")
		_, err = fmt.Scanln(&token)
		if err != nil {
			fmt.Println("There was a problem reading from stdin")
			continue
		}
		if len(token) != 6 {
			fmt.Println("Please make sure your token length is 6")
			continue
		}
		_, err = strconv.Atoi(token)
		if err != nil {
			fmt.Println("Please make sure your token is an integer")
			continue
		}
		return
	}
}

func getCachePath(profile Profile) (path string) {
	path = strings.Replace(profile.RoleArn, ":", "_", -1)
	path = strings.Replace(path, "/", "-", -1)
	path = profile.Name + "--" + path + ".json"
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	path = filepath.Join(usr.HomeDir, ".aws/cli/cache/", path)
	return
}

func loadCachedCreds(profile Profile) (creds *credentials.Credentials) {
	path := getCachePath(profile)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return
	}

	b, err := ioutil.ReadFile(path)

	if err != nil {
		log.Fatalf("Failed to read cache path %s", path)
	}
	assumeRole := new(sts.AssumeRoleOutput)
	err = json.Unmarshal(b, &assumeRole)
	if err != nil {
		log.Fatal(err)
	}
	now := time.Now()
	if now.Unix() > assumeRole.Credentials.Expiration.Unix() {
		return
	}

	creds = credentials.NewStaticCredentials(*assumeRole.Credentials.AccessKeyId, *assumeRole.Credentials.SecretAccessKey, *assumeRole.Credentials.SessionToken)
	creds.Get()
	return
}

func saveCachedCreds(profile Profile, assumeRoleOutput *sts.AssumeRoleOutput) {
	path := getCachePath(profile)
	b, err := json.Marshal(assumeRoleOutput)
	if err != nil {
		log.Fatal(err)
	}
	err = ioutil.WriteFile(path, b, 0600)
	if err != nil {
		log.Fatalf("Failed to write to cache path %s", path)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/anchor/bletchley/dataframe"
	"github.com/anchor/bletchley/framestore"
	"github.com/anchor/bletchley/framestore/vaultaire"
	"io"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"time"
)

const (
	Version = "1.0.0"
)

const (
	DefaultParallelism = 1
	DefaultBatchSize   = 1000
	Version            = "0.1"
)

func WriteFrame(semaphore, resultChan chan int, writer framestore.DataFrameWriter, frame *dataframe.DataFrame) {
	// Block until we're ready to write
	<-semaphore
	err := writer.WriteFrame(frame)
	semaphore <- 1
	if err == nil {
		resultChan <- 1
	} else {
		log.Println("Write failed: ", err)
		resultChan <- 0
	}
}

func main() {
	rcfile := flag.String("cfg", "/etc/bletchley/framestore.gcfg", "Path to configuration file. This file should be in gcfg[0] format. [0] https://code.google.com/p/gcfg/")
	version := flag.Bool("version", false, "Print version number and then exit.")

	flag.Usage = func() {
		helpMessage := "bletchley_analyse will export DataFrames to a storage backend for analysis.\n\n" +
			"If no DataFrame files are passed on the command-line, it will read from stdin.\n\n" +
			fmt.Sprintf("Usage: %s [options] [datafile0 [datafile1 ... ] ]\n\n", os.Args[0]) +
			"Options:\n\n"
		fmt.Fprintf(os.Stderr, helpMessage)
		flag.PrintDefaults()
	}
	flag.Parse()

	if *version {
		fmt.Println(Version)
		os.Exit(0)
	}

	cfgPath := *rcfile
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	if !filepath.IsAbs(cfgPath) {
		cfgPath = filepath.Join(usr.HomeDir, cfgPath)
	}
	cfg, err := InitializeConfig(cfgPath)
	if err != nil {
		log.Fatalf("Cannot initialize configuration "+
			"from config file at %v: %v", cfgPath, err)
	}

	parallelism := DefaultParallelism
	if cfg.General.Parallelism > 0 {
		parallelism = cfg.General.Parallelism
	}

	batchSize := DefaultBatchSize
	if cfg.General.BatchSize > 0 {
		batchSize = cfg.General.BatchSize
	}

	var writer framestore.DataFrameWriter

	if cfg.General.StorageBackend == "file" {
		writer, err = framestore.NewFileWriter(cfg.File.DataFrameFile)
	} else if cfg.General.StorageBackend == "vaultaire" {
		writer, err = vaultaire.NewVaultaireWriter(cfg.Vaultaire.Broker, cfg.Vaultaire.BatchPeriod, cfg.Vaultaire.Origin, "", cfg.Vaultaire.MarquiseDebug)
	} else {
		Log.Infof("No backend specified. Exiting.")
		os.Exit(0)
	}
	if err != nil {
		Log.Fatalf("Couldn't initialize writer: ", err)
	}

	defer writer.Shutdown()

	writerThreads := 0
	framesWritten := 0
	writeStart := time.Now()

	Log.Debugf("Starting writes at %v", writeStart)

	semaphore := make(chan int, parallelism)
	resultChannel := make(chan int, 0)

	for i := 0; i < parallelism; i++ {
		semaphore <- 1
	}
	dataFiles := make([]*os.File, 0)
	if len(os.Args) <= 1 {
		dataFiles = append(dataFiles, os.Stdin)
	} else {
		for _, path := range os.Args[1:] {
			fi, err := os.Open(path)
			if err != nil {
				Log.Fatalf(fmt.Sprintf("Couldn't open file %v: %v", err))
			}
			dataFiles = append(dataFiles, fi)
		}
	}
	for _, file := range dataFiles {
		var buf bytes.Buffer
		_, err = buf.ReadFrom(file)
		if err != nil {
			fmt.Println(err)
		}
		burst, err := dataframe.UnmarshalDataBurst(buf.Bytes())
		if err != nil {
			fmt.Println(err)
		}
		for _, frame := range burst.Frames {
			go WriteFrame(semaphore, resultChannel, writer, frame)
			writerThreads += 1
			// We want to kick off the channel-reader at regular
			// intervals in order to avoid exhausting our stack
			// space with waiting goroutines.
			if writerThreads%batchSize == 0 {
				for i := 0; i < writerThreads; i++ {
					framesWritten += <-resultChannel
					writerThreads--
				}
			}
		}
		for i := 0; i < writerThreads; i++ {
			framesWritten += <-resultChannel
		}
		if err != io.EOF {
			fmt.Println(err)
		}
	}
	writeEnd := time.Now()
	delta := writeEnd.Sub(writeStart)
	Log.Debugf("\nWrote %v frames in %v seconds at %v frames/second.\n", framesWritten, delta.Seconds(), float64(framesWritten)/delta.Seconds())
}
<commit_msg>Oops, declared Version twice<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/anchor/bletchley/dataframe"
	"github.com/anchor/bletchley/framestore"
	"github.com/anchor/bletchley/framestore/vaultaire"
	"io"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"time"
)

const (
	Version = "1.0.0"
)

const (
	DefaultParallelism = 1
	DefaultBatchSize   = 1000
)

func WriteFrame(semaphore, resultChan chan int, writer framestore.DataFrameWriter, frame *dataframe.DataFrame) {
	// Block until we're ready to write
	<-semaphore
	err := writer.WriteFrame(frame)
	semaphore <- 1
	if err == nil {
		resultChan <- 1
	} else {
		log.Println("Write failed: ", err)
		resultChan <- 0
	}
}

func main() {
	rcfile := flag.String("cfg", "/etc/bletchley/framestore.gcfg", "Path to configuration file. This file should be in gcfg[0] format. [0] https://code.google.com/p/gcfg/")
	version := flag.Bool("version", false, "Print version number and then exit.")

	flag.Usage = func() {
		helpMessage := "bletchley_analyse will export DataFrames to a storage backend for analysis.\n\n" +
			"If no DataFrame files are passed on the command-line, it will read from stdin.\n\n" +
			fmt.Sprintf("Usage: %s [options] [datafile0 [datafile1 ... ] ]\n\n", os.Args[0]) +
			"Options:\n\n"
		fmt.Fprintf(os.Stderr, helpMessage)
		flag.PrintDefaults()
	}
	flag.Parse()

	if *version {
		fmt.Println(Version)
		os.Exit(0)
	}

	cfgPath := *rcfile
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	if !filepath.IsAbs(cfgPath) {
		cfgPath = filepath.Join(usr.HomeDir, cfgPath)
	}
	cfg, err := InitializeConfig(cfgPath)
	if err != nil {
		log.Fatalf("Cannot initialize configuration "+
			"from config file at %v: %v", cfgPath, err)
	}

	parallelism := DefaultParallelism
	if cfg.General.Parallelism > 0 {
		parallelism = cfg.General.Parallelism
	}

	batchSize := DefaultBatchSize
	if cfg.General.BatchSize > 0 {
		batchSize = cfg.General.BatchSize
	}

	var writer framestore.DataFrameWriter

	if cfg.General.StorageBackend == "file" {
		writer, err = framestore.NewFileWriter(cfg.File.DataFrameFile)
	} else if cfg.General.StorageBackend == "vaultaire" {
		writer, err = vaultaire.NewVaultaireWriter(cfg.Vaultaire.Broker, cfg.Vaultaire.BatchPeriod, cfg.Vaultaire.Origin, "", cfg.Vaultaire.MarquiseDebug)
	} else {
		Log.Infof("No backend specified. Exiting.")
		os.Exit(0)
	}
	if err != nil {
		Log.Fatalf("Couldn't initialize writer: ", err)
	}

	defer writer.Shutdown()

	writerThreads := 0
	framesWritten := 0
	writeStart := time.Now()

	Log.Debugf("Starting writes at %v", writeStart)

	semaphore := make(chan int, parallelism)
	resultChannel := make(chan int, 0)

	for i := 0; i < parallelism; i++ {
		semaphore <- 1
	}
	dataFiles := make([]*os.File, 0)
	if len(os.Args) <= 1 {
		dataFiles = append(dataFiles, os.Stdin)
	} else {
		for _, path := range os.Args[1:] {
			fi, err := os.Open(path)
			if err != nil {
				Log.Fatalf(fmt.Sprintf("Couldn't open file %v: %v", err))
			}
			dataFiles = append(dataFiles, fi)
		}
	}
	for _, file := range dataFiles {
		var buf bytes.Buffer
		_, err = buf.ReadFrom(file)
		if err != nil {
			fmt.Println(err)
		}
		burst, err := dataframe.UnmarshalDataBurst(buf.Bytes())
		if err != nil {
			fmt.Println(err)
		}
		for _, frame := range burst.Frames {
			go WriteFrame(semaphore, resultChannel, writer, frame)
			writerThreads += 1
			// We want to kick off the channel-reader at regular
			// intervals in order to avoid exhausting our stack
			// space with waiting goroutines.
			if writerThreads%batchSize == 0 {
				for i := 0; i < writerThreads; i++ {
					framesWritten += <-resultChannel
					writerThreads--
				}
			}
		}
		for i := 0; i < writerThreads; i++ {
			framesWritten += <-resultChannel
		}
		if err != io.EOF {
			fmt.Println(err)
		}
	}
	writeEnd := time.Now()
	delta := writeEnd.Sub(writeStart)
	Log.Debugf("\nWrote %v frames in %v seconds at %v frames/second.\n", framesWritten, delta.Seconds(), float64(framesWritten)/delta.Seconds())
}
<|endoftext|>


<commit_before>package testdb

import (
	"database/sql"

	_ "github.com/lib/pq"           // register postgresql driver
	_ "github.com/mattn/go-sqlite3" // register sqlite3 driver
)

const (
	pgTruncateTables = `
CREATE OR REPLACE FUNCTION truncate_tables() RETURNS void AS $$
DECLARE
    statements CURSOR FOR
        SELECT tablename FROM pg_tables
        WHERE tablename != 'goose_db_version'
          AND tableowner = session_user
          AND schemaname = 'public';
BEGIN
    FOR stmt IN statements LOOP
        EXECUTE 'TRUNCATE TABLE ' || quote_ident(stmt.tablename) || ' CASCADE;';
    END LOOP;
END;
$$ LANGUAGE plpgsql;

SELECT truncate_tables();
`

	sqliteTruncateTables = `
DELETE FROM certificates;
DELETE FROM ocsp_responses;
`
)

// PostgreSQLDB returns a PostgreSQL db instance for certdb testing.
func PostgreSQLDB() *sql.DB {
	connStr := "dbname=certdb_development host=/var/run/postgresql sslmode=disable"

	db, err := sql.Open("postgres", connStr)
	if err != nil {
		panic(err)
	}

	if _, err := db.Exec(pgTruncateTables); err != nil {
		panic(err)
	}

	return db
}

// SQLiteDB returns a SQLite db instance for certdb testing.
func SQLiteDB(dbpath string) *sql.DB {
	db, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		panic(err)
	}

	if _, err := db.Exec(sqliteTruncateTables); err != nil {
		panic(err)
	}

	return db
}
<commit_msg>use env variable to override testdb connection string<commit_after>package testdb

import (
	"database/sql"

	_ "github.com/lib/pq"           // register postgresql driver
	_ "github.com/mattn/go-sqlite3" // register sqlite3 driver
)

const (
	pgTruncateTables = `
CREATE OR REPLACE FUNCTION truncate_tables() RETURNS void AS $$
DECLARE
    statements CURSOR FOR
        SELECT tablename FROM pg_tables
        WHERE tablename != 'goose_db_version'
          AND tableowner = session_user
          AND schemaname = 'public';
BEGIN
    FOR stmt IN statements LOOP
        EXECUTE 'TRUNCATE TABLE ' || quote_ident(stmt.tablename) || ' CASCADE;';
    END LOOP;
END;
$$ LANGUAGE plpgsql;

SELECT truncate_tables();
`

	sqliteTruncateTables = `
DELETE FROM certificates;
DELETE FROM ocsp_responses;
`
)

// PostgreSQLDB returns a PostgreSQL db instance for certdb testing.
func PostgreSQLDB() *sql.DB {
	connStr := "dbname=certdb_development host=/var/run/postgresql sslmode=disable"

	if dbURL := os.Getenv("DATABASE_URL"); dbURL != "" {
		connStr = dbURL
	}

	db, err := sql.Open("postgres", connStr)
	if err != nil {
		panic(err)
	}

	if _, err := db.Exec(pgTruncateTables); err != nil {
		panic(err)
	}

	return db
}

// SQLiteDB returns a SQLite db instance for certdb testing.
func SQLiteDB(dbpath string) *sql.DB {
	db, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		panic(err)
	}

	if _, err := db.Exec(sqliteTruncateTables); err != nil {
		panic(err)
	}

	return db
}
<|endoftext|>


<commit_before>package checker

import(
	t "testing"
)

func TestCleanChecker(t *t.T) {
	result := CheckFile("scripts/clean.ts", nil)
	if len(result.Warnings) > 0 {
		t.Error("Shouldn't of had any warnings")
	}
}

func TestSelectTextChecker(t *t.T) {
	result := CheckFile("scripts/select_text.ts", nil)
	count := 18
	if len(result.Warnings) != count {
		t.Errorf("Should have thrown %v warnings only gave %d\n", count, len(result.Warnings))
		for _, warn := range(result.Warnings) {
			t.Error(warn.String())
		}
	}
	//println(result.Warnings[0].String())
}

func TestWithNot(t *t.T) {
	result := CheckFile("scripts/with_not.ts", nil)
	if len(result.Warnings) != 2 {
		t.Error("Should have thrown two warnings")
	}
}
<commit_msg>fixed checker test cases.<commit_after>package checker

import (
	t "testing"
)

func TestCleanChecker(t *t.T) {
	result := CheckFile("scripts/clean.ts")
	if len(result.Warnings) > 0 {
		t.Error("Shouldn't of had any warnings")
	}
}

func TestSelectTextChecker(t *t.T) {
	result := CheckFile("scripts/select_text.ts")
	count := 18
	if len(result.Warnings) != count {
		t.Errorf("Should have thrown %v warnings only gave %d\n", count, len(result.Warnings))
		for _, warn := range result.Warnings {
			t.Error(warn.String())
		}
	}
	//println(result.Warnings[0].String())
}

func TestWithNot(t *t.T) {
	result := CheckFile("scripts/with_not.ts")
	if len(result.Warnings) != 2 {
		t.Error("Should have thrown two warnings")
	}
}
<|endoftext|>


<commit_before>/**
 * @file raw_commands.go
 * @author Mikhail Klementyev jollheef<AT>riseup.net
 * @license GNU AGPLv3
 * @date September, 2015
 * @brief functions for run checkers
 *
 * Provide functions for call checker executables
 */

package checker

import (
	"fmt"
	"io"
	"os/exec"
	"strings"
	"time"
)

import "github.com/jollheef/tin_foil_hat/steward"

var (
	timeout            = "10s" // max checker work time
	connectionAttempts = "2"   // ssh option
	connectTimeout     = "5"   // ssh option
)

// SetTimeout set max checker work time
func SetTimeout(d time.Duration) {
	timeout = fmt.Sprintf("%ds", int(d.Seconds()))
}

func readBytesUntilEOF(pipe io.ReadCloser) (buf []byte, err error) {

	bufSize := 1024

	for err != io.EOF {
		stdout := make([]byte, bufSize)
		var n int

		n, err = pipe.Read(stdout)
		if err != nil && err != io.EOF {
			return
		}

		buf = append(buf, stdout[:n]...)
	}

	if err == io.EOF {
		err = nil
	}

	return
}

func readUntilEOF(pipe io.ReadCloser) (str string, err error) {
	buf, err := readBytesUntilEOF(pipe)
	str = string(buf)
	return
}

func system(name string, arg ...string) (stdout string, stderr string,
	err error) {

	cmd := exec.Command(name, arg...)

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

	errPipe, err := cmd.StderrPipe()
	if err != nil {
		return
	}

	cmd.Start()

	stdout, err = readUntilEOF(outPipe)
	if err != nil {
		return
	}

	stderr, err = readUntilEOF(errPipe)
	if err != nil {
		return
	}

	err = cmd.Wait()

	return
}

func exitStatus(no int) string {
	return fmt.Sprintf("exit status %d", no)
}

func parseState(err error) (steward.ServiceState, error) {

	if err == nil {
		return steward.StatusUP, nil
	}

	switch err.Error() {
	case exitStatus(124): // returns by timeout
		return steward.StatusDown, nil
	case exitStatus(1):
	case exitStatus(255): // Could not resolve hostname
		return steward.StatusError, nil
	case exitStatus(2):
		return steward.StatusMumble, nil
	case exitStatus(3):
		return steward.StatusCorrupt, nil
	case exitStatus(4):
		return steward.StatusDown, nil
	}

	return steward.StatusUnknown, err
}

func put(checker, ip string, port int, flag string) (cred, logs string,
	state steward.ServiceState, err error) {

	cred, logs, err = system("timeout", timeout, checker, "put", ip,
		fmt.Sprintf("%d", port), flag)

	state, err = parseState(err)

	cred = strings.Trim(cred, " \n")

	return
}

func sshPut(host, checker, ip string, port int, flag string) (cred, logs string,
	state steward.ServiceState, err error) {

	cred, logs, err = system("ssh",
		"-o", "ConnectTimeout="+connectTimeout,
		"-o", "ConnectionAttempts="+connectionAttempts,
		host, "timeout", timeout, checker,
		"put", ip, fmt.Sprintf("%d", port), flag)

	state, err = parseState(err)

	cred = strings.Trim(cred, " \n")

	return
}

func get(checker, ip string, port int, cred string) (flag, logs string,
	state steward.ServiceState, err error) {

	flag, logs, err = system("timeout", timeout, checker, "get", ip,
		fmt.Sprintf("%d", port), cred)

	state, err = parseState(err)

	flag = strings.Trim(flag, " \n")

	return
}

func sshGet(host, checker, ip string, port int, cred string) (flag, logs string,
	state steward.ServiceState, err error) {

	flag, logs, err = system("ssh",
		"-o", "ConnectTimeout="+connectTimeout,
		"-o", "ConnectionAttempts="+connectionAttempts,
		host, "timeout", timeout, checker,
		"get", ip, fmt.Sprintf("%d", port), cred)

	state, err = parseState(err)

	flag = strings.Trim(flag, " \n")

	return
}

func check(checker, ip string, port int) (state steward.ServiceState,
	logs string, err error) {

	_, logs, err = system("timeout", timeout, checker, "chk", ip,
		fmt.Sprintf("%d", port))

	state, err = parseState(err)

	return
}

func sshCheck(host, checker, ip string, port int) (state steward.ServiceState,
	logs string, err error) {

	_, logs, err = system("ssh",
		"-o", "ConnectTimeout="+connectTimeout,
		"-o", "ConnectionAttempts="+connectionAttempts,
		host, "timeout", timeout, checker,
		"chk", ip, fmt.Sprintf("%d", port))

	state, err = parseState(err)

	return
}
<commit_msg>Move system to library<commit_after>/**
 * @file raw_commands.go
 * @author Mikhail Klementyev jollheef<AT>riseup.net
 * @license GNU AGPLv3
 * @date September, 2015
 * @brief functions for run checkers
 *
 * Provide functions for call checker executables
 */

package checker

import (
	"fmt"
	"strings"
	"time"

	system "github.com/jollheef/go-system"
	"github.com/jollheef/tin_foil_hat/steward"
)

var (
	timeout            = "10s" // max checker work time
	connectionAttempts = "2"   // ssh option
	connectTimeout     = "5"   // ssh option
)

// SetTimeout set max checker work time
func SetTimeout(d time.Duration) {
	timeout = fmt.Sprintf("%ds", int(d.Seconds()))
}

func parseState(ret int) steward.ServiceState {

	switch ret {
	case 0:
		return steward.StatusUP
	case 124: // returns by timeout
		return steward.StatusDown
	case 1:
	case 255: // Could not resolve hostname
		return steward.StatusError
	case 2:
		return steward.StatusMumble
	case 3:
		return steward.StatusCorrupt
	case 4:
		return steward.StatusDown
	}

	return steward.StatusUnknown
}

func put(checker, ip string, port int, flag string) (cred, logs string,
	state steward.ServiceState, err error) {

	cred, logs, ret, err := system.System("timeout", timeout, checker, "put", ip,
		fmt.Sprintf("%d", port), flag)

	state = parseState(ret)
	if state != steward.StatusUnknown {
		err = nil
	}

	cred = strings.Trim(cred, " \n")

	return
}

func sshPut(host, checker, ip string, port int, flag string) (cred, logs string,
	state steward.ServiceState, err error) {

	cred, logs, ret, err := system.System("ssh",
		"-o", "ConnectTimeout="+connectTimeout,
		"-o", "ConnectionAttempts="+connectionAttempts,
		host, "timeout", timeout, checker,
		"put", ip, fmt.Sprintf("%d", port), flag)

	state = parseState(ret)
	if state != steward.StatusUnknown {
		err = nil
	}

	cred = strings.Trim(cred, " \n")

	return
}

func get(checker, ip string, port int, cred string) (flag, logs string,
	state steward.ServiceState, err error) {

	flag, logs, ret, err := system.System("timeout", timeout, checker, "get", ip,
		fmt.Sprintf("%d", port), cred)

	state = parseState(ret)
	if state != steward.StatusUnknown {
		err = nil
	}

	flag = strings.Trim(flag, " \n")

	return
}

func sshGet(host, checker, ip string, port int, cred string) (flag, logs string,
	state steward.ServiceState, err error) {

	flag, logs, ret, err := system.System("ssh",
		"-o", "ConnectTimeout="+connectTimeout,
		"-o", "ConnectionAttempts="+connectionAttempts,
		host, "timeout", timeout, checker,
		"get", ip, fmt.Sprintf("%d", port), cred)

	state = parseState(ret)
	if state != steward.StatusUnknown {
		err = nil
	}

	flag = strings.Trim(flag, " \n")

	return
}

func check(checker, ip string, port int) (state steward.ServiceState,
	logs string, err error) {

	_, logs, ret, err := system.System("timeout", timeout, checker, "chk", ip,
		fmt.Sprintf("%d", port))

	state = parseState(ret)
	if state != steward.StatusUnknown {
		err = nil
	}

	return
}

func sshCheck(host, checker, ip string, port int) (state steward.ServiceState,
	logs string, err error) {

	_, logs, ret, err := system.System("ssh",
		"-o", "ConnectTimeout="+connectTimeout,
		"-o", "ConnectionAttempts="+connectionAttempts,
		host, "timeout", timeout, checker,
		"chk", ip, fmt.Sprintf("%d", port))

	state = parseState(ret)
	if state != steward.StatusUnknown {
		err = nil
	}

	return
}
<|endoftext|>