<commit_before>package main

import (
	"flag"
	"encoding/json"
	"log"
	"net/http"
	"os/exec"
	"runtime"
	"strconv"
	"sync"
)

var BASE_PORT = 10000
var globalWorldBaseDir string
var globalServerMap = NewServerMap()

type Server struct {
	id int
	creatorId int
	name string
	portOffset int
	handle *exec.Cmd
}

func NewServer(creatorId int, name string) *Server {
	// TODO Generate an id and port offset
	return &Server{
		creatorId: creatorId,
		name: name,
	}
}

type ServerMap struct {
	servers map[int]Server
	mutex   sync.Mutex
}

func (sm *ServerMap) Put(server Server) {
	sm.mutex.Lock()
	sm.servers[server.id] = server
	sm.mutex.Unlock()
}

func NewServerMap() *ServerMap {
	return &ServerMap{
		servers: make(map[int]Server),
	}
}

func runServer(server *Server) {
	app := "./server"

	arg0 := "-client"
	arg1 := "client"
	arg2 := "-world"
	arg3 := globalWorldBaseDir+"world"+strconv.Itoa(server.id)
	arg4 := "-host"
	arg5 := ":" + strconv.Itoa(BASE_PORT+server.id)

	cmd := exec.Command(app, arg0, arg1, arg2, arg3, arg4, arg5)
	server.handle = cmd
	globalServerMap.Put(*server)
	err := cmd.Run()

	if err != nil {
		log.Println(err.Error())
	}
}

type ApiGeneric struct {
	ServerId int
}

type ApiCreate struct {
	CreatorId int
	ServerName string
}

type ServerResponse struct {
	Id int
	CreatorId int
	Name string
	PortOffset int
}

func NewServerResponse(server Server) *ServerResponse {
	return &ServerResponse{
		Id: server.id,
		CreatorId: server.creatorId,
		Name: server.name,
		PortOffset: server.portOffset,
	}
}

type ServerListResponse struct {
	List map[string]ServerResponse
}

func parseGenericRequest(r *http.Request) (ApiGeneric, error) {
	var result ApiGeneric
	err := json.NewDecoder(r.Body).Decode(&result)
	if err != nil {
		log.Println("Error while parsing generic request")
		return result, err
	}
	return result, nil
}

func getHandler(w http.ResponseWriter, r *http.Request) {
	request, err := parseGenericRequest(r)
	if err != nil {
		return
	}

	server, found := globalServerMap.servers[request.ServerId]
	if !found {
		http.NotFound(w, r)
	}

	serverResponse := NewServerResponse(server)
	serverJson, err := json.Marshal(serverResponse)
	if err != nil {
		log.Println("Error marshalling server")
	}

	_, err = w.Write(serverJson)
	if err != nil {
		log.Println("Error sending response")
	}
}

func listHandler(w http.ResponseWriter, r *http.Request) {
}

func createHandler(w http.ResponseWriter, r *http.Request) {
	var request ApiCreate
	err := json.NewDecoder(r.Body).Decode(&request)
	if err != nil {
		log.Println("Error parsing create request", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	server := NewServer(request.CreatorId, request.ServerName)
	go runServer(server)
}

func deleteHandler(w http.ResponseWriter, r *http.Request) {
}

func main() {
	worldFolder := flag.String("worlds", "worlds/", "Sets the base folder used to store the worlds.")
	flag.Parse()
	globalWorldBaseDir = *worldFolder

	runtime.GOMAXPROCS(runtime.NumCPU())

	http.HandleFunc("/get", getHandler)
	http.HandleFunc("/list", listHandler)
	http.HandleFunc("/create", createHandler)
	http.HandleFunc("/delete", deleteHandler)
	err := http.ListenAndServe(":3001", nil)
	if err != nil {
		log.Fatal("ListenAndServe:", err)
	}
}
<commit_msg>Fix a bunch of things<commit_after>package main

import (
	"flag"
	"encoding/json"
	"log"
	"net/http"
	"os/exec"
	"runtime"
	"fmt"
	"sync"
	"path"
)

var BASE_PORT = 10000
var globalWorldBaseDir string
var globalServerMap = NewServerMap()

type Server struct {
	Id int
	CreatorId int
	Name string
	PortOffset int
	Handle *exec.Cmd `json:"-"`
}

func NewServer(creatorId int, name string) *Server {
	// TODO Generate an id and port offset
	return &Server{
		CreatorId: creatorId,
		Name: name,
	}
}

type ServerMap struct {
	servers map[string]*Server
	mutex   sync.Mutex
}

func (sm *ServerMap) Put(server *Server) {
	sm.mutex.Lock()
	defer sm.mutex.Unlock()
	sm.servers[str(server.Id)] = server
}

func (sm *ServerMap) Get(id int) *Server {
	sm.mutex.Lock()
	defer sm.mutex.Unlock()
	return sm.servers[str(id)]
}

func (sm *ServerMap) Encode(w http.ResponseWriter) error {
	sm.mutex.Lock()
	defer sm.mutex.Unlock()
	err := json.NewEncoder(w).Encode(sm.servers)
	return err
}

func NewServerMap() *ServerMap {
	return &ServerMap{
		servers: make(map[string]*Server),
	}
}

func runServer(server *Server) {
	app := "./server"

	arg0 := "-client"
	arg1 := "client"
	arg2 := "-world"
	arg3 := path.Join(globalWorldBaseDir, "world" + str(server.Id))
	arg4 := "-host"
	arg5 := ":" + str(BASE_PORT+server.Id)

	cmd := exec.Command(app, arg0, arg1, arg2, arg3, arg4, arg5)
	server.Handle = cmd
	globalServerMap.Put(server)
	err := cmd.Run()

	if err != nil {
		log.Println(err.Error())
	}
}

type ApiGeneric struct {
	ServerId int
}

type ApiCreate struct {
	CreatorId int
	ServerName string
}

type ServerListResponse struct {
	List map[string]Server
}

func parseGenericRequest(r *http.Request) (ApiGeneric, error) {
	var result ApiGeneric
	err := json.NewDecoder(r.Body).Decode(&result)
	if err != nil {
		log.Println("Error while parsing generic request")
		return result, err
	}
	return result, nil
}

func getHandler(w http.ResponseWriter, r *http.Request) {
	request, err := parseGenericRequest(r)
	if err != nil {
		return
	}

	server := globalServerMap.Get(request.ServerId)
	if server == nil {
		http.NotFound(w, r)
		return
	}

	err = json.NewEncoder(w).Encode(server)
	if err != nil {
		log.Println("Error marshalling server")
	}
}

func listHandler(w http.ResponseWriter, r *http.Request) {
	err := globalServerMap.Encode(w)
	if err != nil {
		http.Error(w, err.Error(), 500)
	}
}

func createHandler(w http.ResponseWriter, r *http.Request) {
	var request ApiCreate
	err := json.NewDecoder(r.Body).Decode(&request)
	if err != nil {
		log.Println("Error parsing create request", err)
		http.Error(w, err.Error(), 500)
	}

	server := NewServer(request.CreatorId, request.ServerName)
	go runServer(server)
}

func deleteHandler(w http.ResponseWriter, r *http.Request) {
}

func main() {
	worldFolder := flag.String("worlds", "worlds/", "Sets the base folder used to store the worlds.")
	flag.Parse()
	globalWorldBaseDir = *worldFolder

	runtime.GOMAXPROCS(runtime.NumCPU())

	http.HandleFunc("/get", getHandler)
	http.HandleFunc("/list", listHandler)
	http.HandleFunc("/create", createHandler)
	http.HandleFunc("/delete", deleteHandler)
	err := http.ListenAndServe(":3001", nil)
	if err != nil {
		log.Fatal("ListenAndServe:", err)
	}
}

// Utility function
func str(i int) string {
	return fmt.Sprintf("%d", i)
}
<|endoftext|>


<commit_before>package death

//Manage the death of your application.

import (
	LOG "github.com/cihub/seelog"
	"io"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"time"
)

//Death manages the death of your application.
type Death struct {
	wg          *sync.WaitGroup
	sigChannel  chan os.Signal
	callChannel chan struct{}
	timeout     time.Duration
	log         Logger
}

var empty struct{}

//Logger interface to log.
type Logger interface {
	Error(v ...interface{}) error
	Debug(v ...interface{})
	Info(v ...interface{})
	Warn(v ...interface{}) error
}

//closer is a wrapper to the struct we are going to close with metadata
//to help with debuging close.
type closer struct {
	Index   int
	C       io.Closer
	Name    string
	PKGPath string
}

//NewDeath Create Death with the signals you want to die from.
func NewDeath(signals ...os.Signal) (death *Death) {
	death = &Death{timeout: 10 * time.Second,
		sigChannel:  make(chan os.Signal, 1),
		callChannel: make(chan struct{}, 1),
		wg:          &sync.WaitGroup{},
		log:         LOG.Current}
	signal.Notify(death.sigChannel, signals...)
	death.wg.Add(1)
	go death.listenForSignal()
	return death
}

//Override the time death is willing to wait for a objects to be closed.
func (d *Death) SetTimeout(t time.Duration) {
	d.timeout = t
}

//setLogger Override the default logger (seelog)
func (d *Death) SetLogger(l Logger) {
	d.log = l
}

//WaitForDeath wait for signal and then kill all items that need to die.
func (d *Death) WaitForDeath(closable ...io.Closer) {
	d.wg.Wait()
	d.log.Info("Shutdown started...")
	count := len(closable)
	d.log.Debug("Closing ", count, " objects")
	if count > 0 {
		d.closeInMass(closable...)
	}
}

//WaitForDeathWithFunc allows you to have a single function get called when it's time to
//kill your application.
func (d *Death) WaitForDeathWithFunc(f func()) {
	d.wg.Wait()
	d.log.Info("Shutdown started...")
	f()
}

//getPkgPath for an io closer.
func getPkgPath(c io.Closer) (name string, pkgPath string) {

	t := reflect.TypeOf(c)
	name, pkgPath = t.Name(), t.PkgPath()
	switch t.Kind() {
	case reflect.Ptr:
		name = t.Elem().Name()
		pkgPath = t.Elem().PkgPath()
	}
	return name, pkgPath

}

//closeInMass Close all the objects at once and wait forr them to finish with a channel.
func (d *Death) closeInMass(closable ...io.Closer) {

	count := len(closable)
	sentToClose := make(map[int]closer)
	//call close async
	doneClosers := make(chan closer, count)
	for i, c := range closable {
		name, pkgPath := getPkgPath(c)
		closer := closer{Index: i, C: c, Name: name, PKGPath: pkgPath}
		go d.closeObjects(closer, doneClosers)
		sentToClose[i] = closer
	}

	//wait on channel for notifications.

	timer := time.NewTimer(d.timeout)
	for {
		select {
		case <-timer.C:
			d.log.Warn(count, " object(s) remaining but timer expired.")
			for _, c := range sentToClose {
				d.log.Error("Failed to close: ", c.PKGPath, "/", c.Name)
			}
			return
		case closer := <-doneClosers:
			delete(sentToClose, closer.Index)
			count--
			d.log.Debug(count, " object(s) left")
			if count == 0 && len(sentToClose) == 0 {
				d.log.Debug("Finished closing objects")
				return
			}
		}
	}
}

//closeObjects and return a bool when finished on a channel.
func (d *Death) closeObjects(closer closer, done chan<- closer) {
	err := closer.C.Close()
	if nil != err {
		d.log.Error(err)
	}
	done <- closer
}

//FallOnSword manually initiates the death process.
func (d *Death) FallOnSword() {
	select {
	case d.callChannel <- struct{}{}:
	default:
	}
}

//ListenForSignal Manage death of application by signal.
func (d *Death) listenForSignal() {
	defer d.wg.Done()
	for {
		select {
		case <-d.sigChannel:
			return
		case <-d.callChannel:
			return
		}
	}
}
<commit_msg>simplify getPkgPath<commit_after>package death

//Manage the death of your application.

import (
	LOG "github.com/cihub/seelog"
	"io"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"time"
)

//Death manages the death of your application.
type Death struct {
	wg          *sync.WaitGroup
	sigChannel  chan os.Signal
	callChannel chan struct{}
	timeout     time.Duration
	log         Logger
}

var empty struct{}

//Logger interface to log.
type Logger interface {
	Error(v ...interface{}) error
	Debug(v ...interface{})
	Info(v ...interface{})
	Warn(v ...interface{}) error
}

//closer is a wrapper to the struct we are going to close with metadata
//to help with debuging close.
type closer struct {
	Index   int
	C       io.Closer
	Name    string
	PKGPath string
}

//NewDeath Create Death with the signals you want to die from.
func NewDeath(signals ...os.Signal) (death *Death) {
	death = &Death{timeout: 10 * time.Second,
		sigChannel:  make(chan os.Signal, 1),
		callChannel: make(chan struct{}, 1),
		wg:          &sync.WaitGroup{},
		log:         LOG.Current}
	signal.Notify(death.sigChannel, signals...)
	death.wg.Add(1)
	go death.listenForSignal()
	return death
}

//Override the time death is willing to wait for a objects to be closed.
func (d *Death) SetTimeout(t time.Duration) {
	d.timeout = t
}

//setLogger Override the default logger (seelog)
func (d *Death) SetLogger(l Logger) {
	d.log = l
}

//WaitForDeath wait for signal and then kill all items that need to die.
func (d *Death) WaitForDeath(closable ...io.Closer) {
	d.wg.Wait()
	d.log.Info("Shutdown started...")
	count := len(closable)
	d.log.Debug("Closing ", count, " objects")
	if count > 0 {
		d.closeInMass(closable...)
	}
}

//WaitForDeathWithFunc allows you to have a single function get called when it's time to
//kill your application.
func (d *Death) WaitForDeathWithFunc(f func()) {
	d.wg.Wait()
	d.log.Info("Shutdown started...")
	f()
}

//getPkgPath for an io closer.
func getPkgPath(c io.Closer) (name string, pkgPath string) {
	t := reflect.TypeOf(c)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return t.Name(), t.PkgPath()
}

//closeInMass Close all the objects at once and wait forr them to finish with a channel.
func (d *Death) closeInMass(closable ...io.Closer) {

	count := len(closable)
	sentToClose := make(map[int]closer)
	//call close async
	doneClosers := make(chan closer, count)
	for i, c := range closable {
		name, pkgPath := getPkgPath(c)
		closer := closer{Index: i, C: c, Name: name, PKGPath: pkgPath}
		go d.closeObjects(closer, doneClosers)
		sentToClose[i] = closer
	}

	//wait on channel for notifications.

	timer := time.NewTimer(d.timeout)
	for {
		select {
		case <-timer.C:
			d.log.Warn(count, " object(s) remaining but timer expired.")
			for _, c := range sentToClose {
				d.log.Error("Failed to close: ", c.PKGPath, "/", c.Name)
			}
			return
		case closer := <-doneClosers:
			delete(sentToClose, closer.Index)
			count--
			d.log.Debug(count, " object(s) left")
			if count == 0 && len(sentToClose) == 0 {
				d.log.Debug("Finished closing objects")
				return
			}
		}
	}
}

//closeObjects and return a bool when finished on a channel.
func (d *Death) closeObjects(closer closer, done chan<- closer) {
	err := closer.C.Close()
	if nil != err {
		d.log.Error(err)
	}
	done <- closer
}

//FallOnSword manually initiates the death process.
func (d *Death) FallOnSword() {
	select {
	case d.callChannel <- struct{}{}:
	default:
	}
}

//ListenForSignal Manage death of application by signal.
func (d *Death) listenForSignal() {
	defer d.wg.Done()
	for {
		select {
		case <-d.sigChannel:
			return
		case <-d.callChannel:
			return
		}
	}
}
<|endoftext|>


<commit_before>package death

import (
	LOG "github.com/cihub/seelog"
	"io"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"time"
)

type Death struct {
	wg         *sync.WaitGroup
	sigChannel chan os.Signal
	timeout    time.Duration
	log        Logger
}

var empty struct{}

//Logger interface to log.
type Logger interface {
	Error(v ...interface{}) error
	Debug(v ...interface{})
	Info(v ...interface{})
	Warn(v ...interface{}) error
}

type closer struct {
	C    io.Closer
	Name string
}

//Create Death with the signals you want to die from.
func NewDeath(signals ...os.Signal) (death *Death) {
	death = &Death{timeout: 10 * time.Second,
		sigChannel: make(chan os.Signal, 1),
		wg:         &sync.WaitGroup{},
		log:        LOG.Current}
	signal.Notify(death.sigChannel, signals...)
	death.wg.Add(1)
	go death.listenForSignal(death.sigChannel)
	return death
}

//Override the time death is willing to wait for a objects to be closed.
func (d *Death) setTimeout(t time.Duration) {
	d.timeout = t
}

//setLogger Override the default logger (seelog)
func (d *Death) setLogger(l Logger) {
	d.log = l
}

//Wait for death and then kill all items that need to die.
func (d *Death) WaitForDeath(closable ...io.Closer) {
	d.wg.Wait()
	d.log.Info("Shutdown started...")
	count := len(closable)
	d.log.Debug("Closing ", count, " objects")
	if count > 0 {
		d.closeInMass(closable...)
	}
}

//Close all the objects at once and wait forr them to finish with a channel.
func (d *Death) closeInMass(closable ...io.Closer) {

	count := len(closable)
	sentToClose := make(map[closer]struct{})
	//call close async
	doneClosers := make(chan closer, count)
	for _, c := range closable {
		elem := reflect.TypeOf(c).Elem()

		closer := closer{C: c, Name: elem.Name()}
		go d.closeObjects(closer, doneClosers)
		sentToClose[closer] = empty
	}

	//wait on channel for notifications.

	timer := time.NewTimer(d.timeout)
	for {
		select {
		case <-timer.C:
			d.log.Warn(count, " object(s) remaining but timer expired.")
			for c, _ := range sentToClose {
				d.log.Error("Failed to close: ", c.Name)
			}
			return
		case closer := <-doneClosers:
			delete(sentToClose, closer)
			count--
			d.log.Debug(count, " object(s) left")
			d.log.Debug("Closers: ", sentToClose)
			if count == 0 && len(sentToClose) == 0 {
				d.log.Debug("Finished closing objects")
				return
			}
		}
	}
}

//closeObjects and return a bool when finished on a channel.
func (d *Death) closeObjects(closer closer, done chan<- closer) {
	err := closer.C.Close()
	if nil != err {
		d.log.Error(err)
	}
	done <- closer
}

//ListenForSignal Manage death of application by signal.
func (d *Death) listenForSignal(c <-chan os.Signal) {
	defer d.wg.Done()
	for range c {
		return
	}
}
<commit_msg>add pkg path<commit_after>package death

import (
	LOG "github.com/cihub/seelog"
	"io"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"time"
)

type Death struct {
	wg         *sync.WaitGroup
	sigChannel chan os.Signal
	timeout    time.Duration
	log        Logger
}

var empty struct{}

//Logger interface to log.
type Logger interface {
	Error(v ...interface{}) error
	Debug(v ...interface{})
	Info(v ...interface{})
	Warn(v ...interface{}) error
}

type closer struct {
	C       io.Closer
	Name    string
	PKGPath string
}

//Create Death with the signals you want to die from.
func NewDeath(signals ...os.Signal) (death *Death) {
	death = &Death{timeout: 10 * time.Second,
		sigChannel: make(chan os.Signal, 1),
		wg:         &sync.WaitGroup{},
		log:        LOG.Current}
	signal.Notify(death.sigChannel, signals...)
	death.wg.Add(1)
	go death.listenForSignal(death.sigChannel)
	return death
}

//Override the time death is willing to wait for a objects to be closed.
func (d *Death) setTimeout(t time.Duration) {
	d.timeout = t
}

//setLogger Override the default logger (seelog)
func (d *Death) setLogger(l Logger) {
	d.log = l
}

//Wait for death and then kill all items that need to die.
func (d *Death) WaitForDeath(closable ...io.Closer) {
	d.wg.Wait()
	d.log.Info("Shutdown started...")
	count := len(closable)
	d.log.Debug("Closing ", count, " objects")
	if count > 0 {
		d.closeInMass(closable...)
	}
}

//Close all the objects at once and wait forr them to finish with a channel.
func (d *Death) closeInMass(closable ...io.Closer) {

	count := len(closable)
	sentToClose := make(map[closer]struct{})
	//call close async
	doneClosers := make(chan closer, count)
	for _, c := range closable {
		elem := reflect.TypeOf(c).Elem()
		closer := closer{C: c, Name: elem.Name(), PKGPath: elem.PkgPath()}
		go d.closeObjects(closer, doneClosers)
		sentToClose[closer] = empty
	}

	//wait on channel for notifications.

	timer := time.NewTimer(d.timeout)
	for {
		select {
		case <-timer.C:
			d.log.Warn(count, " object(s) remaining but timer expired.")
			for c, _ := range sentToClose {
				d.log.Error("Failed to close: ", c.PKGPath, "/", c.Name)
			}
			return
		case closer := <-doneClosers:
			delete(sentToClose, closer)
			count--
			d.log.Debug(count, " object(s) left")
			d.log.Debug("Closers: ", sentToClose)
			if count == 0 && len(sentToClose) == 0 {
				d.log.Debug("Finished closing objects")
				return
			}
		}
	}
}

//closeObjects and return a bool when finished on a channel.
func (d *Death) closeObjects(closer closer, done chan<- closer) {
	err := closer.C.Close()
	if nil != err {
		d.log.Error(err)
	}
	done <- closer
}

//ListenForSignal Manage death of application by signal.
func (d *Death) listenForSignal(c <-chan os.Signal) {
	defer d.wg.Done()
	for range c {
		return
	}
}
<|endoftext|>


<commit_before>package handlers

var validStates = map[string]bool{
	"AL" : true,
	"AK" : true,
	"AZ" : true,
	"AR" : true,
	"CA" : true,
	"CO" : true,
	"CT" : true,
	"DE" : true,
	"FL" : true,
	"GA" : true,
	"HI" : true,
	"ID" : true,
	"IL" : true,
	"IN" : true,
	"IA" : true,
	"KS" : true,
	"KY" : true,
	"LA" : true,
	"ME" : true,
	"MD" : true,
	"MA" : true,
	"MI" : true,
	"MN" : true,
	"MS" : true,
	"MO" : true,
	"MT" : true,
	"NE" : true,
	"NV" : true,
	"NH" : true,
	"NJ" : true,
	"NM" : true,
	"NY" : true,
	"NC" : true,
	"ND" : true,
	"OH" : true,
	"OK" : true,
	"OR" : true,
	"PA" : true,
	"RI" : true,
	"SC" : true,
	"SD" : true,
	"TN" : true,
	"TX" : true,
	"UT" : true,
	"VT" : true,
	"VA" : true,
	"WA" : true,
	"WV" : true,
	"WI" : true,
	"WY" : true,
}

<commit_msg>sort state list<commit_after>package handlers

var validStates = map[string]bool{
	"AK" : true,
	"AL" : true,
	"AR" : true,
	"AZ" : true,
	"CA" : true,
	"CO" : true,
	"CT" : true,
	"DE" : true,
	"FL" : true,
	"GA" : true,
	"HI" : true,
	"IA" : true,
	"ID" : true,
	"IL" : true,
	"IN" : true,
	"KS" : true,
	"KY" : true,
	"LA" : true,
	"MA" : true,
	"MD" : true,
	"ME" : true,
	"MI" : true,
	"MN" : true,
	"MO" : true,
	"MS" : true,
	"MT" : true,
	"NC" : true,
	"ND" : true,
	"NE" : true,
	"NH" : true,
	"NJ" : true,
	"NM" : true,
	"NV" : true,
	"NY" : true,
	"OH" : true,
	"OK" : true,
	"OR" : true,
	"PA" : true,
	"RI" : true,
	"SC" : true,
	"SD" : true,
	"TN" : true,
	"TX" : true,
	"UT" : true,
	"VA" : true,
	"VT" : true,
	"WA" : true,
	"WI" : true,
	"WV" : true,
	"WY" : true,
}

<|endoftext|>


<commit_before>/*
Copyright 2022 The TestGrid Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package queue

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"cloud.google.com/go/storage"
	"github.com/GoogleCloudPlatform/testgrid/util/gcs"
	"github.com/sirupsen/logrus"
)

// PersistClient contains interfaces for reading from and writing to a Path.
type PersistClient interface {
	gcs.Uploader
	gcs.Opener
}

// Fixer will adjust the queue until the context expires.
type Fixer func(context.Context, *Queue) error

// FixPersistent persists the queue to the remote path every tick.
//
// The first time it will load the state. Thereafter it will save the state.
// This includes restarts due to expiring contexts -- it will just load once.
func FixPersistent(logr logrus.FieldLogger, client PersistClient, path gcs.Path, tick <-chan time.Time) Fixer {
	var shouldSave bool
	log := logr.WithField("path", path)
	return func(ctx context.Context, q *Queue) error {
		log.Debug("Using persistent state")

		tryLoad := func() error {
			reader, attrs, err := client.Open(ctx, path)
			if errors.Is(err, storage.ErrObjectNotExist) {
				log.Info("Previous persistent queue state does not exist.")
				return nil
			}
			if err != nil {
				return fmt.Errorf("open: %w", err)
			}

			defer reader.Close()
			dec := json.NewDecoder(reader)
			var whens map[string]time.Time
			if err := dec.Decode(&whens); err != nil {
				return fmt.Errorf("decode: %v", err)
			}

			current := q.Current()

			for name := range whens {
				if _, ok := current[name]; ok {
					continue
				}
				delete(whens, name)
			}

			log.WithField("from", attrs.LastModified).Info("Loaded previous state, syncing queue.")
			if err := q.FixAll(whens, false); err != nil {
				return fmt.Errorf("fix all: %v", err)
			}
			return nil
		}

		logSave := true

		trySave := func() error {
			currently := q.Current()
			buf, err := json.MarshalIndent(currently, "", "  ")
			if err != nil {
				return fmt.Errorf("marshal: %v", err)
			}
			attrs, err := client.Upload(ctx, path, buf, false, "")
			if err == nil && logSave {
				logSave = false
				log.WithField("updated", attrs.Updated).Info("Wrote persistent state")
			}
			return err
		}

		for {
			select {
			case <-ctx.Done():
				log.Debug("Stopped syncing persistent state")
				return ctx.Err()
			case <-tick:
			}

			if shouldSave {
				log.Trace("Saving persistent state...")
				if err := trySave(); err != nil {
					log.WithError(err).Error("Failed to save persistent state.")
					continue
				}
				log.Debug("Saved persistent state.")
			} else {
				log.Trace("Loading persistent state...")
				if err := tryLoad(); err != nil {
					log.WithError(err).Error("Failed to load persistent state.")
					continue
				}
				shouldSave = true
				log.Debug("Loaded persistent state.")
			}
		}
	}
}
<commit_msg>Improve persistent queue routine.<commit_after>/*
Copyright 2022 The TestGrid Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package queue

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"cloud.google.com/go/storage"
	"github.com/GoogleCloudPlatform/testgrid/util/gcs"
	"github.com/sirupsen/logrus"
)

// PersistClient contains interfaces for reading from and writing to a Path.
type PersistClient interface {
	gcs.Uploader
	gcs.Opener
}

// Fixer will adjust the queue until the context expires.
type Fixer func(context.Context, *Queue) error

// FixPersistent persists the queue to the remote path every tick.
//
// The first time it will load the state. Thereafter it will save the state.
// This includes restarts due to expiring contexts -- it will just load once.
func FixPersistent(logr logrus.FieldLogger, client PersistClient, path gcs.Path, tick <-chan time.Time) Fixer {
	var shouldSave bool
	log := logr.WithField("path", path)
	return func(parentCtx context.Context, q *Queue) error {
		log.Debug("Using persistent state")

		ctx, cancel := context.WithCancel(parentCtx)
		defer cancel()
		go func() { // allow a grace period for reads/writes
			select {
			case <-ctx.Done():
				return
			case <-parentCtx.Done():
				timer := time.NewTimer(5 * time.Second)
				select {
				case <-ctx.Done():
					if !timer.Stop() {
						<-timer.C
					}
					return
				case <-timer.C:
					cancel()
					return
				}
			}
		}()

		tryLoad := func() error {
			reader, attrs, err := client.Open(ctx, path)
			if errors.Is(err, storage.ErrObjectNotExist) {
				log.Info("Previous persistent queue state does not exist.")
				return nil
			}
			if err != nil {
				return fmt.Errorf("open: %w", err)
			}

			defer reader.Close()
			dec := json.NewDecoder(reader)
			var whens map[string]time.Time
			if err := dec.Decode(&whens); err != nil {
				return fmt.Errorf("decode: %v", err)
			}

			current := q.Current()

			for name := range whens {
				if _, ok := current[name]; ok {
					continue
				}
				delete(whens, name)
			}

			log.WithField("from", attrs.LastModified).Info("Loaded previous state, syncing queue.")
			if err := q.FixAll(whens, false); err != nil {
				return fmt.Errorf("fix all: %v", err)
			}
			return nil
		}

		logSave := true

		trySave := func() error {
			currently := q.Current()
			buf, err := json.MarshalIndent(currently, "", "  ")
			if err != nil {
				return fmt.Errorf("marshal: %v", err)
			}
			attrs, err := client.Upload(ctx, path, buf, false, "")
			if err == nil && logSave {
				logSave = false
				log.WithField("updated", attrs.Updated).Info("Wrote persistent state")
			}
			return err
		}

		for {
			select {
			case <-parentCtx.Done():
				log.Debug("Stopped syncing persistent state")
				return parentCtx.Err()
			case <-tick:
			}

			if shouldSave {
				log.Trace("Saving persistent state...")
				if err := trySave(); err != nil {
					log := log.WithError(err)
					var chirp func(...interface{})
					if errors.Is(err, context.Canceled) {
						chirp = log.Debug
					} else {
						chirp = log.Warning
					}
					chirp("Failed to save persistent state.")
					continue
				}
				log.Debug("Saved persistent state.")
			} else {
				log.Trace("Loading persistent state...")
				if err := tryLoad(); err != nil {
					log := log.WithError(err)
					var chirp func(...interface{})
					if errors.Is(err, context.Canceled) {
						chirp = log.Debug
					} else {
						chirp = log.Warning
					}
					chirp("Failed to load persistent state.")
					continue
				}
				shouldSave = true
				log.Debug("Loaded persistent state.")
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"strings"
	"text/template"

	"github.com/Masterminds/sprig"
	"github.com/huandu/xstrings"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	options "github.com/grpc-ecosystem/grpc-gateway/third_party/googleapis/google/api"
)

var ProtoHelpersFuncMap = template.FuncMap{
	"string": func(i interface {
		String() string
	}) string {
		return i.String()
	},
	"json": func(v interface{}) string {
		a, _ := json.Marshal(v)
		return string(a)
	},
	"prettyjson": func(v interface{}) string {
		a, _ := json.MarshalIndent(v, "", "  ")
		return string(a)
	},
	"splitArray": func(sep string, s string) []string {
		return strings.Split(s, sep)
	},
	"first": func(a []string) string {
		return a[0]
	},
	"last": func(a []string) string {
		return a[len(a)-1]
	},
	"upperFirst": func(s string) string {
		return strings.ToUpper(s[:1]) + s[1:]
	},
	"lowerFirst": func(s string) string {
		return strings.ToLower(s[:1]) + s[1:]
	},
	"camelCase": func(s string) string {
		return xstrings.ToCamelCase(s)
	},
	"lowerCamelCase": func(s string) string {
		cc := xstrings.ToCamelCase(s)
		return strings.ToLower(cc[:1]) + cc[1:]
	},
	"snakeCase": func(s string) string {
		return xstrings.ToSnakeCase(s)
	},
	"kebabCase": func(s string) string {
		return strings.Replace(xstrings.ToSnakeCase(s), "_", "-", -1)
	},
	"getMessageType":  getMessageType,
	"isFieldMessage":  isFieldMessage,
	"isFieldRepeated": isFieldRepeated,
	"goType":          goType,
	"httpVerb":        httpVerb,
	"httpPath":        httpPath,
}

func init() {
	for k, v := range sprig.TxtFuncMap() {
		ProtoHelpersFuncMap[k] = v
	}
}

func getMessageType(f *descriptor.FileDescriptorProto, name string) *descriptor.DescriptorProto {
	// name is in the form .packageName.MessageTypeName.InnerMessageTypeName...
	// e.g. .article.ProductTag
	splits := strings.Split(name, ".")
	target := splits[len(splits)-1]
	for _, m := range f.MessageType {
		if target == *m.Name {
			return m
		}
	}

	return nil
}

func isFieldMessage(f *descriptor.FieldDescriptorProto) bool {
	if f.Type != nil && *f.Type == descriptor.FieldDescriptorProto_TYPE_MESSAGE {
		return true
	}

	return false
}

func isFieldRepeated(f *descriptor.FieldDescriptorProto) bool {
	if f.Type != nil && f.Label != nil && *f.Label == descriptor.FieldDescriptorProto_LABEL_REPEATED {
		return true
	}

	return false
}

func goType(pkg string, f *descriptor.FieldDescriptorProto) string {
	switch *f.Type {
	case descriptor.FieldDescriptorProto_TYPE_DOUBLE:
		return "float64"
	case descriptor.FieldDescriptorProto_TYPE_FLOAT:
		return "float32"
	case descriptor.FieldDescriptorProto_TYPE_INT64:
		return "int64"
	case descriptor.FieldDescriptorProto_TYPE_UINT64:
		return "uint64"
	case descriptor.FieldDescriptorProto_TYPE_INT32:
		return "uint32"
	case descriptor.FieldDescriptorProto_TYPE_BOOL:
		return "bool"
	case descriptor.FieldDescriptorProto_TYPE_STRING:
		return "string"
	case descriptor.FieldDescriptorProto_TYPE_MESSAGE:
		if *f.Label == descriptor.FieldDescriptorProto_LABEL_REPEATED {
			return fmt.Sprintf("[]*%s.%s", pkg, shortType(*f.TypeName))
		}
		return fmt.Sprintf("*%s.%s", pkg, shortType(*f.TypeName))
	case descriptor.FieldDescriptorProto_TYPE_BYTES:
		return "byte"
	case descriptor.FieldDescriptorProto_TYPE_UINT32:
		return "uint32"
	case descriptor.FieldDescriptorProto_TYPE_ENUM:
		return fmt.Sprintf("*%s.%s", pkg, shortType(*f.TypeName))
	default:
		return "interface{}"
	}
}

func shortType(s string) string {
	t := strings.Split(s, ".")
	return t[len(t)-1]
}

func httpPath(m *descriptor.MethodDescriptorProto) string {

	ext, err := proto.GetExtension(m.Options, options.E_Http)
	if err != nil {
		return err.Error()
	}
	opts, ok := ext.(*options.HttpRule)
	if !ok {
		return fmt.Sprintf("extension is %T; want an HttpRule", ext)
	}

	switch t := opts.Pattern.(type) {
	default:
		return ""
	case *options.HttpRule_Get:
		return t.Get
	case *options.HttpRule_Post:
		return t.Post
	case *options.HttpRule_Put:
		return t.Put
	case *options.HttpRule_Delete:
		return t.Delete
	case *options.HttpRule_Patch:
		return t.Patch
	case *options.HttpRule_Custom:
		return t.Custom.Path
	}
}

func httpVerb(m *descriptor.MethodDescriptorProto) string {

	ext, err := proto.GetExtension(m.Options, options.E_Http)
	if err != nil {
		return err.Error()
	}
	opts, ok := ext.(*options.HttpRule)
	if !ok {
		return fmt.Sprintf("extension is %T; want an HttpRule", ext)
	}

	switch t := opts.Pattern.(type) {
	default:
		return ""
	case *options.HttpRule_Get:
		return "GET"
	case *options.HttpRule_Post:
		return "POST"
	case *options.HttpRule_Put:
		return "PUT"
	case *options.HttpRule_Delete:
		return "DELETE"
	case *options.HttpRule_Patch:
		return "PATCH"
	case *options.HttpRule_Custom:
		return t.Custom.Kind
	}
}
<commit_msg>Fix camelCase when len = 1<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"strings"
	"text/template"

	"github.com/Masterminds/sprig"
	"github.com/huandu/xstrings"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	options "github.com/grpc-ecosystem/grpc-gateway/third_party/googleapis/google/api"
)

var ProtoHelpersFuncMap = template.FuncMap{
	"string": func(i interface {
		String() string
	}) string {
		return i.String()
	},
	"json": func(v interface{}) string {
		a, _ := json.Marshal(v)
		return string(a)
	},
	"prettyjson": func(v interface{}) string {
		a, _ := json.MarshalIndent(v, "", "  ")
		return string(a)
	},
	"splitArray": func(sep string, s string) []string {
		return strings.Split(s, sep)
	},
	"first": func(a []string) string {
		return a[0]
	},
	"last": func(a []string) string {
		return a[len(a)-1]
	},
	"upperFirst": func(s string) string {
		return strings.ToUpper(s[:1]) + s[1:]
	},
	"lowerFirst": func(s string) string {
		return strings.ToLower(s[:1]) + s[1:]
	},
	"camelCase": func(s string) string {
		if len(s) > 1 {
			return xstrings.ToCamelCase(s)
		}

		return strings.ToUpper(s[:1])
	},
	"lowerCamelCase": func(s string) string {
		if len(s) > 1 {
			s = xstrings.ToCamelCase(s)
		}

		return strings.ToLower(s[:1]) + s[1:]
	},
	"snakeCase": func(s string) string {
		return xstrings.ToSnakeCase(s)
	},
	"kebabCase": func(s string) string {
		return strings.Replace(xstrings.ToSnakeCase(s), "_", "-", -1)
	},
	"getMessageType":  getMessageType,
	"isFieldMessage":  isFieldMessage,
	"isFieldRepeated": isFieldRepeated,
	"goType":          goType,
	"httpVerb":        httpVerb,
	"httpPath":        httpPath,
}

func init() {
	for k, v := range sprig.TxtFuncMap() {
		ProtoHelpersFuncMap[k] = v
	}
}

func getMessageType(f *descriptor.FileDescriptorProto, name string) *descriptor.DescriptorProto {
	// name is in the form .packageName.MessageTypeName.InnerMessageTypeName...
	// e.g. .article.ProductTag
	splits := strings.Split(name, ".")
	target := splits[len(splits)-1]
	for _, m := range f.MessageType {
		if target == *m.Name {
			return m
		}
	}

	return nil
}

func isFieldMessage(f *descriptor.FieldDescriptorProto) bool {
	if f.Type != nil && *f.Type == descriptor.FieldDescriptorProto_TYPE_MESSAGE {
		return true
	}

	return false
}

func isFieldRepeated(f *descriptor.FieldDescriptorProto) bool {
	if f.Type != nil && f.Label != nil && *f.Label == descriptor.FieldDescriptorProto_LABEL_REPEATED {
		return true
	}

	return false
}

func goType(pkg string, f *descriptor.FieldDescriptorProto) string {
	switch *f.Type {
	case descriptor.FieldDescriptorProto_TYPE_DOUBLE:
		return "float64"
	case descriptor.FieldDescriptorProto_TYPE_FLOAT:
		return "float32"
	case descriptor.FieldDescriptorProto_TYPE_INT64:
		return "int64"
	case descriptor.FieldDescriptorProto_TYPE_UINT64:
		return "uint64"
	case descriptor.FieldDescriptorProto_TYPE_INT32:
		return "uint32"
	case descriptor.FieldDescriptorProto_TYPE_BOOL:
		return "bool"
	case descriptor.FieldDescriptorProto_TYPE_STRING:
		return "string"
	case descriptor.FieldDescriptorProto_TYPE_MESSAGE:
		if *f.Label == descriptor.FieldDescriptorProto_LABEL_REPEATED {
			return fmt.Sprintf("[]*%s.%s", pkg, shortType(*f.TypeName))
		}
		return fmt.Sprintf("*%s.%s", pkg, shortType(*f.TypeName))
	case descriptor.FieldDescriptorProto_TYPE_BYTES:
		return "byte"
	case descriptor.FieldDescriptorProto_TYPE_UINT32:
		return "uint32"
	case descriptor.FieldDescriptorProto_TYPE_ENUM:
		return fmt.Sprintf("*%s.%s", pkg, shortType(*f.TypeName))
	default:
		return "interface{}"
	}
}

func shortType(s string) string {
	t := strings.Split(s, ".")
	return t[len(t)-1]
}

func httpPath(m *descriptor.MethodDescriptorProto) string {

	ext, err := proto.GetExtension(m.Options, options.E_Http)
	if err != nil {
		return err.Error()
	}
	opts, ok := ext.(*options.HttpRule)
	if !ok {
		return fmt.Sprintf("extension is %T; want an HttpRule", ext)
	}

	switch t := opts.Pattern.(type) {
	default:
		return ""
	case *options.HttpRule_Get:
		return t.Get
	case *options.HttpRule_Post:
		return t.Post
	case *options.HttpRule_Put:
		return t.Put
	case *options.HttpRule_Delete:
		return t.Delete
	case *options.HttpRule_Patch:
		return t.Patch
	case *options.HttpRule_Custom:
		return t.Custom.Path
	}
}

func httpVerb(m *descriptor.MethodDescriptorProto) string {

	ext, err := proto.GetExtension(m.Options, options.E_Http)
	if err != nil {
		return err.Error()
	}
	opts, ok := ext.(*options.HttpRule)
	if !ok {
		return fmt.Sprintf("extension is %T; want an HttpRule", ext)
	}

	switch t := opts.Pattern.(type) {
	default:
		return ""
	case *options.HttpRule_Get:
		return "GET"
	case *options.HttpRule_Post:
		return "POST"
	case *options.HttpRule_Put:
		return "PUT"
	case *options.HttpRule_Delete:
		return "DELETE"
	case *options.HttpRule_Patch:
		return "PATCH"
	case *options.HttpRule_Custom:
		return t.Custom.Kind
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 Will Fitzgerald. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
	Package bitset implements bitsets.

	It provides methods for making a BitSet of an arbitrary
	upper limit, setting and testing bit locations, and clearing
	bit locations as well as the entire set.

	BitSets are implemented as arrays of uint64s, so it may be best
	to limit the upper size to a multiple of 64. It is not an error
	to set, test, or clear a bit within a 64 bit boundary.

	Example use:

	b := MakeBitSet(64000)
	b.SetBit(1000)
	if b.Bit(1000) {
		b.ClearBit(1000)
	}
*/
package bitset

// for MaxUint64
import (
	"math"
)

// BitSets are arrays of uint64.
type BitSet []uint64

// Make a BitSet with an upper limit on size. Note this is the
// number of bits, not the number of uint64s, which is a kind of
// implementation detail.
func MakeBitSet(max_size uint) BitSet {
	return make(BitSet, (max_size+(64-1))/64)
}

/// Test whether bit i is set.
func (set BitSet) Bit(i uint) bool {
	return ((set[i/64] & (1<<(i%64)))!=0)
}

// Set bit i to 1
func (set BitSet) SetBit(i uint) {
	set[i/64]|=(1<<(i%64))
}

// Clear bit i to 0
func (set BitSet) ClearBit( i uint) {
	set[i/64]&=(1<<(i%64))^math.MaxUint64
}

// Clear entire BitSet
func (set BitSet) Clear() {
	for i,_ := range set {
		set[i] = 0
	}
}
<commit_msg>improved ClearBit thanks to zhai<commit_after>// Copyright 2011 Will Fitzgerald. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
	Package bitset implements bitsets.

	It provides methods for making a BitSet of an arbitrary
	upper limit, setting and testing bit locations, and clearing
	bit locations as well as the entire set.

	BitSets are implemented as arrays of uint64s, so it may be best
	to limit the upper size to a multiple of 64. It is not an error
	to set, test, or clear a bit within a 64 bit boundary.

	Example use:

	b := MakeBitSet(64000)
	b.SetBit(1000)
	if b.Bit(1000) {
		b.ClearBit(1000)
	}
*/
package bitset

// BitSets are arrays of uint64.
type BitSet []uint64

// Make a BitSet with an upper limit on size. Note this is the
// number of bits, not the number of uint64s, which is a kind of
// implementation detail.
func MakeBitSet(max_size uint) BitSet {
	return make(BitSet, (max_size+(64-1))/64)
}

/// Test whether bit i is set.
func (set BitSet) Bit(i uint) bool {
	return ((set[i/64] & (1<<(i%64)))!=0)
}

// Set bit i to 1
func (set BitSet) SetBit(i uint) {
	set[i/64]|=(1<<(i%64))
}

// Clear bit i to 0
func (set BitSet) ClearBit(i uint) {
	set[i/64] &^= 1 << (i % 64)
}

// Clear entire BitSet
func (set BitSet) Clear() {
	for i,_ := range set {
		set[i] = 0
	}
}
<|endoftext|>


<commit_before>package gokeepasslib

import (
	"errors"
	"fmt"
	"encoding/binary"
	"crypto/sha256"
	"bytes"
)

//Size in bytes of the data in each block
const blockSplitRate = 16384

type ErrCorruptBlock struct {
	block *Block
	Hash [32]byte
}
func (e ErrCorruptBlock) Error () string {
	return fmt.Sprintf("Hash mismatch at block %d. Should be %x, is %x",e.block.Index,e.Hash,e.block.Hash)
}

var ErrBlockTooSmall = errors.New("gokeepasslib: block is too small")
var ErrEmptyBlock = errors.New("gokeepasslib: block appears to be empty/clossing block")

const minsize int = 40

type Block struct {
	Index  uint32
	Hash   [32]byte
	Length uint32
	Body   []byte
}
func NewBlock (i uint32, body []byte) (b Block) {
	b.Index = i
	b.Body = body
	b.Hash = sha256.Sum256(b.Body)
	b.Length = uint32(len(b.Body))
	return
}
func (u Block) MarshalText () (text []byte,err error) {
	text = make([]byte,u.Size())
	binary.LittleEndian.PutUint32(text[:4],u.Index)
	copy(text[4:36],u.Hash[:])
	binary.LittleEndian.PutUint32(text[36:40],u.Length)
	copy(text[40:],u.Body)
	return
}
func (b *Block) UnmarshalText(text []byte) error {
	if len (text) < minsize {
		return ErrBlockTooSmall
	}
	b.Index = binary.LittleEndian.Uint32(text[:4])
	copy(b.Hash[:],text[4:36])
	b.Length = binary.LittleEndian.Uint32(text[36:40])
	if b.Length == 0 {
		return ErrEmptyBlock
	}
	if len(text)-40 < int(b.Length) {
		return ErrBlockTooSmall
	}
	b.Body = text[40:40+b.Length]
	calculatedHash := sha256.Sum256(b.Body)
	if !bytes.Equal(calculatedHash[:], b.Hash[:]) {
		return ErrCorruptBlock{b,calculatedHash}
	}
	return nil
}
func (b Block) Size () int {
	return 40+len(b.Body)
}
func (b Block) String () string {
	return fmt.Sprintf("Index: %d\nHash: %x\n,Length:%d\n\n",b.Index,b.Hash,b.Length)
}

func DecodeBlocks (body []byte) (result []byte,err error) {
	var block Block
	for len(body) > 0 {
		err = block.UnmarshalText(body)
		if err != nil {
			if err == ErrEmptyBlock {
				return result,nil
			}
			return result,err
		}
		result = append(result, block.Body...)
		body = body[block.Size():]
	}
	return
}

// Converts raw xml data to keepass's block format, which includes a hash of each block to check for data corruption,
// Every block contains the following elements:
// (4 bytes) ID : an unique interger id for this block
// (32 bytes) sha-256 hash of block data
// (4 bytes) size on bytes of the block data
// (Data Size Bytes) the actual xml data of the block, will be blockSplitRate bytes at most
func EncodeBlocks (data []byte) (result []byte,err error) {
	var i uint32
	for len(data) > 0 { //For each block
		var body []byte
		if len(data) >= blockSplitRate { //If there is enough data for another block, use blockSplitRate bytes of data for block
			body = append(body,data[:blockSplitRate]...)
		} else { //Otherwise just use what is remaining and clear data to break from the loop
			body = append(body,data...)
		}
		data = data[len(body):]
		block := NewBlock(i,body)
		content,_ := block.MarshalText()
		result = append(result,content...)
		i++
	}
	endblock,_ := Block{Index:i}.MarshalText()
	result = append(result,endblock...)
	return
}<commit_msg>comments for blocks.go<commit_after>package gokeepasslib

import (
	"errors"
	"fmt"
	"encoding/binary"
	"crypto/sha256"
	"bytes"
)

// Size in bytes of the data in each block
const blockSplitRate = 16384

// ErrCorruptBlock is returned by Decoder when a block in the kdbx file has a different hash from the correct hash
type ErrCorruptBlock struct {
	block *Block
	Hash [32]byte
}
func (e ErrCorruptBlock) Error () string {
	return fmt.Sprintf("Hash mismatch at block %d. Should be %x, is %x",e.block.Index,e.Hash,e.block.Hash)
}

// ErrBlockTooSmall is returned by decoder if the block length does not match the body size
var ErrBlockTooSmall = errors.New("gokeepasslib: block is too small")

// ErrEmptyBlock is returned by decoder when a block is empty,ignored if is the last block in the file
var ErrEmptyBlock = errors.New("gokeepasslib: block appears to be empty/clossing block")

//minsize is the smallest size a block can be
const minsize int = 40

// Block stores an individual block following the kdbx block format
type Block struct {
	Index  uint32
	Hash   [32]byte
	Length uint32
	Body   []byte
}

// NewBlock creates a block with index i and body b, adding length and hash automaticaly
func NewBlock (i uint32, b []byte) (b Block) {
	b.Index = i
	b.Body = b
	b.Hash = sha256.Sum256(b.Body)
	b.Length = uint32(len(b.Body))
	return
}

// MarshalText converts a block into kdbx binary block format
func (u Block) MarshalText () (text []byte,err error) {
	text = make([]byte,u.Size())
	binary.LittleEndian.PutUint32(text[:4],u.Index)
	copy(text[4:36],u.Hash[:])
	binary.LittleEndian.PutUint32(text[36:40],u.Length)
	copy(text[40:],u.Body)
	return
}

// UnmarshalText turns text in binary kdbx format into a block
func (b *Block) UnmarshalText(text []byte) error {
	if len (text) < minsize {
		return ErrBlockTooSmall
	}
	b.Index = binary.LittleEndian.Uint32(text[:4])
	copy(b.Hash[:],text[4:36])
	b.Length = binary.LittleEndian.Uint32(text[36:40])
	if b.Length == 0 {
		return ErrEmptyBlock
	}
	if len(text)-40 < int(b.Length) {
		return ErrBlockTooSmall
	}
	b.Body = text[40:40+b.Length]
	calculatedHash := sha256.Sum256(b.Body)
	if !bytes.Equal(calculatedHash[:], b.Hash[:]) {
		return ErrCorruptBlock{b,calculatedHash}
	}
	return nil
}

// Size returns the number of bytes b will occupy when Marshaled to binary
func (b Block) Size () int {
	return 40+len(b.Body)
}
func (b Block) String () string {
	return fmt.Sprintf("Index: %d\nHash: %x\n,Length:%d\n\n",b.Index,b.Hash,b.Length)
}

// DecodeBlocks converts a []byte in kdbx block format into the xml content, checking for corruption
func DecodeBlocks (body []byte) (result []byte,err error) {
	var block Block
	for len(body) > 0 {
		err = block.UnmarshalText(body)
		if err != nil {
			if err == ErrEmptyBlock {
				return result,nil
			}
			return result,err
		}
		result = append(result, block.Body...)
		body = body[block.Size():]
	}
	return
}

// EncodeBlocks Converts raw xml data to keepass's block format, which includes a hash of each block to check for data corruption,
// Every block contains the following elements:
// (4 bytes) ID : an unique interger id for this block
// (32 bytes) sha-256 hash of block data
// (4 bytes) size on bytes of the block data
// (Data Size Bytes) the actual xml data of the block, will be blockSplitRate bytes at most
func EncodeBlocks (data []byte) (result []byte,err error) {
	var i uint32
	for len(data) > 0 { //For each block
		var body []byte
		if len(data) >= blockSplitRate { //If there is enough data for another block, use blockSplitRate bytes of data for block
			body = append(body,data[:blockSplitRate]...)
		} else { //Otherwise just use what is remaining and clear data to break from the loop
			body = append(body,data...)
		}
		data = data[len(body):]
		block := NewBlock(i,body)
		content,_ := block.MarshalText()
		result = append(result,content...)
		i++
	}
	endblock,_ := Block{Index:i}.MarshalText()
	result = append(result,endblock...)
	return
}<|endoftext|>


<commit_before>package k8s_test

import (
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/onsi/gomega/gbytes"
	"github.com/square/certstrap/pkix"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Web HTTP or HTTPS(TLS) termination at web node", func() {

	var (
		serverCertBytes []byte
		serverKeyBytes  []byte
		caCertFile      *os.File
	)

	BeforeEach(func() {
		var err error

		setReleaseNameAndNamespace("wtt")

		CACert, serverKey, serverCert := generateKeyPairWithCA(namespace, releaseName+"-web")
		CACertBytes, err := CACert.Export()
		Expect(err).NotTo(HaveOccurred())

		caCertFile, err = ioutil.TempFile("", "ca")
		caCertFile.Write(CACertBytes)
		caCertFile.Close()

		serverKeyBytes, err = serverKey.ExportPrivate()
		Expect(err).NotTo(HaveOccurred())

		serverCertBytes, err = serverCert.Export()
		Expect(err).NotTo(HaveOccurred())
	})

	AfterEach(func() {
		os.Remove(caCertFile.Name())
	})

	Context("when configured correctly", func() {

		var (
			atc Endpoint

			chartConfig []string
			proxyPort   string
		)

		JustBeforeEach(func() {
			deployConcourseChart(releaseName, chartConfig...)

			waitAllPodsInNamespaceToBeReady(namespace)

			atc = endpointFactory.NewServiceEndpoint(
				namespace,
				releaseName+"-web",
				proxyPort,
			)
		})

		AfterEach(func() {
			atc.Close()
			cleanup(releaseName, namespace)
		})

		Context("with tls termination at web", func() {

			BeforeEach(func() {
				chartConfig = generateChartConfig(
					"--set=concourse.web.externalUrl=https://test.com",
					"--set=concourse.web.tls.enabled=true",
					"--set=secrets.webTlsCert="+string(serverCertBytes),
					"--set=secrets.webTlsKey="+string(serverKeyBytes),
				)

				proxyPort = "443"
			})

			It("fly login succeeds when using the correct CA and host", func() {
				fly.Login(
					"test",
					"test",
					"https://"+atc.Address(),
					"--ca-cert", caCertFile.Name(),
				)
			})

			It("fly login fails when NOT using the correct CA", func() {
				sess := fly.Start("login", "-u", "test", "-p", "test",
					"--ca-cert", "k8s/certs/wrong-ca.crt",
					"-c", "https://"+atc.Address(),
				)
				<-sess.Exited

				Expect(sess.ExitCode()).ToNot(Equal(0))
				Expect(sess.Err).To(gbytes.Say(`x509: certificate signed by unknown authority`))
			})
		})

	})

	Context("When NOT configured correctly", func() {

		BeforeEach(func() {
			setReleaseNameAndNamespace("wtt")
		})

		It("fails if tls is enabled but externalURL is NOT set", func() {
			expectedErr := "Must specify HTTPS external URL when concourse.web.tls.enabled is true"

			chartConfig := generateChartConfig(
				"--set=concourse.web.tls.enabled=true",
				"--set=secrets.webTlsCert="+string(serverCertBytes),
				"--set=secrets.webTlsKey="+string(serverKeyBytes),
			)

			deployFailingConcourseChart(releaseName, expectedErr,
				chartConfig...,
			)
		})

		It("fails when tls is enabled but ssl cert and ssl key are NOT set", func() {
			expectedErr := "secrets.webTlsCert is required because secrets.create is true and concourse.web.tls.enabled is true"

			chartConfig := generateChartConfig("--set=concourse.web.externalUrl=https://test.com",
				"--set=concourse.web.tls.enabled=true")

			deployFailingConcourseChart(releaseName, expectedErr,
				chartConfig...,
			)
		})
	})

})

func generateChartConfig(args ...string) []string {
	return append(args,
		"--set=worker.enabled=false",
		"--set=concourse.worker.baggageclaim.driver=detect",
		"--set=concourse.web.tls.bindPort=443",
	)
}
func generateKeyPairWithCA(namespace, service string) (*pkix.Certificate, *pkix.Key, *pkix.Certificate) {
	CAKey, err := pkix.CreateRSAKey(1024)
	Expect(err).NotTo(HaveOccurred())

	CACert, err := pkix.CreateCertificateAuthority(CAKey, "", time.Now().Add(time.Hour), "Pivotal", "", "", "", "CA")
	Expect(err).NotTo(HaveOccurred())

	serverKey, err := pkix.CreateRSAKey(1024)
	Expect(err).NotTo(HaveOccurred())

	certificateSigningRequest, err := pkix.CreateCertificateSigningRequest(
		serverKey, "", []net.IP{net.IPv4(127, 0, 0, 1)},
		[]string{serviceAddress(namespace, service)},
		"Pivotal", "", "", "", "127.0.0.1")
	Expect(err).NotTo(HaveOccurred())

	serverCert, err := pkix.CreateCertificateHost(CACert, CAKey,
		certificateSigningRequest, time.Now().Add(time.Hour))
	Expect(err).NotTo(HaveOccurred())

	return CACert, serverKey, serverCert
}
<commit_msg>retry login with incorrect CA<commit_after>package k8s_test

import (
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/onsi/gomega/gbytes"
	"github.com/square/certstrap/pkix"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Web HTTP or HTTPS(TLS) termination at web node", func() {

	var (
		serverCertBytes []byte
		serverKeyBytes  []byte
		caCertFile      *os.File
	)

	BeforeEach(func() {
		var err error

		setReleaseNameAndNamespace("wtt")

		CACert, serverKey, serverCert := generateKeyPairWithCA(namespace, releaseName+"-web")
		CACertBytes, err := CACert.Export()
		Expect(err).NotTo(HaveOccurred())

		caCertFile, err = ioutil.TempFile("", "ca")
		caCertFile.Write(CACertBytes)
		caCertFile.Close()

		serverKeyBytes, err = serverKey.ExportPrivate()
		Expect(err).NotTo(HaveOccurred())

		serverCertBytes, err = serverCert.Export()
		Expect(err).NotTo(HaveOccurred())
	})

	AfterEach(func() {
		os.Remove(caCertFile.Name())
	})

	Context("when configured correctly", func() {

		var (
			atc Endpoint

			chartConfig []string
			proxyPort   string
		)

		JustBeforeEach(func() {
			deployConcourseChart(releaseName, chartConfig...)

			waitAllPodsInNamespaceToBeReady(namespace)

			atc = endpointFactory.NewServiceEndpoint(
				namespace,
				releaseName+"-web",
				proxyPort,
			)
		})

		AfterEach(func() {
			atc.Close()
			cleanup(releaseName, namespace)
		})

		Context("with tls termination at web", func() {

			BeforeEach(func() {
				chartConfig = generateChartConfig(
					"--set=concourse.web.externalUrl=https://test.com",
					"--set=concourse.web.tls.enabled=true",
					"--set=secrets.webTlsCert="+string(serverCertBytes),
					"--set=secrets.webTlsKey="+string(serverKeyBytes),
				)

				proxyPort = "443"
			})

			It("fly login succeeds when using the correct CA and host", func() {
				fly.Login(
					"test",
					"test",
					"https://"+atc.Address(),
					"--ca-cert", caCertFile.Name(),
				)
			})

			It("fly login fails when NOT using the correct CA", func() {
				Eventually(func() *gbytes.Buffer {
					sess := fly.Start("login", "-u", "test", "-p", "test",
						"--ca-cert", "k8s/certs/wrong-ca.crt",
						"-c", "https://"+atc.Address(),
					)
					<-sess.Exited
					return sess.Err
				}, 2*time.Minute, 10*time.Second).
					Should(gbytes.Say(`x509: certificate signed by unknown authority`))
			})
		})

	})

	Context("When NOT configured correctly", func() {

		BeforeEach(func() {
			setReleaseNameAndNamespace("wtt")
		})

		It("fails if tls is enabled but externalURL is NOT set", func() {
			expectedErr := "Must specify HTTPS external URL when concourse.web.tls.enabled is true"

			chartConfig := generateChartConfig(
				"--set=concourse.web.tls.enabled=true",
				"--set=secrets.webTlsCert="+string(serverCertBytes),
				"--set=secrets.webTlsKey="+string(serverKeyBytes),
			)

			deployFailingConcourseChart(releaseName, expectedErr,
				chartConfig...,
			)
		})

		It("fails when tls is enabled but ssl cert and ssl key are NOT set", func() {
			expectedErr := "secrets.webTlsCert is required because secrets.create is true and concourse.web.tls.enabled is true"

			chartConfig := generateChartConfig("--set=concourse.web.externalUrl=https://test.com",
				"--set=concourse.web.tls.enabled=true")

			deployFailingConcourseChart(releaseName, expectedErr,
				chartConfig...,
			)
		})
	})

})

func generateChartConfig(args ...string) []string {
	return append(args,
		"--set=worker.enabled=false",
		"--set=concourse.worker.baggageclaim.driver=detect",
		"--set=concourse.web.tls.bindPort=443",
	)
}
func generateKeyPairWithCA(namespace, service string) (*pkix.Certificate, *pkix.Key, *pkix.Certificate) {
	CAKey, err := pkix.CreateRSAKey(1024)
	Expect(err).NotTo(HaveOccurred())

	CACert, err := pkix.CreateCertificateAuthority(CAKey, "", time.Now().Add(time.Hour), "Pivotal", "", "", "", "CA")
	Expect(err).NotTo(HaveOccurred())

	serverKey, err := pkix.CreateRSAKey(1024)
	Expect(err).NotTo(HaveOccurred())

	certificateSigningRequest, err := pkix.CreateCertificateSigningRequest(
		serverKey, "", []net.IP{net.IPv4(127, 0, 0, 1)},
		[]string{serviceAddress(namespace, service)},
		"Pivotal", "", "", "", "127.0.0.1")
	Expect(err).NotTo(HaveOccurred())

	serverCert, err := pkix.CreateCertificateHost(CACert, CAKey,
		certificateSigningRequest, time.Now().Add(time.Hour))
	Expect(err).NotTo(HaveOccurred())

	return CACert, serverKey, serverCert
}
<|endoftext|>


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

package hyperstart

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math"
	"net"
	"sync"
	"time"

	hyper "github.com/hyperhq/runv/hyperstart/api/json"
)

// Control command IDs
// Need to be in sync with hyperstart/src/api.h
const (
	Version          = "version"
	StartPod         = "startpod"
	GetPod           = "getpod"
	DestroyPod       = "destroypod"
	RestartContainer = "restartcontainer"
	ExecCmd          = "execcmd"
	FinishCmd        = "finishcmd"
	Ready            = "ready"
	Ack              = "ack"
	Error            = "error"
	WinSize          = "winsize"
	Ping             = "ping"
	FinishPod        = "finishpod"
	Next             = "next"
	WriteFile        = "writefile"
	ReadFile         = "readfile"
	NewContainer     = "newcontainer"
	KillContainer    = "killcontainer"
	OnlineCPUMem     = "onlinecpumem"
	SetupInterface   = "setupinterface"
	SetupRoute       = "setuproute"
)

var codeList = map[string]uint32{
	Version:          hyper.INIT_VERSION,
	StartPod:         hyper.INIT_STARTPOD,
	GetPod:           hyper.INIT_GETPOD,
	DestroyPod:       hyper.INIT_DESTROYPOD,
	RestartContainer: hyper.INIT_RESTARTCONTAINER,
	ExecCmd:          hyper.INIT_EXECCMD,
	FinishCmd:        hyper.INIT_FINISHCMD,
	Ready:            hyper.INIT_READY,
	Ack:              hyper.INIT_ACK,
	Error:            hyper.INIT_ERROR,
	WinSize:          hyper.INIT_WINSIZE,
	Ping:             hyper.INIT_PING,
	Next:             hyper.INIT_NEXT,
	WriteFile:        hyper.INIT_WRITEFILE,
	ReadFile:         hyper.INIT_READFILE,
	NewContainer:     hyper.INIT_NEWCONTAINER,
	KillContainer:    hyper.INIT_KILLCONTAINER,
	OnlineCPUMem:     hyper.INIT_ONLINECPUMEM,
	SetupInterface:   hyper.INIT_SETUPINTERFACE,
	SetupRoute:       hyper.INIT_SETUPROUTE,
}

// Values related to the communication on control channel.
const (
	ctlHdrSize      = 8
	ctlHdrLenOffset = 4
)

// Values related to the communication on tty channel.
const (
	ttyHdrSize      = 12
	ttyHdrLenOffset = 8
)

// Hyperstart is the base structure for hyperstart.
type Hyperstart struct {
	ctlSerial, ioSerial string
	sockType            string
	ctl, io             net.Conn

	// ctl access is arbitrated by ctlMutex. We can only allow a single
	// "transaction" (write command + read answer) at a time
	ctlMutex sync.Mutex
}

// NewHyperstart returns a new hyperstart structure.
func NewHyperstart(ctlSerial, ioSerial, sockType string) *Hyperstart {
	return &Hyperstart{
		ctlSerial: ctlSerial,
		ioSerial:  ioSerial,
		sockType:  sockType,
	}
}

// OpenSockets opens both CTL and IO sockets.
func (h *Hyperstart) OpenSockets() error {
	var err error

	h.ctl, err = net.Dial(h.sockType, h.ctlSerial)
	if err != nil {
		return err
	}

	h.io, err = net.Dial(h.sockType, h.ioSerial)
	if err != nil {
		h.ctl.Close()
		return err
	}

	return nil
}

// CloseSockets closes both CTL and IO sockets.
func (h *Hyperstart) CloseSockets() error {
	if h.ctl != nil {
		err := h.ctl.Close()
		if err != nil {
			return err
		}

		h.ctl = nil
	}

	if h.io != nil {
		err := h.io.Close()
		if err != nil {
			return err
		}

		h.io = nil
	}

	return nil
}

// SetDeadline sets a timeout for CTL connection.
func (h *Hyperstart) SetDeadline(t time.Time) error {
	err := h.ctl.SetDeadline(t)
	if err != nil {
		return err
	}

	return nil
}

// IsStarted returns about connection status.
func (h *Hyperstart) IsStarted() bool {
	ret := false
	timeoutDuration := 1 * time.Second

	if h.ctl == nil {
		return ret
	}

	h.SetDeadline(time.Now().Add(timeoutDuration))

	_, err := h.SendCtlMessage(Ping, nil)
	if err == nil {
		ret = true
	}

	h.SetDeadline(time.Time{})

	if ret == false {
		h.CloseSockets()
	}

	return ret
}

// FormatMessage formats hyperstart messages.
func FormatMessage(payload interface{}) ([]byte, error) {
	var payloadSlice []byte
	var err error

	if payload != nil {
		switch p := payload.(type) {
		case string:
			payloadSlice = []byte(p)
		default:
			payloadSlice, err = json.Marshal(p)
			if err != nil {
				return nil, err
			}
		}
	}

	return payloadSlice, nil
}

func (h *Hyperstart) readCtlMessage() (*hyper.DecodedMessage, error) {
	needRead := ctlHdrSize
	length := 0
	read := 0
	buf := make([]byte, 512)
	res := []byte{}
	for read < needRead {
		want := needRead - read
		if want > 512 {
			want = 512
		}
		nr, err := h.ctl.Read(buf[:want])
		if err != nil {
			return nil, err
		}

		res = append(res, buf[:nr]...)
		read = read + nr

		if length == 0 && read >= ctlHdrSize {
			length = int(binary.BigEndian.Uint32(res[ctlHdrLenOffset:ctlHdrSize]))
			if length > ctlHdrSize {
				needRead = length
			}
		}
	}

	return &hyper.DecodedMessage{
		Code:    binary.BigEndian.Uint32(res[:ctlHdrLenOffset]),
		Message: res[ctlHdrSize:],
	}, nil
}

func (h *Hyperstart) writeCtlMessage(m *hyper.DecodedMessage) error {
	length := len(m.Message) + ctlHdrSize
	// XXX: Support sending messages by chunks to support messages over
	// 10240 bytes. That limit is from hyperstart src/init.c,
	// hyper_channel_ops, rbuf_size.
	if length > 10240 {
		return fmt.Errorf("message too long %d", length)
	}
	msg := make([]byte, length)
	binary.BigEndian.PutUint32(msg[:], uint32(m.Code))
	binary.BigEndian.PutUint32(msg[ctlHdrLenOffset:], uint32(length))
	copy(msg[ctlHdrSize:], m.Message)

	_, err := h.ctl.Write(msg)
	if err != nil {
		return err
	}

	return nil
}

// ReadIoMessageWithConn returns data coming from the specified IO channel.
func ReadIoMessageWithConn(conn net.Conn) (*hyper.TtyMessage, error) {
	needRead := ttyHdrSize
	length := 0
	read := 0
	buf := make([]byte, 512)
	res := []byte{}
	for read < needRead {
		want := needRead - read
		if want > 512 {
			want = 512
		}
		nr, err := conn.Read(buf[:want])
		if err != nil {
			return nil, err
		}

		res = append(res, buf[:nr]...)
		read = read + nr

		if length == 0 && read >= ttyHdrSize {
			length = int(binary.BigEndian.Uint32(res[ttyHdrLenOffset:ttyHdrSize]))
			if length > ttyHdrSize {
				needRead = length
			}
		}
	}

	return &hyper.TtyMessage{
		Session: binary.BigEndian.Uint64(res[:ttyHdrLenOffset]),
		Message: res[ttyHdrSize:],
	}, nil
}

// ReadIoMessage returns data coming from the IO channel.
func (h *Hyperstart) ReadIoMessage() (*hyper.TtyMessage, error) {
	return ReadIoMessageWithConn(h.io)
}

// SendIoMessageWithConn sends data to the specified IO channel.
func SendIoMessageWithConn(conn net.Conn, ttyMsg *hyper.TtyMessage) error {
	length := len(ttyMsg.Message) + ttyHdrSize
	// XXX: Support sending messages by chunks to support messages over
	// 10240 bytes. That limit is from hyperstart src/init.c,
	// hyper_channel_ops, rbuf_size.
	if length > 10240 {
		return fmt.Errorf("message too long %d", length)
	}
	msg := make([]byte, length)
	binary.BigEndian.PutUint64(msg[:], ttyMsg.Session)
	binary.BigEndian.PutUint32(msg[ttyHdrLenOffset:], uint32(length))
	copy(msg[ttyHdrSize:], ttyMsg.Message)

	n, err := conn.Write(msg)
	if err != nil {
		return err
	}

	if n != length {
		return fmt.Errorf("%d bytes written out of %d expected", n, length)
	}

	return nil
}

// SendIoMessage sends data to the IO channel.
func (h *Hyperstart) SendIoMessage(ttyMsg *hyper.TtyMessage) error {
	return SendIoMessageWithConn(h.io, ttyMsg)
}

func codeFromCmd(cmd string) (uint32, error) {
	_, ok := codeList[cmd]
	if ok == false {
		return math.MaxUint32, fmt.Errorf("unknown command '%s'", cmd)
	}

	return codeList[cmd], nil
}

func (h *Hyperstart) expectReadingCmd(code uint32) (*hyper.DecodedMessage, error) {
	var msg *hyper.DecodedMessage
	var err error

	for {
		msg, err = h.readCtlMessage()
		if err != nil {
			return nil, nil
		}

		if msg.Code == code {
			break
		}

		if msg.Code == hyper.INIT_NEXT || msg.Code == hyper.INIT_READY {
			continue
		}

		if msg.Code != code {
			if msg.Code == hyper.INIT_ERROR {
				return nil, fmt.Errorf("ERROR received from Hyperstart")
			}

			return nil, fmt.Errorf("CMD ID received %d not matching expected %d", msg.Code, code)
		}
	}

	return msg, nil
}

// WaitForReady waits for a READY message on CTL channel.
func (h *Hyperstart) WaitForReady() error {
	h.ctlMutex.Lock()
	defer h.ctlMutex.Unlock()

	_, err := h.expectReadingCmd(hyper.INIT_READY)
	if err != nil {
		return err
	}

	return nil
}

// SendCtlMessage sends a message to the CTL channel.
func (h *Hyperstart) SendCtlMessage(cmd string, data []byte) (*hyper.DecodedMessage, error) {
	h.ctlMutex.Lock()
	defer h.ctlMutex.Unlock()

	code, err := codeFromCmd(cmd)
	if err != nil {
		return nil, err
	}

	// Write message
	msg := &hyper.DecodedMessage{
		Code:    code,
		Message: data,
	}
	err = h.writeCtlMessage(msg)
	if err != nil {
		return nil, err
	}

	// Wait for answer
	resp, err := h.expectReadingCmd(hyper.INIT_ACK)
	if err != nil {
		return nil, err
	}

	return resp, nil
}
<commit_msg>hyperstart: Don't reset the connections to nil in CloseSockets()<commit_after>//
// Copyright (c) 2016 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package hyperstart

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math"
	"net"
	"sync"
	"time"

	hyper "github.com/hyperhq/runv/hyperstart/api/json"
)

// Control command IDs
// Need to be in sync with hyperstart/src/api.h
const (
	Version          = "version"
	StartPod         = "startpod"
	GetPod           = "getpod"
	DestroyPod       = "destroypod"
	RestartContainer = "restartcontainer"
	ExecCmd          = "execcmd"
	FinishCmd        = "finishcmd"
	Ready            = "ready"
	Ack              = "ack"
	Error            = "error"
	WinSize          = "winsize"
	Ping             = "ping"
	FinishPod        = "finishpod"
	Next             = "next"
	WriteFile        = "writefile"
	ReadFile         = "readfile"
	NewContainer     = "newcontainer"
	KillContainer    = "killcontainer"
	OnlineCPUMem     = "onlinecpumem"
	SetupInterface   = "setupinterface"
	SetupRoute       = "setuproute"
)

var codeList = map[string]uint32{
	Version:          hyper.INIT_VERSION,
	StartPod:         hyper.INIT_STARTPOD,
	GetPod:           hyper.INIT_GETPOD,
	DestroyPod:       hyper.INIT_DESTROYPOD,
	RestartContainer: hyper.INIT_RESTARTCONTAINER,
	ExecCmd:          hyper.INIT_EXECCMD,
	FinishCmd:        hyper.INIT_FINISHCMD,
	Ready:            hyper.INIT_READY,
	Ack:              hyper.INIT_ACK,
	Error:            hyper.INIT_ERROR,
	WinSize:          hyper.INIT_WINSIZE,
	Ping:             hyper.INIT_PING,
	Next:             hyper.INIT_NEXT,
	WriteFile:        hyper.INIT_WRITEFILE,
	ReadFile:         hyper.INIT_READFILE,
	NewContainer:     hyper.INIT_NEWCONTAINER,
	KillContainer:    hyper.INIT_KILLCONTAINER,
	OnlineCPUMem:     hyper.INIT_ONLINECPUMEM,
	SetupInterface:   hyper.INIT_SETUPINTERFACE,
	SetupRoute:       hyper.INIT_SETUPROUTE,
}

// Values related to the communication on control channel.
const (
	ctlHdrSize      = 8
	ctlHdrLenOffset = 4
)

// Values related to the communication on tty channel.
const (
	ttyHdrSize      = 12
	ttyHdrLenOffset = 8
)

// Hyperstart is the base structure for hyperstart.
type Hyperstart struct {
	ctlSerial, ioSerial string
	sockType            string
	ctl, io             net.Conn

	// ctl access is arbitrated by ctlMutex. We can only allow a single
	// "transaction" (write command + read answer) at a time
	ctlMutex sync.Mutex
}

// NewHyperstart returns a new hyperstart structure.
func NewHyperstart(ctlSerial, ioSerial, sockType string) *Hyperstart {
	return &Hyperstart{
		ctlSerial: ctlSerial,
		ioSerial:  ioSerial,
		sockType:  sockType,
	}
}

// OpenSockets opens both CTL and IO sockets.
func (h *Hyperstart) OpenSockets() error {
	var err error

	h.ctl, err = net.Dial(h.sockType, h.ctlSerial)
	if err != nil {
		return err
	}

	h.io, err = net.Dial(h.sockType, h.ioSerial)
	if err != nil {
		h.ctl.Close()
		return err
	}

	return nil
}

// CloseSockets closes both CTL and IO sockets.
func (h *Hyperstart) CloseSockets() error {
	if h.ctl != nil {
		err := h.ctl.Close()
		if err != nil {
			return err
		}
	}

	if h.io != nil {
		err := h.io.Close()
		if err != nil {
			return err
		}
	}

	return nil
}

// SetDeadline sets a timeout for CTL connection.
func (h *Hyperstart) SetDeadline(t time.Time) error {
	err := h.ctl.SetDeadline(t)
	if err != nil {
		return err
	}

	return nil
}

// IsStarted returns about connection status.
func (h *Hyperstart) IsStarted() bool {
	ret := false
	timeoutDuration := 1 * time.Second

	if h.ctl == nil {
		return ret
	}

	h.SetDeadline(time.Now().Add(timeoutDuration))

	_, err := h.SendCtlMessage(Ping, nil)
	if err == nil {
		ret = true
	}

	h.SetDeadline(time.Time{})

	if ret == false {
		h.CloseSockets()
	}

	return ret
}

// FormatMessage formats hyperstart messages.
func FormatMessage(payload interface{}) ([]byte, error) {
	var payloadSlice []byte
	var err error

	if payload != nil {
		switch p := payload.(type) {
		case string:
			payloadSlice = []byte(p)
		default:
			payloadSlice, err = json.Marshal(p)
			if err != nil {
				return nil, err
			}
		}
	}

	return payloadSlice, nil
}

func (h *Hyperstart) readCtlMessage() (*hyper.DecodedMessage, error) {
	needRead := ctlHdrSize
	length := 0
	read := 0
	buf := make([]byte, 512)
	res := []byte{}
	for read < needRead {
		want := needRead - read
		if want > 512 {
			want = 512
		}
		nr, err := h.ctl.Read(buf[:want])
		if err != nil {
			return nil, err
		}

		res = append(res, buf[:nr]...)
		read = read + nr

		if length == 0 && read >= ctlHdrSize {
			length = int(binary.BigEndian.Uint32(res[ctlHdrLenOffset:ctlHdrSize]))
			if length > ctlHdrSize {
				needRead = length
			}
		}
	}

	return &hyper.DecodedMessage{
		Code:    binary.BigEndian.Uint32(res[:ctlHdrLenOffset]),
		Message: res[ctlHdrSize:],
	}, nil
}

func (h *Hyperstart) writeCtlMessage(m *hyper.DecodedMessage) error {
	length := len(m.Message) + ctlHdrSize
	// XXX: Support sending messages by chunks to support messages over
	// 10240 bytes. That limit is from hyperstart src/init.c,
	// hyper_channel_ops, rbuf_size.
	if length > 10240 {
		return fmt.Errorf("message too long %d", length)
	}
	msg := make([]byte, length)
	binary.BigEndian.PutUint32(msg[:], uint32(m.Code))
	binary.BigEndian.PutUint32(msg[ctlHdrLenOffset:], uint32(length))
	copy(msg[ctlHdrSize:], m.Message)

	_, err := h.ctl.Write(msg)
	if err != nil {
		return err
	}

	return nil
}

// ReadIoMessageWithConn returns data coming from the specified IO channel.
func ReadIoMessageWithConn(conn net.Conn) (*hyper.TtyMessage, error) {
	needRead := ttyHdrSize
	length := 0
	read := 0
	buf := make([]byte, 512)
	res := []byte{}
	for read < needRead {
		want := needRead - read
		if want > 512 {
			want = 512
		}
		nr, err := conn.Read(buf[:want])
		if err != nil {
			return nil, err
		}

		res = append(res, buf[:nr]...)
		read = read + nr

		if length == 0 && read >= ttyHdrSize {
			length = int(binary.BigEndian.Uint32(res[ttyHdrLenOffset:ttyHdrSize]))
			if length > ttyHdrSize {
				needRead = length
			}
		}
	}

	return &hyper.TtyMessage{
		Session: binary.BigEndian.Uint64(res[:ttyHdrLenOffset]),
		Message: res[ttyHdrSize:],
	}, nil
}

// ReadIoMessage returns data coming from the IO channel.
func (h *Hyperstart) ReadIoMessage() (*hyper.TtyMessage, error) {
	return ReadIoMessageWithConn(h.io)
}

// SendIoMessageWithConn sends data to the specified IO channel.
func SendIoMessageWithConn(conn net.Conn, ttyMsg *hyper.TtyMessage) error {
	length := len(ttyMsg.Message) + ttyHdrSize
	// XXX: Support sending messages by chunks to support messages over
	// 10240 bytes. That limit is from hyperstart src/init.c,
	// hyper_channel_ops, rbuf_size.
	if length > 10240 {
		return fmt.Errorf("message too long %d", length)
	}
	msg := make([]byte, length)
	binary.BigEndian.PutUint64(msg[:], ttyMsg.Session)
	binary.BigEndian.PutUint32(msg[ttyHdrLenOffset:], uint32(length))
	copy(msg[ttyHdrSize:], ttyMsg.Message)

	n, err := conn.Write(msg)
	if err != nil {
		return err
	}

	if n != length {
		return fmt.Errorf("%d bytes written out of %d expected", n, length)
	}

	return nil
}

// SendIoMessage sends data to the IO channel.
func (h *Hyperstart) SendIoMessage(ttyMsg *hyper.TtyMessage) error {
	return SendIoMessageWithConn(h.io, ttyMsg)
}

func codeFromCmd(cmd string) (uint32, error) {
	_, ok := codeList[cmd]
	if ok == false {
		return math.MaxUint32, fmt.Errorf("unknown command '%s'", cmd)
	}

	return codeList[cmd], nil
}

func (h *Hyperstart) expectReadingCmd(code uint32) (*hyper.DecodedMessage, error) {
	var msg *hyper.DecodedMessage
	var err error

	for {
		msg, err = h.readCtlMessage()
		if err != nil {
			return nil, nil
		}

		if msg.Code == code {
			break
		}

		if msg.Code == hyper.INIT_NEXT || msg.Code == hyper.INIT_READY {
			continue
		}

		if msg.Code != code {
			if msg.Code == hyper.INIT_ERROR {
				return nil, fmt.Errorf("ERROR received from Hyperstart")
			}

			return nil, fmt.Errorf("CMD ID received %d not matching expected %d", msg.Code, code)
		}
	}

	return msg, nil
}

// WaitForReady waits for a READY message on CTL channel.
func (h *Hyperstart) WaitForReady() error {
	h.ctlMutex.Lock()
	defer h.ctlMutex.Unlock()

	_, err := h.expectReadingCmd(hyper.INIT_READY)
	if err != nil {
		return err
	}

	return nil
}

// SendCtlMessage sends a message to the CTL channel.
func (h *Hyperstart) SendCtlMessage(cmd string, data []byte) (*hyper.DecodedMessage, error) {
	h.ctlMutex.Lock()
	defer h.ctlMutex.Unlock()

	code, err := codeFromCmd(cmd)
	if err != nil {
		return nil, err
	}

	// Write message
	msg := &hyper.DecodedMessage{
		Code:    code,
		Message: data,
	}
	err = h.writeCtlMessage(msg)
	if err != nil {
		return nil, err
	}

	// Wait for answer
	resp, err := h.expectReadingCmd(hyper.INIT_ACK)
	if err != nil {
		return nil, err
	}

	return resp, nil
}
<|endoftext|>


<commit_before>package migrations

import (
	"fmt"
	"runtime"
	"time"

	"github.com/cozy/cozy-stack/model/instance"
	"github.com/cozy/cozy-stack/model/job"
	"github.com/cozy/cozy-stack/model/vfs"
	"github.com/cozy/cozy-stack/model/vfs/vfsswift"
	"github.com/cozy/cozy-stack/pkg/config/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/lock"
	"github.com/cozy/cozy-stack/pkg/logger"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/ncw/swift"
)

const (
	swiftV1ToV2 = "swift-v1-to-v2"
	toSwiftV3   = "to-swift-v3"

	swiftV1ContainerPrefixCozy = "cozy-"
	swiftV1ContainerPrefixData = "data-"
	swiftV2ContainerPrefixCozy = "cozy-v2-"
	swiftV2ContainerPrefixData = "data-v2-"
	swiftV3ContainerPrefix     = "cozy-v3-"
)

// maxSimultaneousCalls is the maximal number of simultaneous calls to Swift
// made by a single migration.
const maxSimultaneousCalls = 8

func init() {
	job.AddWorker(&job.WorkerConfig{
		WorkerType:   "migrations",
		Concurrency:  runtime.NumCPU(),
		MaxExecCount: 1,
		Reserved:     true,
		WorkerFunc:   worker,
		WorkerCommit: commit,
		Timeout:      1 * time.Hour,
	})
}

type message struct {
	Type string `json:"type"`
}

func worker(ctx *job.WorkerContext) error {
	var msg message
	if err := ctx.UnmarshalMessage(&msg); err != nil {
		return err
	}

	logger.WithDomain(ctx.Instance.Domain).WithField("nspace", "migration").
		Infof("Start the migration %s", msg.Type)

	switch msg.Type {
	case toSwiftV3:
		return migrateToSwiftV3(ctx.Instance.Domain)
	case swiftV1ToV2:
		return fmt.Errorf("this migration type is no longer supported")
	default:
		return fmt.Errorf("unknown migration type %q", msg.Type)
	}
}

func commit(ctx *job.WorkerContext, err error) error {
	log := logger.WithDomain(ctx.Instance.Domain).WithField("nspace", "migration")
	if err == nil {
		log.Infof("Migration success")
	} else {
		log.Errorf("Migration error: %s", err)
	}
	return err
}

func migrateToSwiftV3(domain string) error {
	c := config.GetSwiftConnection()
	inst, err := instance.GetFromCouch(domain)
	if err != nil {
		return err
	}

	var srcContainer, migratedFrom string
	switch inst.SwiftLayout {
	case 0: // layout v1
		srcContainer = swiftV1ContainerPrefixCozy + domain
		migratedFrom = "v1"
	case 1: // layout v2
		srcContainer = swiftV2ContainerPrefixCozy + inst.DBPrefix()
		migratedFrom = "v2"
	case 2: // layout v3
		return nil // Nothing to do!
	default:
		return instance.ErrInvalidSwiftLayout
	}

	vfs := inst.VFS()
	root, err := vfs.DirByID(consts.RootDirID)
	if err != nil {
		return err
	}

	mutex := lock.ReadWrite(inst, "vfs")
	if err = mutex.Lock(); err != nil {
		return err
	}
	defer mutex.Unlock()

	dstContainer := swiftV3ContainerPrefix + inst.DBPrefix()
	if err = c.ContainerCreate(dstContainer, nil); err != nil {
		return err
	}
	defer func() {
		if err != nil {
			if err := vfsswift.DeleteContainer(c, dstContainer); err != nil {
				log := logger.WithDomain(inst.Domain).WithField("nspace", "migration")
				log.Errorf("Failed to delete v3 container %s: %s", dstContainer, err)
			}
		}
	}()

	if err = copyTheFilesToSwiftV3(inst, c, root, srcContainer, dstContainer); err != nil {
		return err
	}

	meta := &swift.Metadata{"cozy-migrated-from": migratedFrom}
	_ = c.ContainerUpdate(srcContainer, meta.ContainerHeaders())
	if in, err := instance.GetFromCouch(domain); err == nil {
		inst = in
	}
	inst.SwiftLayout = 2
	if err = couchdb.UpdateDoc(couchdb.GlobalDB, inst); err != nil {
		return err
	}
	_ = vfs.Delete()
	return nil
}

func copyTheFilesToSwiftV3(inst *instance.Instance, c *swift.Connection, root *vfs.DirDoc, src, dst string) error {
	nb := 0
	ch := make(chan error)
	log := logger.WithDomain(inst.Domain).
		WithField("nspace", "migration")

	var thumbsContainer string
	switch inst.SwiftLayout {
	case 0: // layout v1
		thumbsContainer = swiftV1ContainerPrefixData + inst.Domain
	case 1: // layout v2
		thumbsContainer = swiftV2ContainerPrefixData + inst.DBPrefix()
	default:
		return instance.ErrInvalidSwiftLayout
	}

	// Use a system of tokens to limit the number of simultaneous calls to
	// Swift: only a goroutine that has a token can make a call.
	tokens := make(chan int, maxSimultaneousCalls)
	for k := 0; k < maxSimultaneousCalls; k++ {
		tokens <- k
	}

	fs := inst.VFS()
	errm := vfs.WalkAlreadyLocked(fs, root, func(_ string, d *vfs.DirDoc, f *vfs.FileDoc, err error) error {
		if err != nil {
			return err
		}
		if f == nil {
			return nil
		}

		srcName := getSrcName(inst, f)
		dstName := getDstName(inst, f)
		if srcName == "" || dstName == "" {
			return fmt.Errorf("Unexpected copy: %q -> %q", srcName, dstName)
		}

		nb++
		go func() {
			k := <-tokens
			_, err := c.ObjectCopy(src, srcName, dst, dstName, nil)
			if err != nil {
				log.Warningf("Cannot copy file from %s %s to %s %s: %s",
					src, srcName, dst, dstName, err)
			}
			ch <- err
			tokens <- k
		}()

		// Copy the thumbnails
		if f.Class == "image" {
			srcSmall, srcMedium, srcLarge := getThumbsSrcNames(inst, f)
			dstSmall, dstMedium, dstLarge := getThumbsDstNames(inst, f)
			nb += 3
			go func() {
				k := <-tokens
				_, err := c.ObjectCopy(thumbsContainer, srcSmall, dst, dstSmall, nil)
				if err != nil {
					log.Infof("Cannot copy thumbnail small from %s %s to %s %s: %s",
						thumbsContainer, srcSmall, dst, dstSmall, err)
				}
				ch <- nil
				_, err = c.ObjectCopy(thumbsContainer, srcMedium, dst, dstMedium, nil)
				if err != nil {
					log.Infof("Cannot copy thumbnail medium from %s %s to %s %s: %s",
						thumbsContainer, srcMedium, dst, dstMedium, err)
				}
				ch <- nil
				_, err = c.ObjectCopy(thumbsContainer, srcLarge, dst, dstLarge, nil)
				if err != nil {
					log.Infof("Cannot copy thumbnail large from %s %s to %s %s: %s",
						thumbsContainer, srcLarge, dst, dstLarge, err)
				}
				ch <- nil
				tokens <- k
			}()
		}
		return nil
	})

	for i := 0; i < nb; i++ {
		if err := <-ch; err != nil {
			errm = multierror.Append(errm, err)
		}
	}
	// Get back the tokens to ensure that each goroutine can finish.
	for k := 0; k < maxSimultaneousCalls; k++ {
		<-tokens
	}
	return errm
}

func getSrcName(inst *instance.Instance, f *vfs.FileDoc) string {
	srcName := ""
	switch inst.SwiftLayout {
	case 0: // layout v1
		srcName = f.DirID + "/" + f.DocName
	case 1: // layout v2
		srcName = vfsswift.MakeObjectName(f.DocID)
	}
	return srcName
}

func getDstName(inst *instance.Instance, f *vfs.FileDoc) string {
	if f.InternalID == "" {
		f.InternalID = vfsswift.NewInternalID()
		if err := couchdb.UpdateDoc(inst, f); err != nil {
			return ""
		}
	}
	return vfsswift.MakeObjectNameV3(f.DocID, f.InternalID)
}

func getThumbsSrcNames(inst *instance.Instance, f *vfs.FileDoc) (string, string, string) {
	var small, medium, large string
	switch inst.SwiftLayout {
	case 0: // layout v1
		small = fmt.Sprintf("thumbs/%s-small", f.DocID)
		medium = fmt.Sprintf("thumbs/%s-medium", f.DocID)
		large = fmt.Sprintf("thumbs/%s-large", f.DocID)
	case 1: // layout v2
		obj := vfsswift.MakeObjectName(f.DocID)
		small = fmt.Sprintf("thumbs/%s-small", obj)
		medium = fmt.Sprintf("thumbs/%s-medium", obj)
		large = fmt.Sprintf("thumbs/%s-large", obj)
	}
	return small, medium, large
}

func getThumbsDstNames(inst *instance.Instance, f *vfs.FileDoc) (string, string, string) {
	obj := vfsswift.MakeObjectName(f.DocID)
	small := fmt.Sprintf("thumbs/%s-small", obj)
	medium := fmt.Sprintf("thumbs/%s-medium", obj)
	large := fmt.Sprintf("thumbs/%s-large", obj)
	return small, medium, large
}
<commit_msg>Increase timeout for migration worker<commit_after>package migrations

import (
	"fmt"
	"runtime"
	"time"

	"github.com/cozy/cozy-stack/model/instance"
	"github.com/cozy/cozy-stack/model/job"
	"github.com/cozy/cozy-stack/model/vfs"
	"github.com/cozy/cozy-stack/model/vfs/vfsswift"
	"github.com/cozy/cozy-stack/pkg/config/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/lock"
	"github.com/cozy/cozy-stack/pkg/logger"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/ncw/swift"
)

const (
	swiftV1ToV2 = "swift-v1-to-v2"
	toSwiftV3   = "to-swift-v3"

	swiftV1ContainerPrefixCozy = "cozy-"
	swiftV1ContainerPrefixData = "data-"
	swiftV2ContainerPrefixCozy = "cozy-v2-"
	swiftV2ContainerPrefixData = "data-v2-"
	swiftV3ContainerPrefix     = "cozy-v3-"
)

// maxSimultaneousCalls is the maximal number of simultaneous calls to Swift
// made by a single migration.
const maxSimultaneousCalls = 8

func init() {
	job.AddWorker(&job.WorkerConfig{
		WorkerType:   "migrations",
		Concurrency:  runtime.NumCPU(),
		MaxExecCount: 1,
		Reserved:     true,
		WorkerFunc:   worker,
		WorkerCommit: commit,
		Timeout:      6 * time.Hour,
	})
}

type message struct {
	Type string `json:"type"`
}

func worker(ctx *job.WorkerContext) error {
	var msg message
	if err := ctx.UnmarshalMessage(&msg); err != nil {
		return err
	}

	logger.WithDomain(ctx.Instance.Domain).WithField("nspace", "migration").
		Infof("Start the migration %s", msg.Type)

	switch msg.Type {
	case toSwiftV3:
		return migrateToSwiftV3(ctx.Instance.Domain)
	case swiftV1ToV2:
		return fmt.Errorf("this migration type is no longer supported")
	default:
		return fmt.Errorf("unknown migration type %q", msg.Type)
	}
}

func commit(ctx *job.WorkerContext, err error) error {
	log := logger.WithDomain(ctx.Instance.Domain).WithField("nspace", "migration")
	if err == nil {
		log.Infof("Migration success")
	} else {
		log.Errorf("Migration error: %s", err)
	}
	return err
}

func migrateToSwiftV3(domain string) error {
	c := config.GetSwiftConnection()
	inst, err := instance.GetFromCouch(domain)
	if err != nil {
		return err
	}

	var srcContainer, migratedFrom string
	switch inst.SwiftLayout {
	case 0: // layout v1
		srcContainer = swiftV1ContainerPrefixCozy + domain
		migratedFrom = "v1"
	case 1: // layout v2
		srcContainer = swiftV2ContainerPrefixCozy + inst.DBPrefix()
		migratedFrom = "v2"
	case 2: // layout v3
		return nil // Nothing to do!
	default:
		return instance.ErrInvalidSwiftLayout
	}

	vfs := inst.VFS()
	root, err := vfs.DirByID(consts.RootDirID)
	if err != nil {
		return err
	}

	mutex := lock.ReadWrite(inst, "vfs")
	if err = mutex.Lock(); err != nil {
		return err
	}
	defer mutex.Unlock()

	dstContainer := swiftV3ContainerPrefix + inst.DBPrefix()
	if err = c.ContainerCreate(dstContainer, nil); err != nil {
		return err
	}
	defer func() {
		if err != nil {
			if err := vfsswift.DeleteContainer(c, dstContainer); err != nil {
				log := logger.WithDomain(inst.Domain).WithField("nspace", "migration")
				log.Errorf("Failed to delete v3 container %s: %s", dstContainer, err)
			}
		}
	}()

	if err = copyTheFilesToSwiftV3(inst, c, root, srcContainer, dstContainer); err != nil {
		return err
	}

	meta := &swift.Metadata{"cozy-migrated-from": migratedFrom}
	_ = c.ContainerUpdate(srcContainer, meta.ContainerHeaders())
	if in, err := instance.GetFromCouch(domain); err == nil {
		inst = in
	}
	inst.SwiftLayout = 2
	if err = couchdb.UpdateDoc(couchdb.GlobalDB, inst); err != nil {
		return err
	}
	_ = vfs.Delete()
	return nil
}

func copyTheFilesToSwiftV3(inst *instance.Instance, c *swift.Connection, root *vfs.DirDoc, src, dst string) error {
	nb := 0
	ch := make(chan error)
	log := logger.WithDomain(inst.Domain).
		WithField("nspace", "migration")

	var thumbsContainer string
	switch inst.SwiftLayout {
	case 0: // layout v1
		thumbsContainer = swiftV1ContainerPrefixData + inst.Domain
	case 1: // layout v2
		thumbsContainer = swiftV2ContainerPrefixData + inst.DBPrefix()
	default:
		return instance.ErrInvalidSwiftLayout
	}

	// Use a system of tokens to limit the number of simultaneous calls to
	// Swift: only a goroutine that has a token can make a call.
	tokens := make(chan int, maxSimultaneousCalls)
	for k := 0; k < maxSimultaneousCalls; k++ {
		tokens <- k
	}

	fs := inst.VFS()
	errm := vfs.WalkAlreadyLocked(fs, root, func(_ string, d *vfs.DirDoc, f *vfs.FileDoc, err error) error {
		if err != nil {
			return err
		}
		if f == nil {
			return nil
		}

		srcName := getSrcName(inst, f)
		dstName := getDstName(inst, f)
		if srcName == "" || dstName == "" {
			return fmt.Errorf("Unexpected copy: %q -> %q", srcName, dstName)
		}

		nb++
		go func() {
			k := <-tokens
			_, err := c.ObjectCopy(src, srcName, dst, dstName, nil)
			if err != nil {
				log.Warningf("Cannot copy file from %s %s to %s %s: %s",
					src, srcName, dst, dstName, err)
			}
			ch <- err
			tokens <- k
		}()

		// Copy the thumbnails
		if f.Class == "image" {
			srcSmall, srcMedium, srcLarge := getThumbsSrcNames(inst, f)
			dstSmall, dstMedium, dstLarge := getThumbsDstNames(inst, f)
			nb += 3
			go func() {
				k := <-tokens
				_, err := c.ObjectCopy(thumbsContainer, srcSmall, dst, dstSmall, nil)
				if err != nil {
					log.Infof("Cannot copy thumbnail small from %s %s to %s %s: %s",
						thumbsContainer, srcSmall, dst, dstSmall, err)
				}
				ch <- nil
				_, err = c.ObjectCopy(thumbsContainer, srcMedium, dst, dstMedium, nil)
				if err != nil {
					log.Infof("Cannot copy thumbnail medium from %s %s to %s %s: %s",
						thumbsContainer, srcMedium, dst, dstMedium, err)
				}
				ch <- nil
				_, err = c.ObjectCopy(thumbsContainer, srcLarge, dst, dstLarge, nil)
				if err != nil {
					log.Infof("Cannot copy thumbnail large from %s %s to %s %s: %s",
						thumbsContainer, srcLarge, dst, dstLarge, err)
				}
				ch <- nil
				tokens <- k
			}()
		}
		return nil
	})

	for i := 0; i < nb; i++ {
		if err := <-ch; err != nil {
			errm = multierror.Append(errm, err)
		}
	}
	// Get back the tokens to ensure that each goroutine can finish.
	for k := 0; k < maxSimultaneousCalls; k++ {
		<-tokens
	}
	return errm
}

func getSrcName(inst *instance.Instance, f *vfs.FileDoc) string {
	srcName := ""
	switch inst.SwiftLayout {
	case 0: // layout v1
		srcName = f.DirID + "/" + f.DocName
	case 1: // layout v2
		srcName = vfsswift.MakeObjectName(f.DocID)
	}
	return srcName
}

func getDstName(inst *instance.Instance, f *vfs.FileDoc) string {
	if f.InternalID == "" {
		f.InternalID = vfsswift.NewInternalID()
		if err := couchdb.UpdateDoc(inst, f); err != nil {
			return ""
		}
	}
	return vfsswift.MakeObjectNameV3(f.DocID, f.InternalID)
}

func getThumbsSrcNames(inst *instance.Instance, f *vfs.FileDoc) (string, string, string) {
	var small, medium, large string
	switch inst.SwiftLayout {
	case 0: // layout v1
		small = fmt.Sprintf("thumbs/%s-small", f.DocID)
		medium = fmt.Sprintf("thumbs/%s-medium", f.DocID)
		large = fmt.Sprintf("thumbs/%s-large", f.DocID)
	case 1: // layout v2
		obj := vfsswift.MakeObjectName(f.DocID)
		small = fmt.Sprintf("thumbs/%s-small", obj)
		medium = fmt.Sprintf("thumbs/%s-medium", obj)
		large = fmt.Sprintf("thumbs/%s-large", obj)
	}
	return small, medium, large
}

func getThumbsDstNames(inst *instance.Instance, f *vfs.FileDoc) (string, string, string) {
	obj := vfsswift.MakeObjectName(f.DocID)
	small := fmt.Sprintf("thumbs/%s-small", obj)
	medium := fmt.Sprintf("thumbs/%s-medium", obj)
	large := fmt.Sprintf("thumbs/%s-large", obj)
	return small, medium, large
}
<|endoftext|>


<commit_before>package matrix

import (
	"fmt"
	"math/rand"
)

type Dense struct {
	v          []float64 // [row, row, ..., row]
	rows, cols int
	stride     int // distance between vertically adjacent values
}

// Creates new matrix that refers to v
func NewDense(rows, cols, stride int, v []float64) *Dense {
	return &Dense{v: v, rows: rows, cols: cols, stride: stride}
}

// Retuns new zero matrix
func DenseZero(rows, cols int) *Dense {
	return &Dense{
		v:      make([]float64, cols*rows),
		rows:   rows,
		cols:   cols,
		stride: cols,
	}
}

// Returns identity matrix
func DenseI(n int) *Dense {
	m := DenseZero(n, n)
	inc := n + 1
	for i := 0; i < len(m.v); i += inc {
		m.v[i] = 1
	}
	return m
}

// Returns dimensions of the matrix (rows, cols)
func (m *Dense) Size() (int, int) {
	return m.rows, m.cols
}

func (m *Dense) Rows() int {
	return m.rows
}

func (m *Dense) Cols() int {
	return m.cols
}

func (m *Dense) Stride() int {
	return m.stride
}

func (m *Dense) Values() []float64 {
	return m.v
}

func (m *Dense) Get(i, k int) float64 {
	m.checkIndexes(i, k)
	return m.v[i*m.stride+k]
}

func (m *Dense) Set(i, k int, a float64) {
	m.checkIndexes(i, k)
	m.v[i*m.stride+k] = a
}

// Sets all elements to a
func (m *Dense) SetAll(a float64) {
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		k := m.cols
		for k >= 2 {
			k--
			row[k] = a
			k--
			row[k] = a
		}
		if k != 0 {
			row[0] = a
		}
	}
}

// Sets all elements to: min + rand.Float64() * (max - min)
func (m *Dense) Rand(min, max float64) {
	s := max - min
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		for k := 0; k < m.cols; k++ {
			row[k] = min + s*rand.Float64()
		}
	}
}

// Sets all elements to: mean + stdDev*rand.NormFloat64()
func (m *Dense) RandNorm(mean, stdDev float64) {
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		for k := 0; k < m.cols; k++ {
			row[k] = mean + stdDev*rand.NormFloat64()
		}
	}
}

// Returns a slice of a matrix that contains rows from start to stop - 1
func (m *Dense) Hslice(start, stop int) *Dense {
	if start > stop || start < 0 || stop > m.rows {
		panic("bad indexes for horizontal slice")
	}
	return &Dense{
		v:      m.v[start*m.stride : stop*m.stride],
		rows:   stop - start,
		cols:   m.cols,
		stride: m.stride,
	}
}

// Returns a slice of a matrix that contains cols from start to stop - 1
func (m *Dense) Vslice(start, stop int) *Dense {
	if start > stop || start < 0 || stop > m.cols {
		panic("bad indexes for vertical slice")
	}
	return &Dense{
		v:      m.v[start : (m.rows-1)*m.stride+stop],
		rows:   m.rows,
		cols:   stop - start,
		stride: m.stride,
	}

}

func (m *Dense) Equal(a *Dense) bool {
	if m.rows != a.rows || m.cols != a.cols {
		return false
	}
	for i := 0; i < m.rows; i++ {
		mr := m.v[i*m.stride:]
		ar := a.v[i*a.stride:]
		for k := 0; k < m.cols; k++ {
			if mr[k] != ar[k] {
				return false
			}
		}
	}
	return true
}

func (m *Dense) String() string {
	rows, cols := m.Size()
	s := "["
	for i := 0; i < rows; i++ {
		if i != 0 {
			s += "\n "
		}
		o := ""
		for k := 0; k < cols; k++ {
			s += fmt.Sprintf("%s%-g", o, m.Get(i, k))
			if k == 0 {
				o = " "
			}
		}
	}
	return s + "]"
}

// Utils

func (m *Dense) checkIndexes(i, k int) {
	if i < 0 || i >= m.rows {
		panic(fmt.Sprintf("row index (%d) out of range [0, %d]", i, m.rows))
	}
	if k < 0 || k >= m.cols {
		panic(fmt.Sprintf("column index (%d) out of range [0, %d]", k, m.cols))
	}
}

func (m *Dense) checkEqualDims(a *Dense) {
	if m.Rows() != a.Rows() || m.Cols() != a.Cols() {
		panic("dimensions of matrices are not equal")
	}
}
<commit_msg>NewMatrix fixed<commit_after>package matrix

import (
	"fmt"
	"math/rand"
)

type Dense struct {
	v          []float64 // [row, row, ..., row]
	rows, cols int
	stride     int // distance between vertically adjacent values
}

// Creates new matrix that refers to v
func NewDense(rows, cols, stride int, v []float64) *Dense {
	if rows * stride > len(v) {
		panic("v is to small")
	}
	return &Dense{v: v, rows: rows, cols: cols, stride: stride}
}

// Retuns new zero matrix
func DenseZero(rows, cols int) *Dense {
	return &Dense{
		v:      make([]float64, cols*rows),
		rows:   rows,
		cols:   cols,
		stride: cols,
	}
}

// Returns identity matrix
func DenseI(n int) *Dense {
	m := DenseZero(n, n)
	inc := n + 1
	for i := 0; i < len(m.v); i += inc {
		m.v[i] = 1
	}
	return m
}

// Returns dimensions of the matrix (rows, cols)
func (m *Dense) Size() (int, int) {
	return m.rows, m.cols
}

func (m *Dense) Rows() int {
	return m.rows
}

func (m *Dense) Cols() int {
	return m.cols
}

func (m *Dense) Stride() int {
	return m.stride
}

func (m *Dense) Values() []float64 {
	return m.v
}

func (m *Dense) Get(i, k int) float64 {
	m.checkIndexes(i, k)
	return m.v[i*m.stride+k]
}

func (m *Dense) Set(i, k int, a float64) {
	m.checkIndexes(i, k)
	m.v[i*m.stride+k] = a
}

// Sets all elements to a
func (m *Dense) SetAll(a float64) {
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		k := m.cols
		for k >= 2 {
			k--
			row[k] = a
			k--
			row[k] = a
		}
		if k != 0 {
			row[0] = a
		}
	}
}

// Sets all elements to: min + rand.Float64() * (max - min)
func (m *Dense) Rand(min, max float64) {
	s := max - min
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		for k := 0; k < m.cols; k++ {
			row[k] = min + s*rand.Float64()
		}
	}
}

// Sets all elements to: mean + stdDev*rand.NormFloat64()
func (m *Dense) RandNorm(mean, stdDev float64) {
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		for k := 0; k < m.cols; k++ {
			row[k] = mean + stdDev*rand.NormFloat64()
		}
	}
}

// Returns a slice of a matrix that contains rows from start to stop - 1
func (m *Dense) Hslice(start, stop int) *Dense {
	if start > stop || start < 0 || stop > m.rows {
		panic("bad indexes for horizontal slice")
	}
	return &Dense{
		v:      m.v[start*m.stride : stop*m.stride],
		rows:   stop - start,
		cols:   m.cols,
		stride: m.stride,
	}
}

// Returns a slice of a matrix that contains cols from start to stop - 1
func (m *Dense) Vslice(start, stop int) *Dense {
	if start > stop || start < 0 || stop > m.cols {
		panic("bad indexes for vertical slice")
	}
	return &Dense{
		v:      m.v[start : (m.rows-1)*m.stride+stop],
		rows:   m.rows,
		cols:   stop - start,
		stride: m.stride,
	}

}

func (m *Dense) Equal(a *Dense) bool {
	if m.rows != a.rows || m.cols != a.cols {
		return false
	}
	for i := 0; i < m.rows; i++ {
		mr := m.v[i*m.stride:]
		ar := a.v[i*a.stride:]
		for k := 0; k < m.cols; k++ {
			if mr[k] != ar[k] {
				return false
			}
		}
	}
	return true
}

func (m *Dense) String() string {
	rows, cols := m.Size()
	s := "["
	for i := 0; i < rows; i++ {
		if i != 0 {
			s += "\n "
		}
		o := ""
		for k := 0; k < cols; k++ {
			s += fmt.Sprintf("%s%-g", o, m.Get(i, k))
			if k == 0 {
				o = " "
			}
		}
	}
	return s + "]"
}

// Utils

func (m *Dense) checkIndexes(i, k int) {
	if i < 0 || i >= m.rows {
		panic(fmt.Sprintf("row index (%d) out of range [0, %d]", i, m.rows))
	}
	if k < 0 || k >= m.cols {
		panic(fmt.Sprintf("column index (%d) out of range [0, %d]", k, m.cols))
	}
}

func (m *Dense) checkEqualDims(a *Dense) {
	if m.Rows() != a.Rows() || m.Cols() != a.Cols() {
		panic("dimensions of matrices are not equal")
	}
}
<|endoftext|>