<commit_before>package longhorn

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/rancher/convoy/convoydriver"
	"github.com/rancher/convoy/util"
	"github.com/rancher/go-rancher-metadata/metadata"

	rancherClient "github.com/rancher/go-rancher/client"
)

const (
	DRIVER_NAME        = "longhorn"
	DRIVER_CONFIG_FILE = "longhorn.cfg"
	MOUNTS_DIR         = "mounts"
	DEV_DIR            = "/dev/longhorn/%s"

	VOLUME_CFG_PREFIX   = "volume_"
	LONGHORN_CFG_PREFIX = DRIVER_NAME + "_"
	CFG_POSTFIX         = ".json"

	SNAPSHOT_PATH = "snapshots"

	DEFAULT_VOLUME_SIZE = "10G"

	RANCHER_METADATA_URL = "http://rancher-metadata/2015-12-19"

	LH_RANCHER_URL         = "lh.rancherurl"
	LH_RANCHER_ACCESS_KEY  = "lh.rancheraccesskey"
	LH_RANCHER_SECRET_KEY  = "lh.ranchersecretkey"
	LH_DEFAULT_VOLUME_SIZE = "lh.defaultvolumesize"
	LH_CONTAINER_NAME      = "lh.containername"

	COMPOSE_VOLUME_NAME = "VOLUME_NAME"
	COMPOSE_VOLUME_SIZE = "VOLUME_SIZE"
	COMPOSE_SLAB_SIZE   = "SLAB_SIZE"
	COMPOSE_CONVOY      = "CONVOY_CONTAINER"

	AFFINITY_LABEL = "io.rancher.scheduler.affinity:container"
)

var (
	log = logrus.WithFields(logrus.Fields{"pkg": "longhorn"})
)

type Driver struct {
	mutex         *sync.RWMutex
	client        *rancherClient.RancherClient
	containerName string
	Device
}

type Device struct {
	Root              string
	DefaultVolumeSize int64
	RancherURL        string
	RancherAccessKey  string
	RancherSecretKey  string
}

func (dev *Device) ConfigFile() (string, error) {
	if dev.Root == "" {
		return "", fmt.Errorf("BUG: Invalid empty device config path")
	}
	return filepath.Join(dev.Root, DRIVER_CONFIG_FILE), nil
}

type Volume struct {
	UUID       string
	Size       int64
	Name       string
	MountPoint string

	configPath string
}

func (v *Volume) Stack(driver *Driver) *Stack {
	sizeString := strconv.FormatInt(v.Size, 10)
	env := map[string]interface{}{
		COMPOSE_SLAB_SIZE:   sizeString,
		COMPOSE_VOLUME_NAME: v.Name,
		COMPOSE_VOLUME_SIZE: sizeString,
		COMPOSE_CONVOY:      driver.containerName,
	}
	return &Stack{
		Client:        driver.client,
		Name:          "longhorn-vol-" + v.Name,
		ExternalId:    "system://longhorn?name=" + v.Name,
		Template:      DockerComposeTemplate,
		Environment:   env,
		ContainerName: driver.containerName,
	}
}

func (v *Volume) ConfigFile() (string, error) {
	if v.UUID == "" {
		return "", fmt.Errorf("BUG: Invalid empty volume UUID")
	}
	if v.configPath == "" {
		return "", fmt.Errorf("BUG: Invalid empty volume config path")
	}
	return filepath.Join(v.configPath, LONGHORN_CFG_PREFIX+VOLUME_CFG_PREFIX+v.UUID+CFG_POSTFIX), nil
}

func (v *Volume) GetDevice() (string, error) {
	return fmt.Sprintf(DEV_DIR, v.Name), nil
}

func (v *Volume) GenerateDefaultMountPoint() string {
	return filepath.Join(v.configPath, MOUNTS_DIR, v.Name)
}

func (d *Driver) blankVolume(id string) *Volume {
	return &Volume{
		configPath: d.Root,
		UUID:       id,
	}
}

func init() {
	convoydriver.Register(DRIVER_NAME, Init)
}

func override(existing, newValue string) string {
	if newValue != "" {
		return newValue
	}
	return existing
}

func Init(root string, config map[string]string) (convoydriver.ConvoyDriver, error) {
	dev := &Device{
		Root: root,
	}
	exists, err := util.ObjectExists(dev)
	if err != nil {
		return nil, err
	}
	if exists {
		if err := util.ObjectLoad(dev); err != nil {
			return nil, err
		}
		dev.RancherURL = override(dev.RancherURL, config[LH_RANCHER_URL])
		dev.RancherAccessKey = override(dev.RancherAccessKey, config[LH_RANCHER_ACCESS_KEY])
		dev.RancherSecretKey = override(dev.RancherSecretKey, config[LH_RANCHER_SECRET_KEY])
	} else {
		if err := util.MkdirIfNotExists(root); err != nil {
			return nil, err
		}

		url := config[LH_RANCHER_URL]
		accessKey := config[LH_RANCHER_ACCESS_KEY]
		secretKey := config[LH_RANCHER_SECRET_KEY]

		if url == "" || accessKey == "" || secretKey == "" {
			return nil, fmt.Errorf("Missing required parameter. lh.rancherurl or lh.rancheraccesskey or lh.ranchersecretkey")
		}

		if _, exists := config[LH_DEFAULT_VOLUME_SIZE]; !exists {
			config[LH_DEFAULT_VOLUME_SIZE] = DEFAULT_VOLUME_SIZE
		}
		volumeSize, err := util.ParseSize(config[LH_DEFAULT_VOLUME_SIZE])
		if err != nil || volumeSize == 0 {
			return nil, fmt.Errorf("Illegal default volume size specified")
		}

		dev = &Device{
			Root:              root,
			RancherURL:        url,
			RancherAccessKey:  accessKey,
			RancherSecretKey:  secretKey,
			DefaultVolumeSize: volumeSize,
		}
	}

	containerName := config[LH_CONTAINER_NAME]
	if containerName == "" {
		handler := metadata.NewClient(RANCHER_METADATA_URL)
		container, err := handler.GetSelfContainer()
		if err != nil {
			return nil, err
		}
		containerName = container.UUID
	}

	log.Debugf("Try to connect to Rancher server at %s [%s:%s]", dev.RancherURL, dev.RancherAccessKey, dev.RancherSecretKey)
	client, err := rancherClient.NewRancherClient(&rancherClient.ClientOpts{
		Url:       dev.RancherURL,
		AccessKey: dev.RancherAccessKey,
		SecretKey: dev.RancherSecretKey,
	})
	if err != nil {
		return nil, fmt.Errorf("Failed to establish connection to Rancher server")
	}

	if err := util.ObjectSave(dev); err != nil {
		return nil, err
	}
	d := &Driver{
		client:        client,
		containerName: containerName,
		Device:        *dev,
	}

	return d, nil
}

func (d *Driver) Name() string {
	return DRIVER_NAME
}

func (d *Driver) Info() (map[string]string, error) {
	return map[string]string{
		"Root":             d.Root,
		"RancherURL":       d.RancherURL,
		"RancherAccessKey": d.RancherAccessKey,
		"RancherSecretKey": d.RancherSecretKey,
	}, nil
}

func (d *Driver) VolumeOps() (convoydriver.VolumeOperations, error) {
	return d, nil
}

func (d *Driver) CreateVolume(id string, opts map[string]string) error {
	size, err := util.ParseSize(opts[convoydriver.OPT_SIZE])
	if err != nil {
		return err
	}
	if size == 0 {
		size = d.DefaultVolumeSize
	}

	volume := d.blankVolume(id)
	volume.Size = size
	volume.Name = opts["VolumeName"]
	stack := volume.Stack(d)

	if err := d.doCreateVolume(volume, stack, id, opts); err != nil {
		stack.Delete()
		return err
	}

	return nil
}

func (d *Driver) doCreateVolume(volume *Volume, stack *Stack, id string, opts map[string]string) error {
	// Doing find just to see if we are creating versus using an existing stack
	env, err := stack.Find()
	if err != nil {
		return err
	}

	// Always run create because it also ensures that things are active
	if _, err := stack.Create(); err != nil {
		return err
	}

	// If env was nil then we created stack so we need to format
	if env == nil {
		dev, _ := volume.GetDevice()
		err := Backoff(5*time.Minute, fmt.Sprintf("Failed to find %s", dev), func() (bool, error) {
			if _, err := os.Stat(dev); err == nil {
				return true, nil
			}
			return false, nil
		})
		if err != nil {
			return err
		}

		log.Infof("Formatting %s", dev)
		if _, err := util.Execute("mkfs.ext4", []string{dev}); err != nil {
			return err
		}
	}

	return util.ObjectSave(volume)
}

func (d *Driver) DeleteVolume(id string, opts map[string]string) error {
	volume := d.blankVolume(id)

	if err := util.ObjectLoad(volume); err != nil {
		return err
	}

	if volume.MountPoint != "" {
		return fmt.Errorf("Cannot delete volume %v. It is still mounted", id)
	}

	referenceOnly, _ := strconv.ParseBool(opts[convoydriver.OPT_REFERENCE_ONLY])
	if !referenceOnly {
		log.Debugf("Deleting stack for volume %v", id)
		if err := volume.Stack(d).Delete(); err != nil {
			return err
		}
	}

	return util.ObjectDelete(volume)
}

func (d *Driver) MountVolume(id string, opts map[string]string) (string, error) {
	volume := d.blankVolume(id)
	if err := util.ObjectLoad(volume); err != nil {
		return "", err
	}

	if err := volume.Stack(d).MoveController(); err != nil {
		log.Errorf("Failed to move controller to %s", d.containerName)
		return "", err
	}

	mountPoint, err := util.VolumeMount(volume, opts[convoydriver.OPT_MOUNT_POINT], false)
	if err != nil {
		return "", err
	}

	if err := util.ObjectSave(volume); err != nil {
		return "", err
	}

	return mountPoint, nil
}

func (d *Driver) UmountVolume(id string) error {
	volume := d.blankVolume(id)
	if err := util.ObjectLoad(volume); err != nil {
		return err
	}

	if err := util.VolumeUmount(volume); err != nil {
		return err
	}

	return util.ObjectSave(volume)
}

func (d *Driver) MountPoint(id string) (string, error) {
	volume := d.blankVolume(id)
	if err := util.ObjectLoad(volume); err != nil {
		return "", err
	}
	return volume.MountPoint, nil
}

func (d *Driver) GetVolumeInfo(id string) (map[string]string, error) {
	volume := d.blankVolume(id)
	if err := util.ObjectLoad(volume); err != nil {
		return nil, err
	}
	return map[string]string{
		"Size": strconv.FormatInt(volume.Size, 10),
	}, nil
}

func (d *Driver) ListVolume(opts map[string]string) (map[string]map[string]string, error) {
	d.mutex.RLock()
	defer d.mutex.RUnlock()

	volumeIDs, err := d.listVolumeIDs()
	if err != nil {
		return nil, err
	}
	result := map[string]map[string]string{}
	for _, id := range volumeIDs {
		result[id], err = d.GetVolumeInfo(id)
		if err != nil {
			return nil, err
		}
	}
	return result, nil
}

func (device *Device) listVolumeIDs() ([]string, error) {
	return util.ListConfigIDs(device.Root, LONGHORN_CFG_PREFIX+VOLUME_CFG_PREFIX, CFG_POSTFIX)
}

func (d *Driver) SnapshotOps() (convoydriver.SnapshotOperations, error) {
	return nil, fmt.Errorf("Longhorn doesn't support snapshot ops")
}

func (d *Driver) BackupOps() (convoydriver.BackupOperations, error) {
	return nil, fmt.Errorf("Longhorn doesn't support backup ops")
}
<commit_msg>Add support for VMs to longhorn driver<commit_after>package longhorn

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/rancher/convoy/convoydriver"
	"github.com/rancher/convoy/util"
	"github.com/rancher/go-rancher-metadata/metadata"

	rancherClient "github.com/rancher/go-rancher/client"
)

const (
	DRIVER_NAME        = "longhorn"
	DRIVER_CONFIG_FILE = "longhorn.cfg"
	MOUNTS_DIR         = "mounts"
	DEV_DIR            = "/dev/longhorn/%s"

	VOLUME_CFG_PREFIX   = "volume_"
	LONGHORN_CFG_PREFIX = DRIVER_NAME + "_"
	CFG_POSTFIX         = ".json"

	SNAPSHOT_PATH = "snapshots"

	DEFAULT_VOLUME_SIZE = "10G"

	RANCHER_METADATA_URL = "http://rancher-metadata/2015-12-19"

	LH_RANCHER_URL         = "lh.rancherurl"
	LH_RANCHER_ACCESS_KEY  = "lh.rancheraccesskey"
	LH_RANCHER_SECRET_KEY  = "lh.ranchersecretkey"
	LH_DEFAULT_VOLUME_SIZE = "lh.defaultvolumesize"
	LH_CONTAINER_NAME      = "lh.containername"

	COMPOSE_VOLUME_NAME = "VOLUME_NAME"
	COMPOSE_VOLUME_SIZE = "VOLUME_SIZE"
	COMPOSE_SLAB_SIZE   = "SLAB_SIZE"
	COMPOSE_CONVOY      = "CONVOY_CONTAINER"

	AFFINITY_LABEL = "io.rancher.scheduler.affinity:container"
)

var (
	log = logrus.WithFields(logrus.Fields{"pkg": "longhorn"})
)

type Driver struct {
	mutex         *sync.RWMutex
	client        *rancherClient.RancherClient
	containerName string
	Device
}

type Device struct {
	Root              string
	DefaultVolumeSize int64
	RancherURL        string
	RancherAccessKey  string
	RancherSecretKey  string
}

func (dev *Device) ConfigFile() (string, error) {
	if dev.Root == "" {
		return "", fmt.Errorf("BUG: Invalid empty device config path")
	}
	return filepath.Join(dev.Root, DRIVER_CONFIG_FILE), nil
}

type Volume struct {
	UUID         string
	Size         int64
	Name         string
	MountPoint   string
	PrepareForVM bool

	configPath string
}

func (v *Volume) Stack(driver *Driver) *Stack {
	sizeString := strconv.FormatInt(v.Size, 10)
	env := map[string]interface{}{
		COMPOSE_SLAB_SIZE:   sizeString,
		COMPOSE_VOLUME_NAME: v.Name,
		COMPOSE_VOLUME_SIZE: sizeString,
		COMPOSE_CONVOY:      driver.containerName,
	}
	return &Stack{
		Client:        driver.client,
		Name:          "longhorn-vol-" + v.Name,
		ExternalId:    "system://longhorn?name=" + v.Name,
		Template:      DockerComposeTemplate,
		Environment:   env,
		ContainerName: driver.containerName,
	}
}

func (v *Volume) ConfigFile() (string, error) {
	if v.UUID == "" {
		return "", fmt.Errorf("BUG: Invalid empty volume UUID")
	}
	if v.configPath == "" {
		return "", fmt.Errorf("BUG: Invalid empty volume config path")
	}
	return filepath.Join(v.configPath, LONGHORN_CFG_PREFIX+VOLUME_CFG_PREFIX+v.UUID+CFG_POSTFIX), nil
}

func (v *Volume) GetDevice() (string, error) {
	return fmt.Sprintf(DEV_DIR, v.Name), nil
}

func (v *Volume) GetMountOpts() []string {
	return []string{}
}

func (v *Volume) GenerateDefaultMountPoint() string {
	return filepath.Join(v.configPath, MOUNTS_DIR, v.Name)
}

func (d *Driver) blankVolume(id string) *Volume {
	return &Volume{
		configPath: d.Root,
		UUID:       id,
	}
}

func init() {
	convoydriver.Register(DRIVER_NAME, Init)
}

func override(existing, newValue string) string {
	if newValue != "" {
		return newValue
	}
	return existing
}

func Init(root string, config map[string]string) (convoydriver.ConvoyDriver, error) {
	dev := &Device{
		Root: root,
	}
	exists, err := util.ObjectExists(dev)
	if err != nil {
		return nil, err
	}
	if exists {
		if err := util.ObjectLoad(dev); err != nil {
			return nil, err
		}
		dev.RancherURL = override(dev.RancherURL, config[LH_RANCHER_URL])
		dev.RancherAccessKey = override(dev.RancherAccessKey, config[LH_RANCHER_ACCESS_KEY])
		dev.RancherSecretKey = override(dev.RancherSecretKey, config[LH_RANCHER_SECRET_KEY])
	} else {
		if err := util.MkdirIfNotExists(root); err != nil {
			return nil, err
		}

		url := config[LH_RANCHER_URL]
		accessKey := config[LH_RANCHER_ACCESS_KEY]
		secretKey := config[LH_RANCHER_SECRET_KEY]

		if url == "" || accessKey == "" || secretKey == "" {
			return nil, fmt.Errorf("Missing required parameter. lh.rancherurl or lh.rancheraccesskey or lh.ranchersecretkey")
		}

		if _, exists := config[LH_DEFAULT_VOLUME_SIZE]; !exists {
			config[LH_DEFAULT_VOLUME_SIZE] = DEFAULT_VOLUME_SIZE
		}
		volumeSize, err := util.ParseSize(config[LH_DEFAULT_VOLUME_SIZE])
		if err != nil || volumeSize == 0 {
			return nil, fmt.Errorf("Illegal default volume size specified")
		}

		dev = &Device{
			Root:              root,
			RancherURL:        url,
			RancherAccessKey:  accessKey,
			RancherSecretKey:  secretKey,
			DefaultVolumeSize: volumeSize,
		}
	}

	containerName := config[LH_CONTAINER_NAME]
	if containerName == "" {
		handler := metadata.NewClient(RANCHER_METADATA_URL)
		container, err := handler.GetSelfContainer()
		if err != nil {
			return nil, err
		}
		containerName = container.UUID
	}

	log.Debugf("Try to connect to Rancher server at %s [%s:%s]", dev.RancherURL, dev.RancherAccessKey, dev.RancherSecretKey)
	client, err := rancherClient.NewRancherClient(&rancherClient.ClientOpts{
		Url:       dev.RancherURL,
		AccessKey: dev.RancherAccessKey,
		SecretKey: dev.RancherSecretKey,
	})
	if err != nil {
		return nil, fmt.Errorf("Failed to establish connection to Rancher server")
	}

	if err := util.ObjectSave(dev); err != nil {
		return nil, err
	}
	d := &Driver{
		client:        client,
		containerName: containerName,
		Device:        *dev,
	}

	return d, nil
}

func (d *Driver) Name() string {
	return DRIVER_NAME
}

func (d *Driver) Info() (map[string]string, error) {
	return map[string]string{
		"Root":             d.Root,
		"RancherURL":       d.RancherURL,
		"RancherAccessKey": d.RancherAccessKey,
		"RancherSecretKey": d.RancherSecretKey,
	}, nil
}

func (d *Driver) VolumeOps() (convoydriver.VolumeOperations, error) {
	return d, nil
}

func (d *Driver) CreateVolume(id string, opts map[string]string) error {
	size, err := util.ParseSize(opts[convoydriver.OPT_SIZE])
	if err != nil {
		return err
	}
	if size == 0 {
		size = d.DefaultVolumeSize
	}

	volume := d.blankVolume(id)
	volume.Size = size
	volume.Name = opts["VolumeName"]
	volume.PrepareForVM, err = strconv.ParseBool(opts[convoydriver.OPT_PREPARE_FOR_VM])
	if err != nil {
		return err
	}

	stack := volume.Stack(d)

	if err := d.doCreateVolume(volume, stack, id, opts); err != nil {
		stack.Delete()
		return err
	}

	return nil
}

func (d *Driver) doCreateVolume(volume *Volume, stack *Stack, id string, opts map[string]string) error {
	// Doing find just to see if we are creating versus using an existing stack
	env, err := stack.Find()
	if err != nil {
		return err
	}

	// Always run create because it also ensures that things are active
	if _, err := stack.Create(); err != nil {
		return err
	}

	// If env was nil then we created stack so we need to format
	if env == nil {
		dev, _ := volume.GetDevice()
		err := Backoff(5*time.Minute, fmt.Sprintf("Failed to find %s", dev), func() (bool, error) {
			if _, err := os.Stat(dev); err == nil {
				return true, nil
			}
			return false, nil
		})
		if err != nil {
			return err
		}

		log.Infof("Formatting %s", dev)
		if _, err := util.Execute("mkfs.ext4", []string{dev}); err != nil {
			return err
		}
	}

	return util.ObjectSave(volume)
}

func (d *Driver) DeleteVolume(id string, opts map[string]string) error {
	volume := d.blankVolume(id)

	if err := util.ObjectLoad(volume); err != nil {
		return err
	}

	if volume.MountPoint != "" {
		return fmt.Errorf("Cannot delete volume %v. It is still mounted", id)
	}

	referenceOnly, _ := strconv.ParseBool(opts[convoydriver.OPT_REFERENCE_ONLY])
	if !referenceOnly {
		log.Debugf("Deleting stack for volume %v", id)
		if err := volume.Stack(d).Delete(); err != nil {
			return err
		}
	}

	return util.ObjectDelete(volume)
}

func (d *Driver) MountVolume(id string, opts map[string]string) (string, error) {
	volume := d.blankVolume(id)
	if err := util.ObjectLoad(volume); err != nil {
		return "", err
	}

	if err := volume.Stack(d).MoveController(); err != nil {
		log.Errorf("Failed to move controller to %s", d.containerName)
		return "", err
	}

	mountPoint, err := util.VolumeMount(volume, opts[convoydriver.OPT_MOUNT_POINT], false)
	if err != nil {
		return "", err
	}

	if volume.PrepareForVM {
		if err := util.MountPointPrepareImageFile(mountPoint, volume.Size); err != nil {
			return "", err
		}
	}

	if err := util.ObjectSave(volume); err != nil {
		return "", err
	}

	return mountPoint, nil
}

func (d *Driver) UmountVolume(id string) error {
	volume := d.blankVolume(id)
	if err := util.ObjectLoad(volume); err != nil {
		return err
	}

	if err := util.VolumeUmount(volume); err != nil {
		return err
	}

	return util.ObjectSave(volume)
}

func (d *Driver) MountPoint(id string) (string, error) {
	volume := d.blankVolume(id)
	if err := util.ObjectLoad(volume); err != nil {
		return "", err
	}
	return volume.MountPoint, nil
}

func (d *Driver) GetVolumeInfo(id string) (map[string]string, error) {
	volume := d.blankVolume(id)
	if err := util.ObjectLoad(volume); err != nil {
		return nil, err
	}
	return map[string]string{
		"Size": strconv.FormatInt(volume.Size, 10),
		convoydriver.OPT_PREPARE_FOR_VM: strconv.FormatBool(volume.PrepareForVM),
	}, nil
}

func (d *Driver) ListVolume(opts map[string]string) (map[string]map[string]string, error) {
	d.mutex.RLock()
	defer d.mutex.RUnlock()

	volumeIDs, err := d.listVolumeIDs()
	if err != nil {
		return nil, err
	}
	result := map[string]map[string]string{}
	for _, id := range volumeIDs {
		result[id], err = d.GetVolumeInfo(id)
		if err != nil {
			return nil, err
		}
	}
	return result, nil
}

func (device *Device) listVolumeIDs() ([]string, error) {
	return util.ListConfigIDs(device.Root, LONGHORN_CFG_PREFIX+VOLUME_CFG_PREFIX, CFG_POSTFIX)
}

func (d *Driver) SnapshotOps() (convoydriver.SnapshotOperations, error) {
	return nil, fmt.Errorf("Longhorn doesn't support snapshot ops")
}

func (d *Driver) BackupOps() (convoydriver.BackupOperations, error) {
	return nil, fmt.Errorf("Longhorn doesn't support backup ops")
}
<|endoftext|>


<commit_before>package blob

import (
	"crypto/sha1"
	"database/sql"
	"encoding/hex"
	"fmt"
	_ "github.com/herenow/go-crate"
	"io"
	"net/http"
	"time"
)

type Driver struct {
	url string
	db  *sql.DB
}

type Table struct {
	Name string
	drv  *Driver
	c    *http.Client
}

// DownloadError represents an error that occurs when downloading a blob.
type DownloadError struct {
	Status string
	StatusCode int
	Message string
}

// Error returns a string representation of the DownloadError.
func (e DownloadError) Error() string {
	return e.Message
}

// New creates a new connection with crate server
func New(crate_url string) (*Driver, error) {
	db, err := sql.Open("crate", crate_url)
	if err != nil {
		return nil, err
	}
	return &Driver{
		url: crate_url,
		db:  db,
	}, nil
}

// NewTable create new blob table with name and extra int to specify
// shards(the second argument) and
// replicas(by the third int argument)
func (d *Driver) NewTable(name string, shards ...int) (*Table, error) {
	sql := fmt.Sprintf(
		"create blob table %s",
		name,
	)
	if len(shards) == 1 {
		sql = fmt.Sprintf(
			"create blob table %s clustered into %d shards",
			name, shards[0],
		)
	}
	if len(shards) >= 2 {
		sql = fmt.Sprintf(
			"create blob table %s clustered into %d shards with (number_of_replicas=%d)",
			name, shards[0], shards[1],
		)
	}
	_, err := d.db.Exec(sql)
	if err != nil {
		return nil, err
	}
	return &Table{
		Name: name,
		drv:  d,
		c:    new(http.Client),
	}, nil
}

// Get an existing table from the crate server
// or error when this table does not exist
func (d *Driver) GetTable(name string) (*Table, error) {
	row := d.db.QueryRow(
		"select table_name from information_schema.tables where table_name = ? and table_schema = 'blob'",
		name,
	)
	if err := row.Scan(&name); err != nil {
		return nil, err
	}
	return &Table{
		Name: name,
		drv:  d,
		c:    new(http.Client),
	}, nil
}

// Close the database connection
func (d *Driver) Close() error {
	if d.db != nil {
		return d.db.Close()
	}
	return nil
}

type Record struct {
	Digest       string
	LastModified time.Time
}

// Sha1Digest calculates the sha1 digest for the io.Reader
func Sha1Digest(r io.Reader) string {
	h := sha1.New()
	io.Copy(h, r)
	return hex.EncodeToString(h.Sum(nil))
}

func (t *Table) url(digest string) string {
	return fmt.Sprintf("%s/_blobs/%s/%s", t.drv.url, t.Name, digest)
}

// Upload upload the blob(r) with sha1 hash(digest)
func (t *Table) Upload(digest string, r io.Reader) (*Record, error) {
	req, err := http.NewRequest("PUT", t.url(digest), r)
	if err != nil {
		return nil, err
	}
	resp, err := t.c.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 201 {
		return nil, fmt.Errorf("Upload failed: %s", resp.Status)
	}
	return &Record{
		Digest: digest,
	}, nil
}

// UploadEx upload a io.ReadSeeker, and computes the sha1 hash automatically
func (t *Table) UploadEx(r io.ReadSeeker) (*Record, error) {
	digest := Sha1Digest(r)
	_, err := r.Seek(0, 0)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("PUT", t.url(digest), r)
	if err != nil {
		return nil, err
	}
	_, err = t.c.Do(req)
	if err != nil {
		return nil, err
	}
	return &Record{
		Digest: digest,
	}, nil
}

// List all blobs inside a blob table
func (t *Table) List() (*sql.Rows, error) {
	query := fmt.Sprintf("select digest, last_modified from blob.%s", t.Name)
	rows, err := t.drv.db.Query(query)
	if err != nil {
		return nil, err
	}
	return rows, err
}

// Is the blob specified by the digest exist in the table
func (t *Table) Has(digest string) (bool, error) {
	req, err := http.NewRequest("HEAD", t.url(digest), nil)
	if err != nil {
		return false, err
	}
	resp, err := t.c.Do(req)
	if err != nil {
		return false, err
	}
	if resp.StatusCode == http.StatusOK {
		return true, nil
	}
	return false, nil
}

// Download a blob in a blob table with the specific digest
func (t *Table) Download(digest string) (io.ReadCloser, error) {
	req, err := http.NewRequest("GET", t.url(digest), nil)
	if err != nil {
		return nil, err
	}
	resp, err := t.c.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return resp.Body, &DownloadError{resp.Status, resp.StatusCode, "problem downloading blob"}
	}
	return resp.Body, nil
}

// Delete a blob in a blob table with the specific digest
func (t *Table) Delete(digest string) error {
	req, err := http.NewRequest("DELETE", t.url(digest), nil)
	if err != nil {
		return err
	}
	resp, err := t.c.Do(req)
	if err != nil {
		return err
	}
	if resp.StatusCode == http.StatusNoContent {
		return nil
	}
	return fmt.Errorf("%s", resp.Status)
}

// Drop the blob table
func (t *Table) Drop() error {
	sql := fmt.Sprintf("drop blob table \"%s\"", t.Name)
	_, err := t.drv.db.Exec(sql)
	return err
}

func (t *Table) newRequest(method, digest string, body io.Reader) (*http.Request, error) {
	req, err := http.NewRequest(method, t.url(digest), body)
	if err != nil {
		return nil ,err
	}
	// in some cases where trying to upload or download a blob resulted in non 2xx response,
	// it seems Crate was not closing the connection (ticket to be opened soon against Crate)
	// which resulted in strange behavior when requests were made after getting an error.
	//
	// For example, if downloading a blob result in a 404, attempting to download another
	// blob would also result ina  404 even if that second blob absolutely was there.
	req.Header.Set("Connection", "close")
	return req, nil
}
<commit_msg>Update blob.go<commit_after>package blob

import (
	"crypto/sha1"
	"database/sql"
	"encoding/hex"
	"fmt"
	_ "github.com/herenow/go-crate"
	"io"
	"net/http"
	"time"
)

type Driver struct {
	url string
	db  *sql.DB
}

type Table struct {
	Name string
	drv  *Driver
	c    *http.Client
}

// DownloadError represents an error that occurs when downloading a blob.
type DownloadError struct {
	Status string
	StatusCode int
	Message string
}

// Error returns a string representation of the DownloadError.
func (e DownloadError) Error() string {
	return e.Message
}

// New creates a new connection with crate server
func New(crate_url string) (*Driver, error) {
	db, err := sql.Open("crate", crate_url)
	if err != nil {
		return nil, err
	}
	return &Driver{
		url: crate_url,
		db:  db,
	}, nil
}

// NewTable create new blob table with name and extra int to specify
// shards(the second argument) and
// replicas(by the third int argument)
func (d *Driver) NewTable(name string, shards ...int) (*Table, error) {
	sql := fmt.Sprintf(
		"create blob table %s",
		name,
	)
	if len(shards) == 1 {
		sql = fmt.Sprintf(
			"create blob table %s clustered into %d shards",
			name, shards[0],
		)
	}
	if len(shards) >= 2 {
		sql = fmt.Sprintf(
			"create blob table %s clustered into %d shards with (number_of_replicas=%d)",
			name, shards[0], shards[1],
		)
	}
	_, err := d.db.Exec(sql)
	if err != nil {
		return nil, err
	}
	return &Table{
		Name: name,
		drv:  d,
		c:    new(http.Client),
	}, nil
}

// Get an existing table from the crate server
// or error when this table does not exist
func (d *Driver) GetTable(name string) (*Table, error) {
	row := d.db.QueryRow(
		"select table_name from information_schema.tables where table_name = ? and table_schema = 'blob'",
		name,
	)
	if err := row.Scan(&name); err != nil {
		return nil, err
	}
	return &Table{
		Name: name,
		drv:  d,
		c:    new(http.Client),
	}, nil
}

// Close the database connection
func (d *Driver) Close() error {
	if d.db != nil {
		return d.db.Close()
	}
	return nil
}

type Record struct {
	Digest       string
	LastModified time.Time
}

// Sha1Digest calculates the sha1 digest for the io.Reader
func Sha1Digest(r io.Reader) string {
	h := sha1.New()
	io.Copy(h, r)
	return hex.EncodeToString(h.Sum(nil))
}

func (t *Table) url(digest string) string {
	return fmt.Sprintf("%s/_blobs/%s/%s", t.drv.url, t.Name, digest)
}

// Upload upload the blob(r) with sha1 hash(digest)
func (t *Table) Upload(digest string, r io.Reader) (*Record, error) {
	req, err := t.newRequest("PUT", digest, r)
	if err != nil {
		return nil, err
	}
	resp, err := t.c.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 201 {
		return nil, fmt.Errorf("Upload failed: %s", resp.Status)
	}
	return &Record{
		Digest: digest,
	}, nil
}

// UploadEx upload a io.ReadSeeker, and computes the sha1 hash automatically
func (t *Table) UploadEx(r io.ReadSeeker) (*Record, error) {
	digest := Sha1Digest(r)
	_, err := r.Seek(0, 0)
	if err != nil {
		return nil, err
	}
	req, err := t.newRequest("PUT", digest, r)
	if err != nil {
		return nil, err
	}
	_, err = t.c.Do(req)
	if err != nil {
		return nil, err
	}
	return &Record{
		Digest: digest,
	}, nil
}

// List all blobs inside a blob table
func (t *Table) List() (*sql.Rows, error) {
	query := fmt.Sprintf("select digest, last_modified from blob.%s", t.Name)
	rows, err := t.drv.db.Query(query)
	if err != nil {
		return nil, err
	}
	return rows, err
}

// Is the blob specified by the digest exist in the table
func (t *Table) Has(digest string) (bool, error) {
	req, err := t.newRequest("HEAD", digest, nil)
	if err != nil {
		return false, err
	}
	resp, err := t.c.Do(req)
	if err != nil {
		return false, err
	}
	if resp.StatusCode == http.StatusOK {
		return true, nil
	}
	return false, nil
}

// Download a blob in a blob table with the specific digest
func (t *Table) Download(digest string) (io.ReadCloser, error) {
	req, err := t.newRequest("GET", digest, nil)
	if err != nil {
		return nil, err
	}
	resp, err := t.c.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return resp.Body, &DownloadError{resp.Status, resp.StatusCode, "problem downloading blob"}
	}
	return resp.Body, nil
}

// Delete a blob in a blob table with the specific digest
func (t *Table) Delete(digest string) error {
	req, err := t.newRequest("DELETE", digest, nil)
	if err != nil {
		return err
	}
	resp, err := t.c.Do(req)
	if err != nil {
		return err
	}
	if resp.StatusCode == http.StatusNoContent {
		return nil
	}
	return fmt.Errorf("%s", resp.Status)
}

// Drop the blob table
func (t *Table) Drop() error {
	sql := fmt.Sprintf("drop blob table \"%s\"", t.Name)
	_, err := t.drv.db.Exec(sql)
	return err
}

func (t *Table) newRequest(method, digest string, body io.Reader) (*http.Request, error) {
	req, err := http.NewRequest(method, t.url(digest), body)
	if err != nil {
		return nil ,err
	}
	// in some cases where trying to upload or download a blob resulted in non 2xx response,
	// it seems Crate was not closing the connection (ticket to be opened soon against Crate)
	// which resulted in strange behavior when requests were made after getting an error.
	//
	// For example, if downloading a blob result in a 404, attempting to download another
	// blob would also result ina  404 even if that second blob absolutely was there.
	req.Header.Set("Connection", "close")
	return req, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/phil-mansfield/gotetra/render/halo"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/gotetra/los/tree"
)

type IDType int
const (
	Rockstar IDType = iota
	M500c
	M200c
	M200m

    finderCells = 150
    overlapMult = 3
)

var (
	ms, rs, xs, ys, zs []float64
	rids []int
)

func main() {
	// Parse command line
	var (
		idTypeStr string
		idStart, idEnd, snap int
		allowSubhalos bool
	)

	flag.StringVar(&idTypeStr, "IDType", "Rockstar",
		"[Rockstar | M500c | M200c | M200m]")
	flag.IntVar(&snap, "Snap", -1, "The snapshot number.")
	flag.IntVar(&idStart, "IDStart", -1, "ID start range (inclusive).")
	flag.IntVar(&idEnd, "IDEnd", -1, "ID stop range (exclusive)")
	flag.BoolVar(&allowSubhalos, "AllowSubhalos", false,
		"Allow subhalo ids to be passed through.")
	flag.Parse()

	if idStart > idEnd {
		log.Fatalf("IDStart, %d, is larger than IDEnd, %d.")
	} else if idStart != idEnd && idStart <= 0 {
		log.Fatalf("Non-positive IDStart %d.")
	}

	idType, err := parseIDType(idTypeStr)
	if err != nil { err.Error() }
	if idType != Rockstar && snap == -1 {
		log.Fatalf("Must set the Snap flag if using a non-default IDType.")
	}

	snapNum, err := snapNum()
	if err != nil {
		log.Fatalf(
			"Error encountered when finding rockstar directory: %s",err.Error(),
		)
	} else if (snap < 1 && idType != Rockstar) || snap > snapNum {
		log.Fatalf("Snap %d is out of bounds for %d snaps.", snap, snapNum)
	}

	// Get IDs and snapshots

	rawIds := getIDs(idStart, idEnd, flag.Args())
	if len(rawIds) == 0 { return }

	fmt.Println(rawIds)

	var ids, snaps []int
	switch idType {
	case Rockstar:
		if snap != -1 {
			snaps = make([]int, len(rawIds))
			for i := range snaps { snaps[i] = snap }
		} else {
			snaps, err = findSnaps(rawIds)
		}
		ids = rawIds
		if err != nil { log.Fatalf(err.Error()) }
	case M200m:
		snaps = make([]int, len(rawIds))
		for i := range snaps { snaps[i] = snap }
		ids, err = convertSortedIDs(rawIds, snap)
		if err != nil { err.Error() }
	default:
		log.Fatal("Unsupported IDType for now. Sorry :3")
	}

	// Tag subhalos, if neccessary.
	var isSub []bool
	if allowSubhalos {
		isSub = make([]bool, len(ids))
	} else {
		isSub, err = findSubs(rawIds, snaps)
		if err != nil { log.Fatal(err.Error()) }
	}

	// Output
	printIds(ids, snaps, isSub)
}

func parseIDType(str string) (IDType, error) {
	switch strings.ToLower(str) {
	case "rockstar": return Rockstar, nil
	case "m500c": return M500c, nil
	case "m200c": return M200c, nil
	case "m200m": return M200m, nil
	}
	return -1, fmt.Errorf("IDType '%s' unrecognized", str)
}

func getIDs(idStart, idEnd int, args []string) []int {
	ids := make([]int, 0, idEnd - idStart + len(args))
	for i := idStart; i < idEnd; i++ { ids = append(ids, i) }
	for _, str := range args {
		i, err := strconv.Atoi(str)
		if err != nil {
			log.Fatalf("Could not parse arg %d: '%s' is not an int.", i+1, str)
		}
		ids = append(ids, i)
	}
	return ids
}

func snapNum() (int, error) {
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	infos, err := ioutil.ReadDir(rockstarDir)
	return len(infos), err
}

func getSnapHaloList(i int) (name string, err error) {
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	infos, err := ioutil.ReadDir(rockstarDir)
	if err != nil { return "", err }
	return path.Join(rockstarDir, infos[i - 1].Name()), nil
}

func findSnaps(ids []int) ([]int, error) {
	treeDir := os.Getenv("GTET_TREE_DIR")
	infos, err := ioutil.ReadDir(treeDir)
	if err != nil { return nil, err }

	names := []string{}
	for _, info := range infos {
		name := info.Name()
		n := len(name)
		if n > 4 && name[:5] == "tree_" && name[n-4:] == ".gtet" {
			names = append(names, name)
		}
	}

	return tree.HaloSnaps(names, ids)
}

// This funciton has side effects :(
func convertSortedIDs(
	rawIDs []int, snap int,
) ([]int, error) {
	list, err := getSnapHaloList(snap)
	if err != nil { return nil, err }

	gtetFmt := os.Getenv("GTET_FMT")
	gtetName := fmt.Sprintf(gtetFmt, snap, 0, 0, 0)
	hd := &io.SheetHeader{}
	err = io.ReadSheetHeaderAt(gtetName, hd)
	if err != nil { return nil, err }
	cosmo := &hd.Cosmo

	rids, xs, ys, zs, ms, rs, err = halo.ReadRockstar(list, halo.R200m, cosmo)
	if err != nil { return nil, err }

	ids := make([]int, len(rawIDs))
	for i := range ids { ids[i] = rids[rawIDs[i]] }
	return ids, nil
}

func findSubs(rawIDs, snaps []int) ([]bool, error) {
	isSub := make([]bool, len(rawIDs))

	// Handle the case where we're indexing by sorted mass ID.
	if xs != nil {
		gtetFmt := os.Getenv("GTET_FMT")
		gtetName := fmt.Sprintf(gtetFmt, snaps[0], 0, 0, 0)
		hd := &io.SheetHeader{}
		err := io.ReadSheetHeaderAt(gtetName, hd)
		if err != nil { return nil, err }

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, zs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, overlapMult)

		for i, id := range rawIDs {
			isSub[i] = sf.HostCount(id) > 0
		}
		return isSub, nil
	}

	// Handle the case where all we have are halo IDs

	// Group by snapshot.
	snapGroups := make(map[int][]int)
	groupIdxs := make(map[int][]int)
	for i, id := range rawIDs {
		snap := snaps[i]
		snapGroups[snap] = append(snapGroups[snap], id)
		groupIdxs[snap] = append(groupIdxs[snap], i)
	}

	// Load each snapshot.
	for snap, group := range snapGroups {
		gtetFmt := os.Getenv("GTET_FMT")
		gtetName := fmt.Sprintf(gtetFmt, snap, 0, 0, 0)
		hd := &io.SheetHeader{}
		err := io.ReadSheetHeaderAt(gtetName, hd)
		if err != nil { return nil, err }

		list, err := getSnapHaloList(snap)
		rids, xs, ys, zs, ms, rs, err = halo.ReadRockstar(
			list, halo.R200m, &hd.Cosmo,
		)

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, zs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, overlapMult)

		for i, id := range group {
			origIdx := groupIdxs[snap][i]
			for j, checkID := range rids {
				if checkID == id {
					isSub[origIdx] = sf.HostCount(j) > 0
					break
				} else if j == len(rids) - 1 {
					return nil, fmt.Errorf("ID %d not in halo list.", id)
				}
			}
		}
	}
	return isSub, nil
}

func printIds(ids []int, snaps []int, isSub []bool) {
	// Find the maximum width of each column.
	idWidth, snapWidth := 0, 0
	for i := range ids {
		if isSub[i] { continue }
		iWidth := len(fmt.Sprintf("%d", ids[i]))
		sWidth := len(fmt.Sprintf("%d", snaps[i]))
		if iWidth > idWidth { idWidth = iWidth }
		if sWidth > snapWidth { snapWidth = sWidth }
	}

	rowFmt := fmt.Sprintf("%%%dd %%%dd\n", idWidth, snapWidth)
	for i := range ids {
		if isSub[i] { continue }
		fmt.Printf(rowFmt, ids[i], snaps[i])
	}
}
<commit_msg>Fixed bug in tree_*.dat searching code.<commit_after>package main

import (
	"fmt"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/phil-mansfield/gotetra/render/halo"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/gotetra/los/tree"
)

type IDType int
const (
	Rockstar IDType = iota
	M500c
	M200c
	M200m

    finderCells = 150
    overlapMult = 3
)

var (
	ms, rs, xs, ys, zs []float64
	rids []int
)

func main() {
	// Parse command line
	var (
		idTypeStr string
		idStart, idEnd, snap int
		allowSubhalos bool
	)

	flag.StringVar(&idTypeStr, "IDType", "Rockstar",
		"[Rockstar | M500c | M200c | M200m]")
	flag.IntVar(&snap, "Snap", -1, "The snapshot number.")
	flag.IntVar(&idStart, "IDStart", -1, "ID start range (inclusive).")
	flag.IntVar(&idEnd, "IDEnd", -1, "ID stop range (exclusive)")
	flag.BoolVar(&allowSubhalos, "AllowSubhalos", false,
		"Allow subhalo ids to be passed through.")
	flag.Parse()

	if idStart > idEnd {
		log.Fatalf("IDStart, %d, is larger than IDEnd, %d.")
	} else if idStart != idEnd && idStart <= 0 {
		log.Fatalf("Non-positive IDStart %d.")
	}

	idType, err := parseIDType(idTypeStr)
	if err != nil { err.Error() }
	if idType != Rockstar && snap == -1 {
		log.Fatalf("Must set the Snap flag if using a non-default IDType.")
	}

	snapNum, err := snapNum()
	if err != nil {
		log.Fatalf(
			"Error encountered when finding rockstar directory: %s",err.Error(),
		)
	} else if (snap < 1 && idType != Rockstar) || snap > snapNum {
		log.Fatalf("Snap %d is out of bounds for %d snaps.", snap, snapNum)
	}

	// Get IDs and snapshots

	rawIds := getIDs(idStart, idEnd, flag.Args())
	if len(rawIds) == 0 { return }

	var ids, snaps []int
	switch idType {
	case Rockstar:
		if snap != -1 {
			snaps = make([]int, len(rawIds))
			for i := range snaps { snaps[i] = snap }
		} else {
			snaps, err = findSnaps(rawIds)
		}
		ids = rawIds
		if err != nil { log.Fatalf(err.Error()) }
	case M200m:
		snaps = make([]int, len(rawIds))
		for i := range snaps { snaps[i] = snap }
		ids, err = convertSortedIDs(rawIds, snap)
		if err != nil { err.Error() }
	default:
		log.Fatal("Unsupported IDType for now. Sorry :3")
	}

	// Tag subhalos, if neccessary.
	var isSub []bool
	if allowSubhalos {
		isSub = make([]bool, len(ids))
	} else {
		isSub, err = findSubs(rawIds, snaps)
		if err != nil { log.Fatal(err.Error()) }
	}

	// Output
	printIds(ids, snaps, isSub)
}

func parseIDType(str string) (IDType, error) {
	switch strings.ToLower(str) {
	case "rockstar": return Rockstar, nil
	case "m500c": return M500c, nil
	case "m200c": return M200c, nil
	case "m200m": return M200m, nil
	}
	return -1, fmt.Errorf("IDType '%s' unrecognized", str)
}

func getIDs(idStart, idEnd int, args []string) []int {
	ids := make([]int, 0, idEnd - idStart + len(args))
	for i := idStart; i < idEnd; i++ { ids = append(ids, i) }
	for _, str := range args {
		i, err := strconv.Atoi(str)
		if err != nil {
			log.Fatalf("Could not parse arg %d: '%s' is not an int.", i+1, str)
		}
		ids = append(ids, i)
	}
	return ids
}

func snapNum() (int, error) {
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	infos, err := ioutil.ReadDir(rockstarDir)
	return len(infos), err
}

func getSnapHaloList(i int) (name string, err error) {
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	infos, err := ioutil.ReadDir(rockstarDir)
	if err != nil { return "", err }
	return path.Join(rockstarDir, infos[i - 1].Name()), nil
}

func findSnaps(ids []int) ([]int, error) {
	treeDir := os.Getenv("GTET_TREE_DIR")
	infos, err := ioutil.ReadDir(treeDir)
	if err != nil { return nil, err }

	names := []string{}
	for _, info := range infos {
		name := info.Name()
		n := len(name)
		if n > 4 && name[:5] == "tree_" && name[n-4:] == ".dat" {
			names = append(names, path.Join(treeDir, name))
		}
	}

	return tree.HaloSnaps(names, ids)
}

// This funciton has side effects :(
func convertSortedIDs(
	rawIDs []int, snap int,
) ([]int, error) {
	list, err := getSnapHaloList(snap)
	if err != nil { return nil, err }

	gtetFmt := os.Getenv("GTET_FMT")
	gtetName := fmt.Sprintf(gtetFmt, snap, 0, 0, 0)
	hd := &io.SheetHeader{}
	err = io.ReadSheetHeaderAt(gtetName, hd)
	if err != nil { return nil, err }
	cosmo := &hd.Cosmo

	rids, xs, ys, zs, ms, rs, err = halo.ReadRockstar(list, halo.R200m, cosmo)
	if err != nil { return nil, err }

	ids := make([]int, len(rawIDs))
	for i := range ids { ids[i] = rids[rawIDs[i]] }
	return ids, nil
}

func findSubs(rawIDs, snaps []int) ([]bool, error) {
	isSub := make([]bool, len(rawIDs))

	// Handle the case where we're indexing by sorted mass ID.
	if xs != nil {
		gtetFmt := os.Getenv("GTET_FMT")
		gtetName := fmt.Sprintf(gtetFmt, snaps[0], 0, 0, 0)
		hd := &io.SheetHeader{}
		err := io.ReadSheetHeaderAt(gtetName, hd)
		if err != nil { return nil, err }

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, zs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, overlapMult)

		for i, id := range rawIDs {
			isSub[i] = sf.HostCount(id) > 0
		}
		return isSub, nil
	}

	// Handle the case where all we have are halo IDs

	// Group by snapshot.
	snapGroups := make(map[int][]int)
	groupIdxs := make(map[int][]int)
	for i, id := range rawIDs {
		snap := snaps[i]
		snapGroups[snap] = append(snapGroups[snap], id)
		groupIdxs[snap] = append(groupIdxs[snap], i)
	}

	// Load each snapshot.
	for snap, group := range snapGroups {
		gtetFmt := os.Getenv("GTET_FMT")
		gtetName := fmt.Sprintf(gtetFmt, snap, 0, 0, 0)
		hd := &io.SheetHeader{}
		err := io.ReadSheetHeaderAt(gtetName, hd)
		if err != nil { return nil, err }

		list, err := getSnapHaloList(snap)
		rids, xs, ys, zs, ms, rs, err = halo.ReadRockstar(
			list, halo.R200m, &hd.Cosmo,
		)

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, zs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, overlapMult)

		for i, id := range group {
			origIdx := groupIdxs[snap][i]
			for j, checkID := range rids {
				if checkID == id {
					isSub[origIdx] = sf.HostCount(j) > 0
					break
				} else if j == len(rids) - 1 {
					return nil, fmt.Errorf("ID %d not in halo list.", id)
				}
			}
		}
	}
	return isSub, nil
}

func printIds(ids []int, snaps []int, isSub []bool) {
	// Find the maximum width of each column.
	idWidth, snapWidth := 0, 0
	for i := range ids {
		if isSub[i] { continue }
		iWidth := len(fmt.Sprintf("%d", ids[i]))
		sWidth := len(fmt.Sprintf("%d", snaps[i]))
		if iWidth > idWidth { idWidth = iWidth }
		if sWidth > snapWidth { snapWidth = sWidth }
	}

	rowFmt := fmt.Sprintf("%%%dd %%%dd\n", idWidth, snapWidth)
	for i := range ids {
		if isSub[i] { continue }
		fmt.Printf(rowFmt, ids[i], snaps[i])
	}
}
<|endoftext|>


<commit_before>// Utility package for sending json responses.
// Able to send to client any JSON serializable object. Object that is sent
// as response can be transformed via Transfrorm function which can be modified
// via SetTransformer function. Default implementation just wraps provided value
// in map with key "data" so it is pretty much useless, but wrapping is done
// because there is no restriction on type of data in response (just string or
// int is not valid JSON).
//
// Example of usage:
//
//     func someHandler(w http.ResponseWriter, r *Request) {
//         // create some object to return
//         obj := ...
//	       jsonResponse.New(obj).OK(w)
//     }
//
// Example of usage with custom header:
//
//     func someHandler(w http.ResponseWriter, r *Request) {
//         // create some object to return
//         obj := ...
//	       jsonResponse.New(obj).Header("X-Custom-Header", "this is cool").OK(w)
//     }

package jsonresponse

import (
	"encoding/json"
	"net/http"
	"fmt"
	"bytes"
)

// Response object, only contains object to return.
type Response struct {
	// Object to return, should be json serializable, or serialization will panic.
	Data    interface{}
	Headers map[string]string
	Excuse  string
}

func serializeToString(data interface{}) (s string ) {
	b, err := json.Marshal(data)
	if err != nil {
		panic(err)
	}
	if indent {
		var out bytes.Buffer
		json.Indent(&out, b, "", "\t")
		return out.String()
	} else {
		return string(b)
	}
}
// New creates response object with provided data and returns it.
func New(data interface{}) (r Response) {
	return Response{Data: data, Headers: map[string]string{}}
}

// Empty creates response object with not data. This can be useful since some
// http responses does not require data as response, only status code,
// like 204 (No Content)
func Empty() (r Response) {
	return Response{Data: nil, Headers: map[string]string{}}
}

// Response transforms body, sets headers and writes encoded body (to JSON) to
// provided writer.
func (r Response) Response(w http.ResponseWriter, httpCode int) {
	var headers map[string]string
	var body interface{}
	if transformer != nil && r.Data != nil {
		headers, body = transformer(r, httpCode)
	} else {
		headers = map[string]string{}
		body = r.Data
	}

	// if we have headers for this response, include it (and override transformer headers)
	for k, v := range r.Headers {
		headers[k] = v
	}

	// set headers to response writer
	responseHeaders := w.Header()
	for k, v := range headers {
		responseHeaders.Set(k, v)
	}

	// if Content-Type is not already included - add it here
	if _, ok := headers["Content-Type"]; !ok {
		responseHeaders.Set("Content-Type", defaultContentTypeHeader)
	}

	// write headers
	w.WriteHeader(httpCode)

	if body != nil {
		fmt.Fprint(w, serializeToString(body) + "\n")
	}
}

// Header adds header to response.
func (r Response) Header(key, value string) Response {
	r.Headers[key] = value
	return r
}

// WithProgrammingExcuse adds random programming excuse. This works with default
// transformer and Excuse field of Response has to be added to custom transfomer
// if it is used.
func (r Response) WithProgrammingExcuse() Response {
	r.Excuse = randomExcuse()
	return r
}

// 1xx

// Continue sends response to client with HTTP status 100.
// This means that server has received the request headers and that the client
// should proceed to send the request body
func (r Response) Continue(w http.ResponseWriter) {
	r.Response(w, http.StatusContinue)
}

// SwitchingProtocols sends response to client with HTTP status 101.
// This means the requester has asked the server to switch protocols and the
// server is acknowledging that it will do so.
func (r Response) SwitchingProtocols(w http.ResponseWriter) {
	r.Response(w, http.StatusSwitchingProtocols)
}

// 2xx

// OK sends response to client with HTTP status 200.
// Standard response for successful HTTP requests.
func (r Response) OK(w http.ResponseWriter) {
	r.Response(w, http.StatusOK)
}

// Created sends response to client with HTTP status 201.
// The request has been fulfilled and resulted in a new resource being created.
func (r Response) Created(w http.ResponseWriter) {
	r.Response(w, http.StatusCreated)
}

// Accepted sends response to client with HTTP status 202.
// The request has been accepted for processing, but the processing has not
// been completed.
func (r Response) Accepted(w http.ResponseWriter) {
	r.Response(w, http.StatusAccepted)
}

// NonAuthoritativeInfo sends response to client with HTTP status 203.
// The server successfully processed the request, but is returning information
// that may be from another source.
func (r Response) NonAuthoritativeInfo(w http.ResponseWriter) {
	r.Response(w, http.StatusNonAuthoritativeInfo)
}

// NoContent sends response to client with HTTP status 204.
// The server successfully processed the request, but is not returning any content.
func (r Response) NoContent(w http.ResponseWriter) {
	r.Response(w, http.StatusNoContent)
}

// ResetContent sends response to client with HTTP status 205.
// The server successfully processed the request, but is not returning any content.
// Unlike a NoContent response, this response requires that the requester reset
// the document view.
func (r Response) ResetContent(w http.ResponseWriter) {
	r.Response(w, http.StatusResetContent)
}

// PartialContent sends response to client with HTTP status 206.
// The server is delivering only part of the resource (byte serving) due to a
// range header sent by the client.
func (r Response) PartialContent(w http.ResponseWriter) {
	r.Response(w, http.StatusPartialContent)
}

// 3xx

// MultipleChoices sends response to client with HTTP status 300.
// Indicates multiple options for the resource that the client may follow.
func (r Response) MultipleChoices(w http.ResponseWriter) {
	r.Response(w, http.StatusMultipleChoices)
}

// MovedPermanently sends response to client with HTTP status 301.
// This and all future requests should be directed to the given URI.
func (r Response) MovedPermanently(w http.ResponseWriter) {
	r.Response(w, http.StatusMovedPermanently)
}

// Found sends response to client with HTTP status 302.
func (r Response) Found(w http.ResponseWriter) {
	r.Response(w, http.StatusFound)
}

// SeeOther sends response to client with HTTP status 303.
// The response to the request can be found under another URI using a GET method.
func (r Response) SeeOther(w http.ResponseWriter) {
	r.Response(w, http.StatusSeeOther)
}

// NotModified sends response to client with HTTP status 304.
// Indicates that the resource has not been modified since the version specified
// by the request headers If-Modified-Since or If-None-Match.
func (r Response) NotModified(w http.ResponseWriter) {
	r.Response(w, http.StatusNotModified)
}

// UseProxy sends response to client with HTTP status 305.
// The requested resource is only available through a proxy, whose address is
// provided in the response.
func (r Response) UseProxy(w http.ResponseWriter) {
	r.Response(w, http.StatusUseProxy)
}

// TemporaryRedirect sends response to client with HTTP status 307.
// In this case, the request should be repeated with another URI; however,
// future requests should still use the original URI.
func (r Response) TemporaryRedirect(w http.ResponseWriter) {
	r.Response(w, http.StatusTemporaryRedirect)
}

// 4xx

// BadRequest sends response to client with HTTP status 400.
// The server cannot or will not process the request due to something that is
// perceived to be a client error
func (r Response) BadRequest(w http.ResponseWriter) {
	r.Response(w, http.StatusBadRequest)
}

// Unauthorized sends response to client with HTTP status 401.
// Similar to 403 Forbidden, but specifically for use when authentication
// is required and has failed or has not yet been provided.
func (r Response) Unauthorized(w http.ResponseWriter) {
	r.Response(w, http.StatusUnauthorized)
}

// PaymentRequired sends response to client with HTTP status 402.
// Reserved for future use.
func (r Response) PaymentRequired(w http.ResponseWriter) {
	r.Response(w, http.StatusPaymentRequired)
}

// Forbidden sends response to client with HTTP status 403.
// The request was a valid request, but the server is refusing to respond to it.
// Unlike a 401 Unauthorized response, authenticating will make no difference.
func (r Response) Forbidden(w http.ResponseWriter) {
	r.Response(w, http.StatusForbidden)
}

// NotFound sends response to client with HTTP status 404.
// The requested resource could not be found but may be available again in the future.
func (r Response) NotFound(w http.ResponseWriter) {
	r.Response(w, http.StatusNotFound)
}

// MethodNotAllowed sends response to client with HTTP status 405.
// A request was made of a resource using a request method not supported by that
// resource; for example, using GET on a form which requires data to be presented
// via POST, or using PUT on a read-only resource.
func (r Response) MethodNotAllowed(w http.ResponseWriter) {
	r.Response(w, http.StatusMethodNotAllowed)
}

// NotAcceptable sends response to client with HTTP status 406.
// The requested resource is only capable of generating content not acceptable
// according to the Accept headers sent in the request.
func (r Response) NotAcceptable(w http.ResponseWriter) {
	r.Response(w, http.StatusNotAcceptable)
}

// ProxyAuthRequired sends response to client with HTTP status 407.
// The client must first authenticate itself with the proxy.
func (r Response) ProxyAuthRequired(w http.ResponseWriter) {
	r.Response(w, http.StatusProxyAuthRequired)
}

// RequestTimeout sends response to client with HTTP status 408.
// The server timed out waiting for the request.
func (r Response) RequestTimeout(w http.ResponseWriter) {
	r.Response(w, http.StatusRequestTimeout)
}

// Conflict sends response to client with HTTP status 409.
// Indicates that the request could not be processed because of conflict
// in the request.
func (r Response) Conflict(w http.ResponseWriter) {
	r.Response(w, http.StatusConflict)
}

// Gone sends response to client with HTTP status 410.
// Indicates that the resource requested is no longer available and will not
// be available again.
func (r Response) Gone(w http.ResponseWriter) {
	r.Response(w, http.StatusGone)
}

// LengthRequired sends response to client with HTTP status 411.
// The request did not specify the length of its content, which is
// required by the requested resource.
func (r Response) LengthRequired(w http.ResponseWriter) {
	r.Response(w, http.StatusLengthRequired)
}

// PreconditionFailed sends response to client with HTTP status 412.
// The server does not meet one of the preconditions that the requester put
// on the request.
func (r Response) PreconditionFailed(w http.ResponseWriter) {
	r.Response(w, http.StatusPreconditionFailed)
}

// RequestEntityTooLarge sends response to client with HTTP status 413.
// The request is larger than the server is willing or able to process.
func (r Response) RequestEntityTooLarge(w http.ResponseWriter) {
	r.Response(w, http.StatusRequestEntityTooLarge)
}

// RequestURITooLong sends response to client with HTTP status 414.
// The URI provided was too long for the server to process.
func (r Response) RequestURITooLong(w http.ResponseWriter) {
	r.Response(w, http.StatusRequestURITooLong)
}

// UnsupportedMediaType sends response to client with HTTP status 415.
// The request entity has a media type which the server or resource does
// not support.
func (r Response) UnsupportedMediaType(w http.ResponseWriter) {
	r.Response(w, http.StatusUnsupportedMediaType)
}

// RequestedRangeNotSatisfiable sends response to client with HTTP status 416.
// The client has asked for a portion of the file (byte serving), but the
// server cannot supply that portion.
func (r Response) RequestedRangeNotSatisfiable(w http.ResponseWriter) {
	r.Response(w, http.StatusRequestedRangeNotSatisfiable)
}

// ExpectationFailed sends response to client with HTTP status 417.
// The server cannot meet the requirements of the Expect request-header field.
func (r Response) ExpectationFailed(w http.ResponseWriter) {
	r.Response(w, http.StatusExpectationFailed)
}

// Teapot sends response to client with HTTP status 418.
// This code should be returned by tea pots requested to brew coffee.
func (r Response) Teapot(w http.ResponseWriter) {
	r.Response(w, http.StatusTeapot)
}

// 5xx

// InternalServerError sends response to client with HTTP status 500.
// A generic error message, given when an unexpected condition was
// encountered and no more specific message is suitable.
func (r Response) InternalServerError(w http.ResponseWriter) {
	r.Response(w, http.StatusInternalServerError)
}

// NotImplemented sends response to client with HTTP status 501.
// The server either does not recognize the request method, or it lacks the
// ability to fulfill the request.
func (r Response) NotImplemented(w http.ResponseWriter) {
	r.Response(w, http.StatusNotImplemented)
}

// BadGateway sends response to client with HTTP status 502.
// The server was acting as a gateway or proxy and received an invalid
// response from the upstream server.
func (r Response) BadGateway(w http.ResponseWriter) {
	r.Response(w, http.StatusBadGateway)
}

// ServiceUnavailable sends response to client with HTTP status 503.
// The server is currently unavailable (because it is overloaded or down
// for maintenance). Generally, this is a temporary state.
func (r Response) ServiceUnavailable(w http.ResponseWriter) {
	r.Response(w, http.StatusServiceUnavailable)
}

// GatewayTimeout sends response to client with HTTP status 504.
// The server was acting as a gateway or proxy and did not receive a timely
// response from the upstream server.
func (r Response) GatewayTimeout(w http.ResponseWriter) {
	r.Response(w, http.StatusGatewayTimeout)
}

// HTTPVersionNotSupported sends response to client with HTTP status 505.
// The server does not support the HTTP protocol version used in the request.
func (r Response) HTTPVersionNotSupported(w http.ResponseWriter) {
	r.Response(w, http.StatusHTTPVersionNotSupported)
}
<commit_msg>Do not indent in separate call<commit_after>// Utility package for sending json responses.
// Able to send to client any JSON serializable object. Object that is sent
// as response can be transformed via Transfrorm function which can be modified
// via SetTransformer function. Default implementation just wraps provided value
// in map with key "data" so it is pretty much useless, but wrapping is done
// because there is no restriction on type of data in response (just string or
// int is not valid JSON).
//
// Example of usage:
//
//     func someHandler(w http.ResponseWriter, r *Request) {
//         // create some object to return
//         obj := ...
//	       jsonResponse.New(obj).OK(w)
//     }
//
// Example of usage with custom header:
//
//     func someHandler(w http.ResponseWriter, r *Request) {
//         // create some object to return
//         obj := ...
//	       jsonResponse.New(obj).Header("X-Custom-Header", "this is cool").OK(w)
//     }

package jsonresponse

import (
	"encoding/json"
	"net/http"
	"fmt"
)

// Response object, only contains object to return.
type Response struct {
	// Object to return, should be json serializable, or serialization will panic.
	Data    interface{}
	Headers map[string]string
	Excuse  string
}

func serializeToString(data interface{}) (s string ) {
	var b []byte
	var err error
	if indent {
		b, err = json.MarshalIndent(data, "", "\t")
	} else {
		b, err = json.Marshal(data)
	}
	if err != nil {
		panic(err)
	}
	return string(b)
}

// New creates response object with provided data and returns it.
func New(data interface{}) (r Response) {
	return Response{Data: data, Headers: map[string]string{}}
}

// Empty creates response object with not data. This can be useful since some
// http responses does not require data as response, only status code,
// like 204 (No Content)
func Empty() (r Response) {
	return Response{Data: nil, Headers: map[string]string{}}
}

// Response transforms body, sets headers and writes encoded body (to JSON) to
// provided writer.
func (r Response) Response(w http.ResponseWriter, httpCode int) {
	var headers map[string]string
	var body interface{}
	if transformer != nil && r.Data != nil {
		headers, body = transformer(r, httpCode)
	} else {
		headers = map[string]string{}
		body = r.Data
	}

	// if we have headers for this response, include it (and override transformer headers)
	for k, v := range r.Headers {
		headers[k] = v
	}

	// set headers to response writer
	responseHeaders := w.Header()
	for k, v := range headers {
		responseHeaders.Set(k, v)
	}

	// if Content-Type is not already included - add it here
	if _, ok := headers["Content-Type"]; !ok {
		responseHeaders.Set("Content-Type", defaultContentTypeHeader)
	}

	// write headers
	w.WriteHeader(httpCode)

	if body != nil {
		fmt.Fprint(w, serializeToString(body) + "\n")
	}
}

// Header adds header to response.
func (r Response) Header(key, value string) Response {
	r.Headers[key] = value
	return r
}

// WithProgrammingExcuse adds random programming excuse. This works with default
// transformer and Excuse field of Response has to be added to custom transfomer
// if it is used.
func (r Response) WithProgrammingExcuse() Response {
	r.Excuse = randomExcuse()
	return r
}

// 1xx

// Continue sends response to client with HTTP status 100.
// This means that server has received the request headers and that the client
// should proceed to send the request body
func (r Response) Continue(w http.ResponseWriter) {
	r.Response(w, http.StatusContinue)
}

// SwitchingProtocols sends response to client with HTTP status 101.
// This means the requester has asked the server to switch protocols and the
// server is acknowledging that it will do so.
func (r Response) SwitchingProtocols(w http.ResponseWriter) {
	r.Response(w, http.StatusSwitchingProtocols)
}

// 2xx

// OK sends response to client with HTTP status 200.
// Standard response for successful HTTP requests.
func (r Response) OK(w http.ResponseWriter) {
	r.Response(w, http.StatusOK)
}

// Created sends response to client with HTTP status 201.
// The request has been fulfilled and resulted in a new resource being created.
func (r Response) Created(w http.ResponseWriter) {
	r.Response(w, http.StatusCreated)
}

// Accepted sends response to client with HTTP status 202.
// The request has been accepted for processing, but the processing has not
// been completed.
func (r Response) Accepted(w http.ResponseWriter) {
	r.Response(w, http.StatusAccepted)
}

// NonAuthoritativeInfo sends response to client with HTTP status 203.
// The server successfully processed the request, but is returning information
// that may be from another source.
func (r Response) NonAuthoritativeInfo(w http.ResponseWriter) {
	r.Response(w, http.StatusNonAuthoritativeInfo)
}

// NoContent sends response to client with HTTP status 204.
// The server successfully processed the request, but is not returning any content.
func (r Response) NoContent(w http.ResponseWriter) {
	r.Response(w, http.StatusNoContent)
}

// ResetContent sends response to client with HTTP status 205.
// The server successfully processed the request, but is not returning any content.
// Unlike a NoContent response, this response requires that the requester reset
// the document view.
func (r Response) ResetContent(w http.ResponseWriter) {
	r.Response(w, http.StatusResetContent)
}

// PartialContent sends response to client with HTTP status 206.
// The server is delivering only part of the resource (byte serving) due to a
// range header sent by the client.
func (r Response) PartialContent(w http.ResponseWriter) {
	r.Response(w, http.StatusPartialContent)
}

// 3xx

// MultipleChoices sends response to client with HTTP status 300.
// Indicates multiple options for the resource that the client may follow.
func (r Response) MultipleChoices(w http.ResponseWriter) {
	r.Response(w, http.StatusMultipleChoices)
}

// MovedPermanently sends response to client with HTTP status 301.
// This and all future requests should be directed to the given URI.
func (r Response) MovedPermanently(w http.ResponseWriter) {
	r.Response(w, http.StatusMovedPermanently)
}

// Found sends response to client with HTTP status 302.
func (r Response) Found(w http.ResponseWriter) {
	r.Response(w, http.StatusFound)
}

// SeeOther sends response to client with HTTP status 303.
// The response to the request can be found under another URI using a GET method.
func (r Response) SeeOther(w http.ResponseWriter) {
	r.Response(w, http.StatusSeeOther)
}

// NotModified sends response to client with HTTP status 304.
// Indicates that the resource has not been modified since the version specified
// by the request headers If-Modified-Since or If-None-Match.
func (r Response) NotModified(w http.ResponseWriter) {
	r.Response(w, http.StatusNotModified)
}

// UseProxy sends response to client with HTTP status 305.
// The requested resource is only available through a proxy, whose address is
// provided in the response.
func (r Response) UseProxy(w http.ResponseWriter) {
	r.Response(w, http.StatusUseProxy)
}

// TemporaryRedirect sends response to client with HTTP status 307.
// In this case, the request should be repeated with another URI; however,
// future requests should still use the original URI.
func (r Response) TemporaryRedirect(w http.ResponseWriter) {
	r.Response(w, http.StatusTemporaryRedirect)
}

// 4xx

// BadRequest sends response to client with HTTP status 400.
// The server cannot or will not process the request due to something that is
// perceived to be a client error
func (r Response) BadRequest(w http.ResponseWriter) {
	r.Response(w, http.StatusBadRequest)
}

// Unauthorized sends response to client with HTTP status 401.
// Similar to 403 Forbidden, but specifically for use when authentication
// is required and has failed or has not yet been provided.
func (r Response) Unauthorized(w http.ResponseWriter) {
	r.Response(w, http.StatusUnauthorized)
}

// PaymentRequired sends response to client with HTTP status 402.
// Reserved for future use.
func (r Response) PaymentRequired(w http.ResponseWriter) {
	r.Response(w, http.StatusPaymentRequired)
}

// Forbidden sends response to client with HTTP status 403.
// The request was a valid request, but the server is refusing to respond to it.
// Unlike a 401 Unauthorized response, authenticating will make no difference.
func (r Response) Forbidden(w http.ResponseWriter) {
	r.Response(w, http.StatusForbidden)
}

// NotFound sends response to client with HTTP status 404.
// The requested resource could not be found but may be available again in the future.
func (r Response) NotFound(w http.ResponseWriter) {
	r.Response(w, http.StatusNotFound)
}

// MethodNotAllowed sends response to client with HTTP status 405.
// A request was made of a resource using a request method not supported by that
// resource; for example, using GET on a form which requires data to be presented
// via POST, or using PUT on a read-only resource.
func (r Response) MethodNotAllowed(w http.ResponseWriter) {
	r.Response(w, http.StatusMethodNotAllowed)
}

// NotAcceptable sends response to client with HTTP status 406.
// The requested resource is only capable of generating content not acceptable
// according to the Accept headers sent in the request.
func (r Response) NotAcceptable(w http.ResponseWriter) {
	r.Response(w, http.StatusNotAcceptable)
}

// ProxyAuthRequired sends response to client with HTTP status 407.
// The client must first authenticate itself with the proxy.
func (r Response) ProxyAuthRequired(w http.ResponseWriter) {
	r.Response(w, http.StatusProxyAuthRequired)
}

// RequestTimeout sends response to client with HTTP status 408.
// The server timed out waiting for the request.
func (r Response) RequestTimeout(w http.ResponseWriter) {
	r.Response(w, http.StatusRequestTimeout)
}

// Conflict sends response to client with HTTP status 409.
// Indicates that the request could not be processed because of conflict
// in the request.
func (r Response) Conflict(w http.ResponseWriter) {
	r.Response(w, http.StatusConflict)
}

// Gone sends response to client with HTTP status 410.
// Indicates that the resource requested is no longer available and will not
// be available again.
func (r Response) Gone(w http.ResponseWriter) {
	r.Response(w, http.StatusGone)
}

// LengthRequired sends response to client with HTTP status 411.
// The request did not specify the length of its content, which is
// required by the requested resource.
func (r Response) LengthRequired(w http.ResponseWriter) {
	r.Response(w, http.StatusLengthRequired)
}

// PreconditionFailed sends response to client with HTTP status 412.
// The server does not meet one of the preconditions that the requester put
// on the request.
func (r Response) PreconditionFailed(w http.ResponseWriter) {
	r.Response(w, http.StatusPreconditionFailed)
}

// RequestEntityTooLarge sends response to client with HTTP status 413.
// The request is larger than the server is willing or able to process.
func (r Response) RequestEntityTooLarge(w http.ResponseWriter) {
	r.Response(w, http.StatusRequestEntityTooLarge)
}

// RequestURITooLong sends response to client with HTTP status 414.
// The URI provided was too long for the server to process.
func (r Response) RequestURITooLong(w http.ResponseWriter) {
	r.Response(w, http.StatusRequestURITooLong)
}

// UnsupportedMediaType sends response to client with HTTP status 415.
// The request entity has a media type which the server or resource does
// not support.
func (r Response) UnsupportedMediaType(w http.ResponseWriter) {
	r.Response(w, http.StatusUnsupportedMediaType)
}

// RequestedRangeNotSatisfiable sends response to client with HTTP status 416.
// The client has asked for a portion of the file (byte serving), but the
// server cannot supply that portion.
func (r Response) RequestedRangeNotSatisfiable(w http.ResponseWriter) {
	r.Response(w, http.StatusRequestedRangeNotSatisfiable)
}

// ExpectationFailed sends response to client with HTTP status 417.
// The server cannot meet the requirements of the Expect request-header field.
func (r Response) ExpectationFailed(w http.ResponseWriter) {
	r.Response(w, http.StatusExpectationFailed)
}

// Teapot sends response to client with HTTP status 418.
// This code should be returned by tea pots requested to brew coffee.
func (r Response) Teapot(w http.ResponseWriter) {
	r.Response(w, http.StatusTeapot)
}

// 5xx

// InternalServerError sends response to client with HTTP status 500.
// A generic error message, given when an unexpected condition was
// encountered and no more specific message is suitable.
func (r Response) InternalServerError(w http.ResponseWriter) {
	r.Response(w, http.StatusInternalServerError)
}

// NotImplemented sends response to client with HTTP status 501.
// The server either does not recognize the request method, or it lacks the
// ability to fulfill the request.
func (r Response) NotImplemented(w http.ResponseWriter) {
	r.Response(w, http.StatusNotImplemented)
}

// BadGateway sends response to client with HTTP status 502.
// The server was acting as a gateway or proxy and received an invalid
// response from the upstream server.
func (r Response) BadGateway(w http.ResponseWriter) {
	r.Response(w, http.StatusBadGateway)
}

// ServiceUnavailable sends response to client with HTTP status 503.
// The server is currently unavailable (because it is overloaded or down
// for maintenance). Generally, this is a temporary state.
func (r Response) ServiceUnavailable(w http.ResponseWriter) {
	r.Response(w, http.StatusServiceUnavailable)
}

// GatewayTimeout sends response to client with HTTP status 504.
// The server was acting as a gateway or proxy and did not receive a timely
// response from the upstream server.
func (r Response) GatewayTimeout(w http.ResponseWriter) {
	r.Response(w, http.StatusGatewayTimeout)
}

// HTTPVersionNotSupported sends response to client with HTTP status 505.
// The server does not support the HTTP protocol version used in the request.
func (r Response) HTTPVersionNotSupported(w http.ResponseWriter) {
	r.Response(w, http.StatusHTTPVersionNotSupported)
}
<|endoftext|>


<commit_before>// Package replicat is a server for n way synchronization of content (Replication for the cloud).
// Copyright 2016 Jacob Taylor jacob@replic.at       More Info: http://replic.at
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/goji/httpauth"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"sync"
	"time"
)

// ReplicatServer is a structure that contains the definition of the servers in a cluster. Each node has a name and this
// node (as determined by globalSettings.name at the moment) also has a StorageTracker interface.
type ReplicatServer struct {
	ClusterKey    string
	Name          string
	Address       string
	Status        string
	CurrentState  DirTreeMap
	PreviousState DirTreeMap
	Lock          sync.Mutex
	storage       StorageTracker
}

// GetStatus - get the current status of the server
func (server *ReplicatServer) GetStatus() string {
	return server.Status
}

// SetStatus - set the current status of the server
func (server *ReplicatServer) SetStatus(status string) {
	// adding an inline check to make sure the server is valid
	// if the server is being shut down or at the end of a unit test)
	// the assignment can crash without this check
	if server != nil {
		server.Status = status
		sendConfigToServer()
	}
}

const (
	// REPLICAT_STATUS_INITIAL_SCAN - The server is scanning its local storage
	REPLICAT_STATUS_INITIAL_SCAN = "Initial Scan"
	// REPLICAT_STATUS_JOINING_CLUSTER - The server has started to join the cluster and get up to date
	REPLICAT_STATUS_JOINING_CLUSTER = "Joining Cluster"
	// REPLICAT_STATUS_ONLINE - The server is up to date and part of the cluster
	REPLICAT_STATUS_ONLINE = "Online"
)

var serverMap = make(map[string]*ReplicatServer)
var serverMapLock sync.RWMutex

var lastConfigPing = time.Now()

// BootstrapAndServe - Start the server
func BootstrapAndServe(address string) {
	//trackerTestDual()
	//trackerTestSmallFileInSubfolder()
	//trackerTestEmptyDirectoryMovesInOutAround()
	//trackerTestFileChangeTrackerAddFolders()
	//trackerTestSmallFileCreationAndRename()
	//trackerTestSmallFileCreationAndUpdate()
	//trackerTestSmallFileMovesInOutAround()
	//trackerTestDirectoryCreation()
	//trackerTestNestedDirectoryCreation()
	//trackerTestDirectoryStorage()
	//trackerTestFileChangeTrackerAutoCreateFolderAndCleanup()
	//trackerTestNestedFastDirectoryCreation()

	// testing code to enable debugger use
	http.Handle("/event/", httpauth.SimpleBasicAuth("replicat", "isthecat")(http.HandlerFunc(eventHandler)))
	http.Handle("/tree/", httpauth.SimpleBasicAuth("replicat", "isthecat")(http.HandlerFunc(folderTreeHandler)))
	http.Handle("/config/", httpauth.SimpleBasicAuth("replicat", "isthecat")(http.HandlerFunc(configHandler)))
	http.Handle("/upload/", httpauth.SimpleBasicAuth("replicat", "isthecat")(http.HandlerFunc(uploadHandler)))

	// quick test of Minio
	suffix := rand.Int31n(10000)
	tempFolder := fmt.Sprintf("replicat-test-%05d", suffix)
	fmt.Printf("temporary path name: %s\n", tempFolder)
	tracker2 := &minioTracker{}
	tracker2.Initialize()
	folderList, err := tracker2.ListFolders(false)
	for _, v := range folderList {
		fmt.Printf("bucket: %s\n", v)
	}

	objectName := "babySloth"

	err = tracker2.CreatePath(tempFolder, true)
	if err != nil {
		panic(err)
	}

	folderlist, err := tracker2.ListFolders(false)

	found := false
	for _, v := range folderlist {
		//fmt.Printf("Comparing %s and %s\n", v, tempFolder)
		if v == tempFolder {
			found = true
			break
		}
	}
	if found == false {
		fmt.Printf("Could not find the directory we should have just created: %s\n", tempFolder)
	}

	// Time to create file!
	err = tracker2.CreateObject(tempFolder, objectName, "/testdata/LargeSampleImages/baby sloth on a rail.jpg", "")
	if err != nil {
		panic(err)
	}

	err = tracker2.DeleteObject(tempFolder, objectName)
	if err != nil {
		panic(err)
	}

	err = tracker2.DeleteFolder(tempFolder)
	if err != nil {
		panic(err)
	}

	folderlist, err = tracker2.ListFolders(false)

	found = false
	for _, v := range folderlist {
		if v == tempFolder {
			found = true
			break
		}
	}

	if found == true {
		fmt.Printf("We found our folder after it should have been deleted: %s\n", tempFolder)
	}

	fmt.Println("YAY - Made it to the end")

	if tracker2 != nil {
		panic("nobody knows the trouble I've seen?\n")
	}

	lsnr, err := net.Listen("tcp4", address)
	if err != nil {
		panic(fmt.Sprintf("Error listening: %v\nAddress: %s\n", err, address))
	}
	fmt.Println("Listening on:", lsnr.Addr().String())

	logOnlyHandler := LogOnlyChangeHandler{}
	tracker := FilesystemTracker{}
	fmt.Printf("Looking up settings for node: %s\n", globalSettings.Name)

	directory := globalSettings.Directory

	fmt.Printf("GlobalSettings directory retrieved for this node: %s\n", directory)
	server := &ReplicatServer{Name: globalSettings.Name, ClusterKey: globalSettings.ClusterKey, Address: lsnr.Addr().String(), storage: &tracker, Status: REPLICAT_STATUS_INITIAL_SCAN}
	serverMap[globalSettings.Name] = server
	tracker.init(directory, server)

	go func(tracker StorageTracker) {
		for true {
			tracker.GetStatistics()
			time.Sleep(30 * time.Second)
		}
	}(&tracker)

	var c ChangeHandler
	c = &logOnlyHandler
	tracker.watchDirectory(&c)

	go func(listener net.Listener) {
		err = http.Serve(listener, nil)
		if err != nil {
			panic(err)
		}
	}(lsnr)

	fmt.Println("Starting config update processor")
	go configUpdateProcessor(configUpdateChannel)

	if globalSettings.ManagerAddress != "" {
		fmt.Printf("about to send config to server (%s)\nOur address is: (%s)\n", globalSettings.ManagerAddress, lsnr.Addr())
	}

	// we need something to keep us running.
	keepConfigCurrent()
}

const (
	managerOverdueSeconds = 30
	managerCheckSleepTime = 45
)

func keepConfigCurrent() {
	for {
		if time.Since(lastConfigPing) > managerOverdueSeconds*time.Second {
			log.Printf("Manager Contact Overdue, attempting to contact: %s\n", globalSettings.ManagerAddress)
			sendConfigToServer()
		} else {
			fmt.Println("No Update Required")
		}
		time.Sleep(time.Duration(rand.Intn(10)-5+managerCheckSleepTime) * time.Second)
	}
}

func sendConfigToServer() {
	// This field will be empty during testing. We have to abort during testing since the manager is not reachable.
	if globalSettings.ManagerAddress == "" {
		return
	}

	url := "https://" + globalSettings.ManagerAddress + "/config/"
	log.Printf("sendConfigToServer: Manager location: %s\n", url)

	server := serverMap[globalSettings.Name]
	jsonStr, _ := json.Marshal(server)
	jsonStr2, _ := json.Marshal(server.storage.GetStatistics())

	byteArraysToMarshal := [][]byte{jsonStr, jsonStr2}
	byteData := bytes.Join(byteArraysToMarshal, []byte{})
	log.Printf("jasonstr: %s\n", string(byteData))
	jsonData := bytes.NewBuffer(byteData)

	req, err := http.NewRequest("POST", url, jsonData)
	req.Header.Set("Content-Type", "application/json")

	data := []byte(globalSettings.ManagerCredentials)
	authHash := base64.StdEncoding.EncodeToString(data)
	req.Header.Add("Authorization", "Basic "+authHash)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("Unable to reach server. Data could be lost: %s\n", err)
		return
	}

	newServerMap, err := extractServerMapFromConfig(resp.Body)
	if err == io.EOF {
		log.Printf("EOF unexpected Config update failed\n")
	} else if err != nil {
		log.Printf("Config update failed due to error: %s\n", err)
	} else {
		configUpdateChannel <- newServerMap
	}

}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		r.ParseMultipartForm(32 << 20)
		file, handler, err := r.FormFile("uploadfile")
		if err != nil {
			fmt.Println(err)
			return
		}
		defer file.Close()
		fmt.Fprint(w, handler.Header)

		hash := []byte(r.Form.Get("HASH"))

		storage := serverMap[globalSettings.Name].storage
		local, _ := storage.getEntryJSON(handler.Filename)

		if !bytes.Equal(hash, local.Hash) {
			fullPath := globalSettings.Directory + "/" + handler.Filename
			f, err := os.OpenFile(fullPath, os.O_WRONLY|os.O_CREATE, 0666)

			if err != nil {
				fmt.Println(err)
				return
			}

			bytesWritten, err := io.Copy(f, file)
			if err != nil {
				log.Printf("Error copying file: %s, error(%#v)\n", handler.Filename, err)
				w.WriteHeader(http.StatusInternalServerError)
				w.Write([]byte("500 - Error copying file"))
				f.Close()
				return
			}
			fmt.Printf("Wrote out (%s) bytes (%d)\n", handler.Filename, bytesWritten)
			f.Close()

			// Get the old entry
			entryString := r.Form.Get("entryJSON")
			if entryString != "" {
				var entry EntryJSON
				err := json.Unmarshal([]byte(entryString), &entry)
				if err != nil {
					log.Fatalf("Error copying file (Entry handling): %s, error(%#v)\n", handler.Filename, err)
					w.WriteHeader(http.StatusInternalServerError)
					w.Write([]byte("500 - Error copying file"))
					return
				}

				err = os.Chtimes(fullPath, time.Now(), entry.ModTime)
				if err != nil {
					log.Fatalf("Error copying file (Changing times): %s, error(%#v)\n", handler.Filename, err)
					w.WriteHeader(http.StatusInternalServerError)
					w.Write([]byte("500 - Error copying file"))
					return
				}
			}

		}
	}
}

var configUpdateChannel = make(chan *map[string]*ReplicatServer, 100)

func configHandler(_ http.ResponseWriter, r *http.Request) {
	newServerMap, err := extractServerMapFromConfig(r.Body)
	if err != nil {
		fmt.Printf("Config update failed due to error: %s\n", err)
		return
	}
	configUpdateChannel <- newServerMap
}

func extractServerMapFromConfig(sourceReader io.Reader) (newServerMap *map[string]*ReplicatServer, err error) {
	decoder := json.NewDecoder(sourceReader)
	err = decoder.Decode(&newServerMap)
	return
}

func configUpdateProcessor(c chan *map[string]*ReplicatServer) {
	for {
		newServerMap := <-c
		serverMapLock.Lock()

		lastConfigPing = time.Now()
		sendData := false

		// find any nodes that have been deleted and update ones that have changed
		for name, serverData := range serverMap {
			newServerData, exists := (*newServerMap)[name]
			if !exists {
				fmt.Printf("No longer found config for: %s deleting\n", name)
				delete(serverMap, name)
				continue
			}

			if serverData.Address != newServerData.Address || serverData.Name != newServerData.Name || serverData.ClusterKey != newServerData.ClusterKey || serverData.Status != serverData.Status {
				fmt.Printf("Server data is changed. Replacing.\nold: %v\nnew: %v\n", &serverData, &newServerData)
				if serverData.Status != REPLICAT_STATUS_JOINING_CLUSTER && newServerData.Status == REPLICAT_STATUS_JOINING_CLUSTER {
					fmt.Printf("Decided to send data to: %s\n", serverData.Name)
					sendData = true
				}
				serverMap[name] = newServerData
				fmt.Println("Server data replaced with new server data")
			}
		}

		// find any new nodes
		for name, newServerData := range *newServerMap {
			_, exists := serverMap[name]
			if !exists {
				fmt.Printf("New server configuration for %s: %v\n", name, newServerData)

				// If this server map is for ourselves, build a list of folder if needed and notify others
				if name == globalSettings.Name {
					listOfFileInfo, err := scanDirectoryContents()
					if err != nil {
						log.Fatal(err)
					}
					newServerData.CurrentState = listOfFileInfo
					// Tell all of our friends that we exist and our current state for them to compare against.
					go func(tree DirTreeMap) {
						sendFolderTree(tree)
					}(listOfFileInfo)
				} else {
					fmt.Printf("New Node Decided to send data to: %s\n", name)
					sendData = true
				}

				fmt.Printf("New server configuration provided. Copying: %s\n", name)
				serverMap[name] = newServerData
			}
		}

		if sendData {
			server := serverMap[globalSettings.Name]
			fmt.Println("about to send existing files")
			server.storage.SendCatalog()
			fmt.Println("done sending existing files")
		}

		serverMapLock.Unlock()
	}
}
<commit_msg>minor cleanup<commit_after>// Package replicat is a server for n way synchronization of content (Replication for the cloud).
// Copyright 2016 Jacob Taylor jacob@replic.at       More Info: http://replic.at
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/goji/httpauth"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"sync"
	"time"
)

// ReplicatServer is a structure that contains the definition of the servers in a cluster. Each node has a name and this
// node (as determined by globalSettings.name at the moment) also has a StorageTracker interface.
type ReplicatServer struct {
	ClusterKey    string
	Name          string
	Address       string
	Status        string
	CurrentState  DirTreeMap
	PreviousState DirTreeMap
	Lock          sync.Mutex
	storage       StorageTracker
}

// GetStatus - get the current status of the server
func (server *ReplicatServer) GetStatus() string {
	return server.Status
}

// SetStatus - set the current status of the server
func (server *ReplicatServer) SetStatus(status string) {
	// adding an inline check to make sure the server is valid
	// if the server is being shut down or at the end of a unit test)
	// the assignment can crash without this check
	if server != nil {
		server.Status = status
		sendConfigToServer()
	}
}

const (
	// REPLICAT_STATUS_INITIAL_SCAN - The server is scanning its local storage
	REPLICAT_STATUS_INITIAL_SCAN = "Initial Scan"
	// REPLICAT_STATUS_JOINING_CLUSTER - The server has started to join the cluster and get up to date
	REPLICAT_STATUS_JOINING_CLUSTER = "Joining Cluster"
	// REPLICAT_STATUS_ONLINE - The server is up to date and part of the cluster
	REPLICAT_STATUS_ONLINE = "Online"
)

var serverMap = make(map[string]*ReplicatServer)
var serverMapLock sync.RWMutex

var lastConfigPing = time.Now()

// BootstrapAndServe - Start the server
func BootstrapAndServe(address string) {
	//trackerTestDual()
	//trackerTestSmallFileInSubfolder()
	//trackerTestEmptyDirectoryMovesInOutAround()
	//trackerTestFileChangeTrackerAddFolders()
	//trackerTestSmallFileCreationAndRename()
	//trackerTestSmallFileCreationAndUpdate()
	//trackerTestSmallFileMovesInOutAround()
	//trackerTestDirectoryCreation()
	//trackerTestNestedDirectoryCreation()
	//trackerTestDirectoryStorage()
	//trackerTestFileChangeTrackerAutoCreateFolderAndCleanup()
	//trackerTestNestedFastDirectoryCreation()

	// testing code to enable debugger use
	http.Handle("/event/", httpauth.SimpleBasicAuth("replicat", "isthecat")(http.HandlerFunc(eventHandler)))
	http.Handle("/tree/", httpauth.SimpleBasicAuth("replicat", "isthecat")(http.HandlerFunc(folderTreeHandler)))
	http.Handle("/config/", httpauth.SimpleBasicAuth("replicat", "isthecat")(http.HandlerFunc(configHandler)))
	http.Handle("/upload/", httpauth.SimpleBasicAuth("replicat", "isthecat")(http.HandlerFunc(uploadHandler)))

	// quick test of Minio
	suffix := rand.Int31n(10000)
	tempFolder := fmt.Sprintf("replicat-test-%05d", suffix)
	fmt.Printf("temporary path name: %s\n", tempFolder)
	tracker2 := &minioTracker{}
	tracker2.Initialize()
	//folderList, err := tracker2.ListFolders(false)
	//for _, v := range folderList {
	//	fmt.Printf("bucket: %s\n", v)
	//}

	objectName := "babySloth"

	err := tracker2.CreatePath(tempFolder, true)
	if err != nil {
		panic(err)
	}

	folderlist, err := tracker2.ListFolders(false)

	found := false
	for _, v := range folderlist {
		//fmt.Printf("Comparing %s and %s\n", v, tempFolder)
		if v == tempFolder {
			found = true
			break
		}
	}
	if found == false {
		fmt.Printf("Could not find the directory we should have just created: %s\n", tempFolder)
	}

	// Time to create file!
	err = tracker2.CreateObject(tempFolder, objectName, "/testdata/LargeSampleImages/baby sloth on a rail.jpg", "")
	if err != nil {
		panic(err)
	}

	err = tracker2.DeleteObject(tempFolder, objectName)
	if err != nil {
		panic(err)
	}

	err = tracker2.DeleteFolder(tempFolder)
	if err != nil {
		panic(err)
	}

	folderlist, err = tracker2.ListFolders(false)

	found = false
	for _, v := range folderlist {
		if v == tempFolder {
			found = true
			break
		}
	}

	if found == true {
		fmt.Printf("We found our folder after it should have been deleted: %s\n", tempFolder)
	}

	fmt.Println("YAY - Made it to the end")

	if tracker2 != nil {
		panic("nobody knows the trouble I've seen?\n")
	}

	lsnr, err := net.Listen("tcp4", address)
	if err != nil {
		panic(fmt.Sprintf("Error listening: %v\nAddress: %s\n", err, address))
	}
	fmt.Println("Listening on:", lsnr.Addr().String())

	logOnlyHandler := LogOnlyChangeHandler{}
	tracker := FilesystemTracker{}
	fmt.Printf("Looking up settings for node: %s\n", globalSettings.Name)

	directory := globalSettings.Directory

	fmt.Printf("GlobalSettings directory retrieved for this node: %s\n", directory)
	server := &ReplicatServer{Name: globalSettings.Name, ClusterKey: globalSettings.ClusterKey, Address: lsnr.Addr().String(), storage: &tracker, Status: REPLICAT_STATUS_INITIAL_SCAN}
	serverMap[globalSettings.Name] = server
	tracker.init(directory, server)

	go func(tracker StorageTracker) {
		for true {
			tracker.GetStatistics()
			time.Sleep(30 * time.Second)
		}
	}(&tracker)

	var c ChangeHandler
	c = &logOnlyHandler
	tracker.watchDirectory(&c)

	go func(listener net.Listener) {
		err = http.Serve(listener, nil)
		if err != nil {
			panic(err)
		}
	}(lsnr)

	fmt.Println("Starting config update processor")
	go configUpdateProcessor(configUpdateChannel)

	if globalSettings.ManagerAddress != "" {
		fmt.Printf("about to send config to server (%s)\nOur address is: (%s)\n", globalSettings.ManagerAddress, lsnr.Addr())
	}

	// we need something to keep us running.
	keepConfigCurrent()
}

const (
	managerOverdueSeconds = 30
	managerCheckSleepTime = 45
)

func keepConfigCurrent() {
	for {
		if time.Since(lastConfigPing) > managerOverdueSeconds*time.Second {
			log.Printf("Manager Contact Overdue, attempting to contact: %s\n", globalSettings.ManagerAddress)
			sendConfigToServer()
		} else {
			fmt.Println("No Update Required")
		}
		time.Sleep(time.Duration(rand.Intn(10)-5+managerCheckSleepTime) * time.Second)
	}
}

func sendConfigToServer() {
	// This field will be empty during testing. We have to abort during testing since the manager is not reachable.
	if globalSettings.ManagerAddress == "" {
		return
	}

	url := "https://" + globalSettings.ManagerAddress + "/config/"
	log.Printf("sendConfigToServer: Manager location: %s\n", url)

	server := serverMap[globalSettings.Name]
	jsonStr, _ := json.Marshal(server)
	jsonStr2, _ := json.Marshal(server.storage.GetStatistics())

	byteArraysToMarshal := [][]byte{jsonStr, jsonStr2}
	byteData := bytes.Join(byteArraysToMarshal, []byte{})
	log.Printf("jasonstr: %s\n", string(byteData))
	jsonData := bytes.NewBuffer(byteData)

	req, err := http.NewRequest("POST", url, jsonData)
	req.Header.Set("Content-Type", "application/json")

	data := []byte(globalSettings.ManagerCredentials)
	authHash := base64.StdEncoding.EncodeToString(data)
	req.Header.Add("Authorization", "Basic "+authHash)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("Unable to reach server. Data could be lost: %s\n", err)
		return
	}

	newServerMap, err := extractServerMapFromConfig(resp.Body)
	if err == io.EOF {
		log.Printf("EOF unexpected Config update failed\n")
	} else if err != nil {
		log.Printf("Config update failed due to error: %s\n", err)
	} else {
		configUpdateChannel <- newServerMap
	}

}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		r.ParseMultipartForm(32 << 20)
		file, handler, err := r.FormFile("uploadfile")
		if err != nil {
			fmt.Println(err)
			return
		}
		defer file.Close()
		fmt.Fprint(w, handler.Header)

		hash := []byte(r.Form.Get("HASH"))

		storage := serverMap[globalSettings.Name].storage
		local, _ := storage.getEntryJSON(handler.Filename)

		if !bytes.Equal(hash, local.Hash) {
			fullPath := globalSettings.Directory + "/" + handler.Filename
			f, err := os.OpenFile(fullPath, os.O_WRONLY|os.O_CREATE, 0666)

			if err != nil {
				fmt.Println(err)
				return
			}

			bytesWritten, err := io.Copy(f, file)
			if err != nil {
				log.Printf("Error copying file: %s, error(%#v)\n", handler.Filename, err)
				w.WriteHeader(http.StatusInternalServerError)
				w.Write([]byte("500 - Error copying file"))
				f.Close()
				return
			}
			fmt.Printf("Wrote out (%s) bytes (%d)\n", handler.Filename, bytesWritten)
			f.Close()

			// Get the old entry
			entryString := r.Form.Get("entryJSON")
			if entryString != "" {
				var entry EntryJSON
				err := json.Unmarshal([]byte(entryString), &entry)
				if err != nil {
					log.Fatalf("Error copying file (Entry handling): %s, error(%#v)\n", handler.Filename, err)
					w.WriteHeader(http.StatusInternalServerError)
					w.Write([]byte("500 - Error copying file"))
					return
				}

				err = os.Chtimes(fullPath, time.Now(), entry.ModTime)
				if err != nil {
					log.Fatalf("Error copying file (Changing times): %s, error(%#v)\n", handler.Filename, err)
					w.WriteHeader(http.StatusInternalServerError)
					w.Write([]byte("500 - Error copying file"))
					return
				}
			}

		}
	}
}

var configUpdateChannel = make(chan *map[string]*ReplicatServer, 100)

func configHandler(_ http.ResponseWriter, r *http.Request) {
	newServerMap, err := extractServerMapFromConfig(r.Body)
	if err != nil {
		fmt.Printf("Config update failed due to error: %s\n", err)
		return
	}
	configUpdateChannel <- newServerMap
}

func extractServerMapFromConfig(sourceReader io.Reader) (newServerMap *map[string]*ReplicatServer, err error) {
	decoder := json.NewDecoder(sourceReader)
	err = decoder.Decode(&newServerMap)
	return
}

func configUpdateProcessor(c chan *map[string]*ReplicatServer) {
	for {
		newServerMap := <-c
		serverMapLock.Lock()

		lastConfigPing = time.Now()
		sendData := false

		// find any nodes that have been deleted and update ones that have changed
		for name, serverData := range serverMap {
			newServerData, exists := (*newServerMap)[name]
			if !exists {
				fmt.Printf("No longer found config for: %s deleting\n", name)
				delete(serverMap, name)
				continue
			}

			if serverData.Address != newServerData.Address || serverData.Name != newServerData.Name || serverData.ClusterKey != newServerData.ClusterKey || serverData.Status != serverData.Status {
				fmt.Printf("Server data is changed. Replacing.\nold: %v\nnew: %v\n", &serverData, &newServerData)
				if serverData.Status != REPLICAT_STATUS_JOINING_CLUSTER && newServerData.Status == REPLICAT_STATUS_JOINING_CLUSTER {
					fmt.Printf("Decided to send data to: %s\n", serverData.Name)
					sendData = true
				}
				serverMap[name] = newServerData
				fmt.Println("Server data replaced with new server data")
			}
		}

		// find any new nodes
		for name, newServerData := range *newServerMap {
			_, exists := serverMap[name]
			if !exists {
				fmt.Printf("New server configuration for %s: %v\n", name, newServerData)

				// If this server map is for ourselves, build a list of folder if needed and notify others
				if name == globalSettings.Name {
					listOfFileInfo, err := scanDirectoryContents()
					if err != nil {
						log.Fatal(err)
					}
					newServerData.CurrentState = listOfFileInfo
					// Tell all of our friends that we exist and our current state for them to compare against.
					go func(tree DirTreeMap) {
						sendFolderTree(tree)
					}(listOfFileInfo)
				} else {
					fmt.Printf("New Node Decided to send data to: %s\n", name)
					sendData = true
				}

				fmt.Printf("New server configuration provided. Copying: %s\n", name)
				serverMap[name] = newServerData
			}
		}

		if sendData {
			server := serverMap[globalSettings.Name]
			fmt.Println("about to send existing files")
			server.storage.SendCatalog()
			fmt.Println("done sending existing files")
		}

		serverMapLock.Unlock()
	}
}
<|endoftext|>


<commit_before>package api

const (
	ReqError = "request error"
	NoError = "Success"
)

func GetGeneralStatusError(code int, module *ModuleAPI) (res string) {
	switch code{
	case 200, 204:
		res = NoError
	case 500:
		res = "server error"
	case 400:
		res = "bad parameter"
	case 404:
		if module.Module == Containers {
			res = "no such container or request 404"
		}else if module.Module == Images {
			res = "no such image or request 404"
		}else {
			res = "no possible error"
		}
	case 406:
		if module.Module == Containers {
			res = "impossible to attach (container not running)"
		}else {
			res = "no possible error"
		}
	case 304:
		if module.Module == Containers && module.Api == Start {
			res = "container already started"
		}else if module.Module == Containers && module.Api == Stop {
			res = "container already stopped"
		}
	case 409:
		res = "conflict"
	default:
		res = "no possible error"
	}
	return
}
<commit_msg>add APIError for docker-go<commit_after>package api

type APIError struct {
	ErrorMessage string
	StatusCode int
}

func (err APIError) Error() string {
	return string(err.StatusCode)+": "+err.ErrorMessage
}

func (err APIError) IsClientError() bool {
	return err.StatusCode >= 400 && err.StatusCode < 500
}

func (err APIError) IsClientError() bool {
	return err.StatusCode >= 500 && err.StatusCode < 600
}
<|endoftext|>