<commit_before>/*
 * Copyright (c) 2017 AlexRuzin (stan.ruzin@gmail.com)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package govfs

import (
    "testing"
    "time"
    "os"
    "io"
    _"bytes"
    "bytes"
)

const FS_DATABASE_FILE string = "test_db"

func TestIOSanity(t *testing.T) {
    /*
     * This test will generate the raw fs stream file, along with some contents
     *  that will be later loaded by the TestFSReader() method
     */
    out("[+] Running Standard I/O Sanity Test...")

    /* Remove the test database if it exists */
    var filename = gen_raw_filename(FS_DATABASE_FILE)
    if _, err := os.Stat(filename); os.IsExist(err) {
        os.Remove(filename)
    }

    var header = CreateDatabase(filename, 0)
    if header == nil {
        drive_fail("TEST1: Failed to obtain header", t)
    }
    out("[+] Test 1 PASS")

    // The root file "/" must at least exist
    if file, err := header.Create("/"); file != nil && err == nil {
        drive_fail("TEST2: Failed to return root handle", t)
    }
    out("[+] Test 2 PASS")

    /*
     * Try to delete the root file "/"
     */
    if header.Delete("/") == nil {
        drive_fail("TEST3: Cannot delete root -- critical", t)
    }
    out("[+] Test 3 PASS")

    /*
     * Attempt to write to a nonexistant file
     */
    var data = []byte{ 1, 2 }
    if header.Write("/folder5/folder5/file5", data) == nil {
        drive_fail("TEST4: Cannot write to a nonexistant file", t)
    }
    out("[+] Test 4 PASS")

    /*
     * Create empty file file9
     */
    if file, err := header.Create("/folder5/folder4/folder2/file9"); file == nil || err != nil {
        drive_fail("TEST4.1: file9 cannot be created", t)
    }
    out("[+] Test 4.1 PASS")

    /*
     * Attempt to create a new file0
     */
    if file, err := header.Create("/folder0/folder0/file0"); file == nil || err != nil {
        drive_fail("TEST5.0: file0 cannot be created", t)
    }
    out("[+] Test 5.0 PASS")

    /*
     * Attempt to create a new file0, this will fail since it should already exist
     */
    if file, err := header.Create("/folder0/folder0/file0"); file != nil && err == nil {
        drive_fail("TEST5.1: file0 cannot be created twice", t)
    }
    out("[+] Test 5.1 PASS")

    /*
     * Write some data into file0
     */
    data = []byte{ 1, 2, 3, 4 }
    if header.Write("/folder0/folder0/file0", data) != nil {
        drive_fail("TEST6: Failed to write data in file0", t)
    }
    out("[+] Test 6 PASS")

    /*
     * Check that the size of file0 is 4
     */
    if k, _ := header.get_file_size("/folder0/folder0/file0"); k != uint(len(data)) {
        drive_fail("TEST6.1: The size of data does not match", t)
    }
    out("[+] Test 6.1 PASS")

    /*
     * Attempt to create a new file3
     */
    if file, err := header.Create("/folder1/folder0/file3"); file == nil || err != nil {
        drive_fail("TEST7: file3 cannot be created", t)
    }
    out("[+] Test 7 PASS")

    /*
     * Write some data into file3
     */
    var data2 = []byte{ 1, 2, 3, 4, 5, 6, 7 }
    if header.Write("/folder1/folder0/file3", data2) != nil {
        drive_fail("TEST8: Failed to write data in file3", t)
    }
    out("[+] Test 8 PASS")

    /*
     * Write some data into file3
     */
    if header.Write("/folder1/folder0/file3", data2) != nil {
        drive_fail("TEST8.1: Failed to write data in file3", t)
    }
    out("[+] Test 8.1 PASS")

    /*
     * Read the written data from file0 and compare
     */
    output_data, _ := header.Read("/folder0/folder0/file0")
    if output_data == nil || len(output_data) != len(data) || header.t_size - 7 /* len(file3) */ != uint(len(data)) {
        drive_fail("TEST9: Failed to read data from file0", t)
    }
    out("[+] Test 9 PASS")

    /*
     * Read the written data from file3 and compare
     */
    output_data, _ = header.Read("/folder1/folder0/file3")
    if output_data == nil || len(output_data) != len(data2) || header.t_size - 4 /* len(file0) */ != uint(len(data2)) {
        drive_fail("TEST10: Failed to read data from file3", t)
    }
    out("[+] Test 10 PASS")

    /*
     * Write other data to file0
     */
    data = []byte{ 1, 2, 3 }
    if header.Write("/folder0/folder0/file0", data) != nil {
        drive_fail("TEST11: Failed to write data in file1", t)
    }
    out("[+] Test 11 PASS")

    /*
     * Read the new data from file0
     */
    output_data, _ = header.Read("/folder0/folder0/file0")
    if output_data == nil || len(output_data) != len(data) {
        drive_fail("TEST12: Failed to read data from file1", t)
    }
    out("[+] Test 12 PASS")

    /*
     * Attempt to create a new file5. This will be a blank file
     */
    if file, err := header.Create("/folder2/file7"); file == nil || err != nil {
        drive_fail("TEST13: file3 cannot be created", t)
    }
    out("[+] Test 13 PASS")

    /*
     * Delete file0 -- complete this
     */
    // FIXME/ADDME

    /*
     * Create just a folder
     */
    if file, err := header.Create("/folder2/file5/"); file == nil || err != nil {
        drive_fail("TEST15: folder file5 cannot be created", t)
    }
    out("[+] Test 15 PASS")

    /*
     * Tests the Reader interface
     */
    reader, err := header.NewReader("/folder0/folder0/file0")
    if err != nil {
        drive_fail("TEST15.1: Failed to create Reader", t)
    }

    file0data := make([]byte, 3)
    data_read, err := reader.Read(file0data)
    if data_read != len(data) || err != io.EOF || bytes.Compare(file0data, data) != 0 {
        drive_fail("TEST15.2: Failed to read from NewReader", t)
    }

    /* Test the reader interface again */
    file0data = make([]byte, 1)
    data_read, err = reader.Read(file0data)
    if data_read != 1 || err != nil || file0data[0] != 1 {
        drive_fail("TEST15.3: Invalid Reader interface behaviour", t)
    }
    out("[+] Test 15.1, 15.2, 15.3 PASS -- Reader interface")

    /*
     * Tests the Writer interface
     */
    file0data = []byte{1, 2, 3, 4, 5, 6, 7, 8}
    writer, err := header.NewWriter("/folder0/folder0/file0")
    if writer == nil || err != nil {
        drive_fail("TEST15.4: Invalid Writer object", t)
    }

    written, err := writer.Write(file0data)
    if written != len(file0data) || err != io.EOF {
        drive_fail("TEST15.4: Invalid Writer response", t)
    }

    file0data = make([]byte, 8)
    data_read, err = reader.Read(file0data)
    if data_read != 8 || err != io.EOF || file0data[0] != 1 || file0data[1] != 2 {
        drive_fail("TEST15.5: Invalid Reader data",t )
    }
    out("[+] Test 15.4, 15.5 PASS -- Writer interface")

    /*
     * Print out files
     */
    file_list := header.get_file_list()
    for _, e := range file_list {
        out(e)
    }

    /*
     * Unmount/commit database to file
     */
    if err := header.unmount_db(); err != nil {
        drive_fail("TEST16: Failed to commit database", t)
    }
    out("[+] Test 16 PASS. Raw FS stream written to: " + header.filename)

    time.Sleep(10000)
}

func TestFSReader(t *testing.T) {
    /*
     * Read in FS_DATABASE_FILE and do basic tests
     */
    var filename = gen_raw_filename(FS_DATABASE_FILE)
    out("[+] Loading Raw FS stream file: " + filename)

    /* Remove the test database if it exists */
    if _, err := os.Stat(filename); os.IsNotExist(err) {
        drive_fail("error: Standard raw fs stream " + filename + " does not exist", t)
    }

    var header = CreateDatabase(filename, 0)
    if header == nil {
        drive_fail("TEST1: Failed to obtain header", t)
    }
    out("[+] Test 1 PASS (Loaded FS stream)")
}

func gen_raw_filename(suffix string) string {
    return os.Getenv("TEMP") + "\\" + suffix + ".db"
}

func drive_fail(output string, t *testing.T) {
    t.Errorf(output)
    t.FailNow()
}
<commit_msg>Changed TestIOSanity to TestFSWriter to conform with naming convention<commit_after>/*
 * Copyright (c) 2017 AlexRuzin (stan.ruzin@gmail.com)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package govfs

import (
    "testing"
    "time"
    "os"
    "io"
    _"bytes"
    "bytes"
)

const FS_DATABASE_FILE string = "test_db"

func TestFSWriter(t *testing.T) {
    /*
     * This test will generate the raw fs stream file, along with some contents
     *  that will be later loaded by the TestFSReader() method
     */
    out("[+] Running Standard I/O Sanity Test...")

    /* Remove the test database if it exists */
    var filename = gen_raw_filename(FS_DATABASE_FILE)
    if _, err := os.Stat(filename); os.IsExist(err) {
        os.Remove(filename)
    }

    var header = CreateDatabase(filename, 0)
    if header == nil {
        drive_fail("TEST1: Failed to obtain header", t)
    }
    out("[+] Test 1 PASS")

    // The root file "/" must at least exist
    if file, err := header.Create("/"); file != nil && err == nil {
        drive_fail("TEST2: Failed to return root handle", t)
    }
    out("[+] Test 2 PASS")

    /*
     * Try to delete the root file "/"
     */
    if header.Delete("/") == nil {
        drive_fail("TEST3: Cannot delete root -- critical", t)
    }
    out("[+] Test 3 PASS")

    /*
     * Attempt to write to a nonexistant file
     */
    var data = []byte{ 1, 2 }
    if header.Write("/folder5/folder5/file5", data) == nil {
        drive_fail("TEST4: Cannot write to a nonexistant file", t)
    }
    out("[+] Test 4 PASS")

    /*
     * Create empty file file9
     */
    if file, err := header.Create("/folder5/folder4/folder2/file9"); file == nil || err != nil {
        drive_fail("TEST4.1: file9 cannot be created", t)
    }
    out("[+] Test 4.1 PASS")

    /*
     * Attempt to create a new file0
     */
    if file, err := header.Create("/folder0/folder0/file0"); file == nil || err != nil {
        drive_fail("TEST5.0: file0 cannot be created", t)
    }
    out("[+] Test 5.0 PASS")

    /*
     * Attempt to create a new file0, this will fail since it should already exist
     */
    if file, err := header.Create("/folder0/folder0/file0"); file != nil && err == nil {
        drive_fail("TEST5.1: file0 cannot be created twice", t)
    }
    out("[+] Test 5.1 PASS")

    /*
     * Write some data into file0
     */
    data = []byte{ 1, 2, 3, 4 }
    if header.Write("/folder0/folder0/file0", data) != nil {
        drive_fail("TEST6: Failed to write data in file0", t)
    }
    out("[+] Test 6 PASS")

    /*
     * Check that the size of file0 is 4
     */
    if k, _ := header.get_file_size("/folder0/folder0/file0"); k != uint(len(data)) {
        drive_fail("TEST6.1: The size of data does not match", t)
    }
    out("[+] Test 6.1 PASS")

    /*
     * Attempt to create a new file3
     */
    if file, err := header.Create("/folder1/folder0/file3"); file == nil || err != nil {
        drive_fail("TEST7: file3 cannot be created", t)
    }
    out("[+] Test 7 PASS")

    /*
     * Write some data into file3
     */
    var data2 = []byte{ 1, 2, 3, 4, 5, 6, 7 }
    if header.Write("/folder1/folder0/file3", data2) != nil {
        drive_fail("TEST8: Failed to write data in file3", t)
    }
    out("[+] Test 8 PASS")

    /*
     * Write some data into file3
     */
    if header.Write("/folder1/folder0/file3", data2) != nil {
        drive_fail("TEST8.1: Failed to write data in file3", t)
    }
    out("[+] Test 8.1 PASS")

    /*
     * Read the written data from file0 and compare
     */
    output_data, _ := header.Read("/folder0/folder0/file0")
    if output_data == nil || len(output_data) != len(data) || header.t_size - 7 /* len(file3) */ != uint(len(data)) {
        drive_fail("TEST9: Failed to read data from file0", t)
    }
    out("[+] Test 9 PASS")

    /*
     * Read the written data from file3 and compare
     */
    output_data, _ = header.Read("/folder1/folder0/file3")
    if output_data == nil || len(output_data) != len(data2) || header.t_size - 4 /* len(file0) */ != uint(len(data2)) {
        drive_fail("TEST10: Failed to read data from file3", t)
    }
    out("[+] Test 10 PASS")

    /*
     * Write other data to file0
     */
    data = []byte{ 1, 2, 3 }
    if header.Write("/folder0/folder0/file0", data) != nil {
        drive_fail("TEST11: Failed to write data in file1", t)
    }
    out("[+] Test 11 PASS")

    /*
     * Read the new data from file0
     */
    output_data, _ = header.Read("/folder0/folder0/file0")
    if output_data == nil || len(output_data) != len(data) {
        drive_fail("TEST12: Failed to read data from file1", t)
    }
    out("[+] Test 12 PASS")

    /*
     * Attempt to create a new file5. This will be a blank file
     */
    if file, err := header.Create("/folder2/file7"); file == nil || err != nil {
        drive_fail("TEST13: file3 cannot be created", t)
    }
    out("[+] Test 13 PASS")

    /*
     * Delete file0 -- complete this
     */
    // FIXME/ADDME

    /*
     * Create just a folder
     */
    if file, err := header.Create("/folder2/file5/"); file == nil || err != nil {
        drive_fail("TEST15: folder file5 cannot be created", t)
    }
    out("[+] Test 15 PASS")

    /*
     * Tests the Reader interface
     */
    reader, err := header.NewReader("/folder0/folder0/file0")
    if err != nil {
        drive_fail("TEST15.1: Failed to create Reader", t)
    }

    file0data := make([]byte, 3)
    data_read, err := reader.Read(file0data)
    if data_read != len(data) || err != io.EOF || bytes.Compare(file0data, data) != 0 {
        drive_fail("TEST15.2: Failed to read from NewReader", t)
    }

    /* Test the reader interface again */
    file0data = make([]byte, 1)
    data_read, err = reader.Read(file0data)
    if data_read != 1 || err != nil || file0data[0] != 1 {
        drive_fail("TEST15.3: Invalid Reader interface behaviour", t)
    }
    out("[+] Test 15.1, 15.2, 15.3 PASS -- Reader interface")

    /*
     * Tests the Writer interface
     */
    file0data = []byte{1, 2, 3, 4, 5, 6, 7, 8}
    writer, err := header.NewWriter("/folder0/folder0/file0")
    if writer == nil || err != nil {
        drive_fail("TEST15.4: Invalid Writer object", t)
    }

    written, err := writer.Write(file0data)
    if written != len(file0data) || err != io.EOF {
        drive_fail("TEST15.4: Invalid Writer response", t)
    }

    file0data = make([]byte, 8)
    data_read, err = reader.Read(file0data)
    if data_read != 8 || err != io.EOF || file0data[0] != 1 || file0data[1] != 2 {
        drive_fail("TEST15.5: Invalid Reader data",t )
    }
    out("[+] Test 15.4, 15.5 PASS -- Writer interface")

    /*
     * Print out files
     */
    file_list := header.get_file_list()
    for _, e := range file_list {
        out(e)
    }

    /*
     * Unmount/commit database to file
     */
    if err := header.unmount_db(); err != nil {
        drive_fail("TEST16: Failed to commit database", t)
    }
    out("[+] Test 16 PASS. Raw FS stream written to: " + header.filename)

    time.Sleep(10000)
}

func TestFSReader(t *testing.T) {
    /*
     * Read in FS_DATABASE_FILE and do basic tests
     */
    var filename = gen_raw_filename(FS_DATABASE_FILE)
    out("[+] Loading Raw FS stream file: " + filename)

    /* Remove the test database if it exists */
    if _, err := os.Stat(filename); os.IsNotExist(err) {
        drive_fail("error: Standard raw fs stream " + filename + " does not exist", t)
    }

    var header = CreateDatabase(filename, 0)
    if header == nil {
        drive_fail("TEST1: Failed to obtain header", t)
    }
    out("[+] Test 1 PASS (Loaded FS stream)")
}

func gen_raw_filename(suffix string) string {
    return os.Getenv("TEMP") + "\\" + suffix + ".db"
}

func drive_fail(output string, t *testing.T) {
    t.Errorf(output)
    t.FailNow()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/twitter/gozer/mesos"
)

var (
	user       = flag.String("user", "", "The user to register as")
	port       = flag.Int("port", 4343, "Port to listen on for HTTP endpoint")
	master     = flag.String("master", "localhost", "Hostname of the master")
	masterPort = flag.Int("masterPort", 5050, "Port of the master")

	taskstore = &TaskStore{
		Tasks: make(map[string]*Task),
	}
)

type TaskState int

const (
	TaskState_UNKNOWN TaskState = iota
	TaskState_INIT
	TaskState_STARTING
	TaskState_RUNNING
	TaskState_FAILED
	TaskState_FINISHED
)

func (t TaskState) String() string {
	switch t {
	case TaskState_INIT:
		return "INIT"
	case TaskState_STARTING:
		return "STARTING"
	case TaskState_RUNNING:
		return "RUNNING"
	case TaskState_FAILED:
		return "FAILED"
	case TaskState_FINISHED:
		return "FINISHED"
	default:
		return "UNKNOWN"
	}
}

type Task struct {
	Id        string           `json:"id"`
	Command   string           `json:"command"`
	State     TaskState        `json:"-"`
	MesosTask *mesos.MesosTask `json:"-"`
	// TODO(dhamon): resource requirements
}

type TaskStore struct {
	sync.RWMutex
	Tasks map[string]*Task
}

func (t *TaskStore) AddTask(task *Task) error {
	t.Lock()
	defer t.Unlock()

	if _, ok := t.Tasks[task.Id]; ok {
		return fmt.Errorf("Task Id '%s' already exists, addition ignored", task.Id)
	}

	task.State = TaskState_INIT
	task.MesosTask = &mesos.MesosTask{
		Id:      task.Id,
		Command: task.Command,
	}
	t.Tasks[task.Id] = task
	log.Print("TASK '%s' State * -> %s", task.Id, task.State)

	return nil
}

func (t *TaskStore) UpdateTask(taskId string, state TaskState) error {
	t.Lock()
	defer t.Unlock()

	task, ok := t.Tasks[taskId]
	if !ok {
		return fmt.Errorf("Task Id '%s' not found, update ignored", taskId)
	}

	log.Printf("Task '%s' State %s -> %s", taskId, task.State, state)
	task.State = state

	return nil
}

func (t *TaskStore) TaskIds() []string {
	t.RLock()
	defer t.RUnlock()

	keys := make([]string, 0)
	for key := range t.Tasks {
		keys = append(keys, key)
	}

	return keys
}

func (t *TaskStore) TaskState(taskId string) (TaskState, error) {
	t.RLock()
	defer t.RUnlock()

	task, ok := t.Tasks[taskId]
	if !ok {
		return TaskState_UNKNOWN, fmt.Errorf("Task Id '%s' not found", taskId)
	}

	return task.State, nil
}

func (t *TaskStore) MesosTask(taskId string) (*mesos.MesosTask, error) {
	t.RLock()
	defer t.RUnlock()

	task, ok := t.Tasks[taskId]
	if !ok {
		return nil, fmt.Errorf("Task Id '%s' not found", taskId)
	}

	return task.MesosTask, nil
}

func startAPI() {
	log.Printf("api listening on port %d", *port)
	if err := http.ListenAndServe(fmt.Sprintf(":%d", *port), nil); err != nil {
		log.Fatalf("failed to start listening on port %d", *port)
	}
}

func addTaskHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Add("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		log.Printf("received addtask request with unexpected method. want %q, got %q: %+v", "POST", r.Method, r)
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Printf("ERROR: failed to read body from addtask request %+v: %+v", r, err)
		// TODO(dhamon): Better error for this case.
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	var task Task
	err = json.Unmarshal(body, &task)

	if err != nil {
		log.Printf("ERROR: failed to parse JSON body from addtask request %+v: %+v", r, err)
		// TODO(dhamon): Better error for this case.
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	taskstore.AddTask(&task)

	w.WriteHeader(http.StatusOK)
}

func main() {
	flag.Parse()

	http.HandleFunc("/api/addtask", addTaskHandler)
	go startAPI()

	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

	log.Printf("Registering...")
	master, err := mesos.New(&mesos.MesosMasterConfig{
		FrameworkName:  "gozer",
		RegisteredUser: "nobody",
		Masters: []mesos.MesosMasterLocation{mesos.MesosMasterLocation{
			Hostname: *master,
			Port:     *masterPort,
		}},
	})
	if err != nil {
		log.Fatal(err)
	}

	// Start Framework engine
	go master.Run()

	// Launch simple /bin/true after a little while
	go func() {
		time.Sleep(10 * time.Second)

		trueTask := &Task{
			Id:      "gozer-bin-true",
			Command: "/bin/true",
		}

		taskstore.AddTask(trueTask)
	}()

	// Shephard all our tasks
	//
	// Note: This will require a significant re-architecting, most likely to break out the
	// gozer based tasks and their state transitions, possibly using a go-routine per task,
	// which may limit the total number of tasks we can handle (100k go-routines might be
	// too much).  It should make for a simple abstraction, where the update routine should
	// then be able to simply use a channel to post a state transition to the gozer task,
	// which the gozer task manager (per task) go routine would then use to transition the
	// task through its state diagram.  This should also make it very simple to detect bad
	// transitions.
	//
	// It would be nice if we could just only use the mesos.TaskState_TASK_* states, however,
	// they do not encompass the ideas of PENDING (waiting for offers), and ASSIGNED (offer
	// selected, waiting for running), nor do they encompass tear-down and death.
	//
	// For now we use a simple loop to do a very naiive management of tasks, updates, events,
	// errors, etc.
	for {

		select {
		// case <-master.Events:
		// case <-master.Updates:

		case <-time.After(5 * time.Second):
			log.Print("Gozer: timeout, check tasks")
			// After a timeout, see if there any tasks to launch
			taskIds := taskstore.TaskIds()
			for _, taskId := range taskIds {
				state, err := taskstore.TaskState(taskId)
				if err != nil || state != TaskState_INIT {
					continue
				}

				mesosTask, err := taskstore.MesosTask(taskId)
				if err != nil {
					log.Print(err)
					continue
				}

				// Start this task (very naiive method)
				offer := <-master.Offers
				err = master.LaunchTask(offer, mesosTask)
				if err != nil {
					log.Print(err)
					continue
				}

				taskstore.UpdateTask(taskId, TaskState_STARTING)
			}
		}
	}
}
<commit_msg>Whoops, forgot to use user from flags. :)<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/twitter/gozer/mesos"
)

var (
	user       = flag.String("user", "", "The user to register as")
	port       = flag.Int("port", 4343, "Port to listen on for HTTP endpoint")
	master     = flag.String("master", "localhost", "Hostname of the master")
	masterPort = flag.Int("masterPort", 5050, "Port of the master")

	taskstore = &TaskStore{
		Tasks: make(map[string]*Task),
	}
)

type TaskState int

const (
	TaskState_UNKNOWN TaskState = iota
	TaskState_INIT
	TaskState_STARTING
	TaskState_RUNNING
	TaskState_FAILED
	TaskState_FINISHED
)

func (t TaskState) String() string {
	switch t {
	case TaskState_INIT:
		return "INIT"
	case TaskState_STARTING:
		return "STARTING"
	case TaskState_RUNNING:
		return "RUNNING"
	case TaskState_FAILED:
		return "FAILED"
	case TaskState_FINISHED:
		return "FINISHED"
	default:
		return "UNKNOWN"
	}
}

type Task struct {
	Id        string           `json:"id"`
	Command   string           `json:"command"`
	State     TaskState        `json:"-"`
	MesosTask *mesos.MesosTask `json:"-"`
	// TODO(dhamon): resource requirements
}

type TaskStore struct {
	sync.RWMutex
	Tasks map[string]*Task
}

func (t *TaskStore) AddTask(task *Task) error {
	t.Lock()
	defer t.Unlock()

	if _, ok := t.Tasks[task.Id]; ok {
		return fmt.Errorf("Task Id '%s' already exists, addition ignored", task.Id)
	}

	task.State = TaskState_INIT
	task.MesosTask = &mesos.MesosTask{
		Id:      task.Id,
		Command: task.Command,
	}
	t.Tasks[task.Id] = task
	log.Print("TASK '%s' State * -> %s", task.Id, task.State)

	return nil
}

func (t *TaskStore) UpdateTask(taskId string, state TaskState) error {
	t.Lock()
	defer t.Unlock()

	task, ok := t.Tasks[taskId]
	if !ok {
		return fmt.Errorf("Task Id '%s' not found, update ignored", taskId)
	}

	log.Printf("Task '%s' State %s -> %s", taskId, task.State, state)
	task.State = state

	return nil
}

func (t *TaskStore) TaskIds() []string {
	t.RLock()
	defer t.RUnlock()

	keys := make([]string, 0)
	for key := range t.Tasks {
		keys = append(keys, key)
	}

	return keys
}

func (t *TaskStore) TaskState(taskId string) (TaskState, error) {
	t.RLock()
	defer t.RUnlock()

	task, ok := t.Tasks[taskId]
	if !ok {
		return TaskState_UNKNOWN, fmt.Errorf("Task Id '%s' not found", taskId)
	}

	return task.State, nil
}

func (t *TaskStore) MesosTask(taskId string) (*mesos.MesosTask, error) {
	t.RLock()
	defer t.RUnlock()

	task, ok := t.Tasks[taskId]
	if !ok {
		return nil, fmt.Errorf("Task Id '%s' not found", taskId)
	}

	return task.MesosTask, nil
}

func startAPI() {
	log.Printf("api listening on port %d", *port)
	if err := http.ListenAndServe(fmt.Sprintf(":%d", *port), nil); err != nil {
		log.Fatalf("failed to start listening on port %d", *port)
	}
}

func addTaskHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Add("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		log.Printf("received addtask request with unexpected method. want %q, got %q: %+v", "POST", r.Method, r)
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Printf("ERROR: failed to read body from addtask request %+v: %+v", r, err)
		// TODO(dhamon): Better error for this case.
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	var task Task
	err = json.Unmarshal(body, &task)

	if err != nil {
		log.Printf("ERROR: failed to parse JSON body from addtask request %+v: %+v", r, err)
		// TODO(dhamon): Better error for this case.
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	taskstore.AddTask(&task)

	w.WriteHeader(http.StatusOK)
}

func main() {
	flag.Parse()

	http.HandleFunc("/api/addtask", addTaskHandler)
	go startAPI()

	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

	log.Printf("Registering...")
	master, err := mesos.New(&mesos.MesosMasterConfig{
		FrameworkName:  "gozer",
		RegisteredUser: *user,
		Masters: []mesos.MesosMasterLocation{mesos.MesosMasterLocation{
			Hostname: *master,
			Port:     *masterPort,
		}},
	})
	if err != nil {
		log.Fatal(err)
	}

	// Start Framework engine
	go master.Run()

	// Launch simple /bin/true after a little while
	go func() {
		time.Sleep(10 * time.Second)

		trueTask := &Task{
			Id:      "gozer-bin-true",
			Command: "/bin/true",
		}

		taskstore.AddTask(trueTask)
	}()

	// Shephard all our tasks
	//
	// Note: This will require a significant re-architecting, most likely to break out the
	// gozer based tasks and their state transitions, possibly using a go-routine per task,
	// which may limit the total number of tasks we can handle (100k go-routines might be
	// too much).  It should make for a simple abstraction, where the update routine should
	// then be able to simply use a channel to post a state transition to the gozer task,
	// which the gozer task manager (per task) go routine would then use to transition the
	// task through its state diagram.  This should also make it very simple to detect bad
	// transitions.
	//
	// It would be nice if we could just only use the mesos.TaskState_TASK_* states, however,
	// they do not encompass the ideas of PENDING (waiting for offers), and ASSIGNED (offer
	// selected, waiting for running), nor do they encompass tear-down and death.
	//
	// For now we use a simple loop to do a very naiive management of tasks, updates, events,
	// errors, etc.
	for {

		select {
		// case <-master.Events:
		// case <-master.Updates:

		case <-time.After(5 * time.Second):
			log.Print("Gozer: timeout, check tasks")
			// After a timeout, see if there any tasks to launch
			taskIds := taskstore.TaskIds()
			for _, taskId := range taskIds {
				state, err := taskstore.TaskState(taskId)
				if err != nil || state != TaskState_INIT {
					continue
				}

				mesosTask, err := taskstore.MesosTask(taskId)
				if err != nil {
					log.Print(err)
					continue
				}

				// Start this task (very naiive method)
				offer := <-master.Offers
				err = master.LaunchTask(offer, mesosTask)
				if err != nil {
					log.Print(err)
					continue
				}

				taskstore.UpdateTask(taskId, TaskState_STARTING)
			}
		}
	}
}
<|endoftext|>


<commit_before>package flow

import (
	"testing"
)

type echo struct {
	In  <-chan int
	Out chan<- int
}

func (c *echo) Process() {
	for i := range c.In {
		c.Out <- i
	}
	close(c.Out)
}

func newDoubleEcho() (*Graph, error) {
	n := NewGraph()
	// Components
	e1 := new(echo)
	e2 := new(echo)
	// Structure
	if err := n.Add("e1", e1); err != nil {
		return nil, err
	}
	if err := n.Add("e2", e2); err != nil {
		return nil, err
	}
	if err := n.Connect("e1", "Out", "e2", "In"); err != nil {
		return nil, err
	}
	// Ports
	if err := n.MapInPort("In", "e1", "In"); err != nil {
		return nil, err
	}
	if err := n.MapOutPort("Out", "e2", "Out"); err != nil {
		return nil, err
	}
	return n, nil
}

func TestSimpleGraph(t *testing.T) {
	data := []int{7, 97, 16, 356, 81}

	n, err := newDoubleEcho()
	if err != nil {
		t.Error(err)
		return
	}

	in := make(chan int)
	out := make(chan int)
	n.SetInPort("In", in)
	n.SetOutPort("Out", out)

	wait := Run(n)

	go func() {
		for _, n := range data {
			in <- n
		}
		close(in)
	}()

	i := 0
	for actual := range out {
		expected := data[i]
		if actual != expected {
			t.Errorf("%d != %d", actual, expected)
		}
		i++
	}

	<-wait
}
<commit_msg>Add test for adding invalid process type<commit_after>package flow

import (
	"testing"
)

type echo struct {
	In  <-chan int
	Out chan<- int
}

func (c *echo) Process() {
	for i := range c.In {
		c.Out <- i
	}
	close(c.Out)
}

func newDoubleEcho() (*Graph, error) {
	n := NewGraph()
	// Components
	e1 := new(echo)
	e2 := new(echo)
	// Structure
	if err := n.Add("e1", e1); err != nil {
		return nil, err
	}
	if err := n.Add("e2", e2); err != nil {
		return nil, err
	}
	if err := n.Connect("e1", "Out", "e2", "In"); err != nil {
		return nil, err
	}
	// Ports
	if err := n.MapInPort("In", "e1", "In"); err != nil {
		return nil, err
	}
	if err := n.MapOutPort("Out", "e2", "Out"); err != nil {
		return nil, err
	}
	return n, nil
}

func TestSimpleGraph(t *testing.T) {
	data := []int{7, 97, 16, 356, 81}

	n, err := newDoubleEcho()
	if err != nil {
		t.Error(err)
		return
	}

	in := make(chan int)
	out := make(chan int)
	n.SetInPort("In", in)
	n.SetOutPort("Out", out)

	wait := Run(n)

	go func() {
		for _, n := range data {
			in <- n
		}
		close(in)
	}()

	i := 0
	for actual := range out {
		expected := data[i]
		if actual != expected {
			t.Errorf("%d != %d", actual, expected)
		}
		i++
	}

	<-wait
}

func TestAddInvalidProcess(t *testing.T) {
	s := struct{ Name string }{"This is not a Component"}
	n := NewGraph()
	err := n.Add("wrong", s)
	if err == nil {
		t.Errorf("Expected an error")
	}
}
<|endoftext|>


<commit_before>package graval

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"strings"
)

const (
	welcomeMessage = "Welcome to the Go FTP Server"
)

type FTPConn struct {
	cwd           string
	conn          *net.TCPConn
	controlReader *bufio.Reader
	controlWriter *bufio.Writer
	data          *net.TCPConn
	driver        FTPDriver
	reqUser       string
	user          string
}

// NewFTPConn constructs a new object that will handle the FTP protocol over
// an active net.TCPConn. The TCP connection should already be open before
// it is handed to this functions. driver is an instance of FTPDrive that
// will handle all auth and persistence details.
func NewFTPConn(tcpConn *net.TCPConn, driver FTPDriver) *FTPConn {
	c := new(FTPConn)
	c.cwd = "/"
	c.conn = tcpConn
	c.controlReader = bufio.NewReader(tcpConn)
	c.controlWriter = bufio.NewWriter(tcpConn)
	c.driver = driver
	return c
}

// Serve starts an endless loop that reads FTP commands from the client and
// responds appropriately. terminated is a channel that will receive a true
// message when the connection closes. This loop will be running inside a
// goroutine, so use this channel to be notified when the connection can be
// cleaned up.
func (ftpConn *FTPConn) Serve(terminated chan bool) {
	log.Print("Connection Established")
	// send welcome
	ftpConn.writeMessage(220, welcomeMessage)
	// read commands
	for {
		line, err := ftpConn.controlReader.ReadString('\n')
		if err == nil {
			ftpConn.receiveLine(line)
		}
	}
	terminated <- true
	log.Print("Connection Terminated")
}

// Close will manually close this connection, even if the client isn't ready.
func (ftpConn *FTPConn) Close() {
	ftpConn.conn.Close()
	if ftpConn.data != nil {
		ftpConn.data.Close()
	}
}

// receiveLine accepts a single line FTP command and co-ordinates an
// appropriate response.
func (ftpConn *FTPConn) receiveLine(line string) {
	log.Print(line)
	command, param := ftpConn.parseLine(line)
	switch command {
	case "PASS":
		ftpConn.cmdPass(param)
		break
	case "QUIT":
		ftpConn.Close()
		break
	case "SYST":
		ftpConn.cmdSyst()
		break
	case "USER":
		ftpConn.cmdUser(param)
		break
	default:
		ftpConn.writeMessage(500, "Command not found")
	}
}

// cmdUser responds to the USER FTP command by asking for the password
func (ftpConn *FTPConn) cmdUser(param string) {
	ftpConn.reqUser = param
	ftpConn.writeMessage(331, "User name ok, password required")
}

// cmdPass respond to the PASS FTP command by asking the driver if the supplied
// username and password are valid
func (ftpConn *FTPConn) cmdPass(param string) {
	if ftpConn.driver.Authenticate(ftpConn.reqUser, param) {
		ftpConn.user = ftpConn.reqUser
		ftpConn.reqUser = ""
		ftpConn.writeMessage(230, "Password ok, continue")
	} else {
		ftpConn.writeMessage(530, "Incorrect password, not logged in")
	}
}

// cmdSyst responds to the SYST FTP command by providing a canned response.
func (ftpConn *FTPConn) cmdSyst() {
	ftpConn.writeMessage(215, "UNIX Type: L8")
}

func (ftpConn *FTPConn) parseLine(line string) (string, string) {
	params := strings.SplitN(strings.Trim(line, "\r\n"), " ", 2)
	if len(params) == 1 {
		return params[0], ""
	}
	return params[0], params[1]
}

// writeMessage will send a standard FTP response back to the client.
func (ftpConn *FTPConn) writeMessage(code int, message string) (wrote int, err error) {
	line := fmt.Sprintf("%d %s\r\n", code, message)
	wrote, err = ftpConn.controlWriter.WriteString(line)
	ftpConn.controlWriter.Flush()
	log.Print(message)
	return
}
<commit_msg>support the MODE command<commit_after>package graval

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"strings"
)

const (
	welcomeMessage = "Welcome to the Go FTP Server"
)

type FTPConn struct {
	cwd           string
	conn          *net.TCPConn
	controlReader *bufio.Reader
	controlWriter *bufio.Writer
	data          *net.TCPConn
	driver        FTPDriver
	reqUser       string
	user          string
}

// NewFTPConn constructs a new object that will handle the FTP protocol over
// an active net.TCPConn. The TCP connection should already be open before
// it is handed to this functions. driver is an instance of FTPDrive that
// will handle all auth and persistence details.
func NewFTPConn(tcpConn *net.TCPConn, driver FTPDriver) *FTPConn {
	c := new(FTPConn)
	c.cwd = "/"
	c.conn = tcpConn
	c.controlReader = bufio.NewReader(tcpConn)
	c.controlWriter = bufio.NewWriter(tcpConn)
	c.driver = driver
	return c
}

// Serve starts an endless loop that reads FTP commands from the client and
// responds appropriately. terminated is a channel that will receive a true
// message when the connection closes. This loop will be running inside a
// goroutine, so use this channel to be notified when the connection can be
// cleaned up.
func (ftpConn *FTPConn) Serve(terminated chan bool) {
	log.Print("Connection Established")
	// send welcome
	ftpConn.writeMessage(220, welcomeMessage)
	// read commands
	for {
		line, err := ftpConn.controlReader.ReadString('\n')
		if err == nil {
			ftpConn.receiveLine(line)
		}
	}
	terminated <- true
	log.Print("Connection Terminated")
}

// Close will manually close this connection, even if the client isn't ready.
func (ftpConn *FTPConn) Close() {
	ftpConn.conn.Close()
	if ftpConn.data != nil {
		ftpConn.data.Close()
	}
}

// receiveLine accepts a single line FTP command and co-ordinates an
// appropriate response.
func (ftpConn *FTPConn) receiveLine(line string) {
	log.Print(line)
	command, param := ftpConn.parseLine(line)
	switch command {
	case "MODE":
		ftpConn.cmdMode(param)
		break
	case "PASS":
		ftpConn.cmdPass(param)
		break
	case "QUIT":
		ftpConn.Close()
		break
	case "SYST":
		ftpConn.cmdSyst()
		break
	case "USER":
		ftpConn.cmdUser(param)
		break
	default:
		ftpConn.writeMessage(500, "Command not found")
	}
}

// cmdMode responds to the MODE FTP command.
//
// the original FTP spec had various options for hosts to negotiate how data
// would be sent over the data socket, In reality these days (S)tream mode
// is all that is used for the mode - data is just streamed down the data
// socket unchanged.
func (ftpConn *FTPConn) cmdMode(param string) {
	if param == "S" {
		ftpConn.writeMessage(200, "OK")
	} else {
		ftpConn.writeMessage(504, "MODE is an obsolete command")
	}
}

// cmdUser responds to the USER FTP command by asking for the password
func (ftpConn *FTPConn) cmdUser(param string) {
	ftpConn.reqUser = param
	ftpConn.writeMessage(331, "User name ok, password required")
}

// cmdPass respond to the PASS FTP command by asking the driver if the supplied
// username and password are valid
func (ftpConn *FTPConn) cmdPass(param string) {
	if ftpConn.driver.Authenticate(ftpConn.reqUser, param) {
		ftpConn.user = ftpConn.reqUser
		ftpConn.reqUser = ""
		ftpConn.writeMessage(230, "Password ok, continue")
	} else {
		ftpConn.writeMessage(530, "Incorrect password, not logged in")
	}
}

// cmdSyst responds to the SYST FTP command by providing a canned response.
func (ftpConn *FTPConn) cmdSyst() {
	ftpConn.writeMessage(215, "UNIX Type: L8")
}

func (ftpConn *FTPConn) parseLine(line string) (string, string) {
	params := strings.SplitN(strings.Trim(line, "\r\n"), " ", 2)
	if len(params) == 1 {
		return params[0], ""
	}
	return params[0], params[1]
}

// writeMessage will send a standard FTP response back to the client.
func (ftpConn *FTPConn) writeMessage(code int, message string) (wrote int, err error) {
	line := fmt.Sprintf("%d %s\r\n", code, message)
	wrote, err = ftpConn.controlWriter.WriteString(line)
	ftpConn.controlWriter.Flush()
	log.Print(message)
	return
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package crypto

import (
	"bytes"
	"errors"
	"sort"

	stdecdsa "crypto/ecdsa"

	"github.com/decred/dcrd/dcrec/secp256k1/v3/ecdsa"

	secp256k1 "github.com/decred/dcrd/dcrec/secp256k1/v3"

	"github.com/ava-labs/gecko/cache"
	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/utils"
	"github.com/ava-labs/gecko/utils/hashing"
)

const (
	// SECP256K1RSigLen is the number of bytes in a secp2561k recoverable
	// signature
	SECP256K1RSigLen = 65

	// SECP256K1RSKLen is the number of bytes in a secp2561k recoverable private
	// key
	SECP256K1RSKLen = 32

	// SECP256K1RPKLen is the number of bytes in a secp2561k recoverable public
	// key
	SECP256K1RPKLen = 33

	// from the decred library:
	// compactSigMagicOffset is a value used when creating the compact signature
	// recovery code inherited from Bitcoin and has no meaning, but has been
	// retained for compatibility.  For historical purposes, it was originally
	// picked to avoid a binary representation that would allow compact
	// signatures to be mistaken for other components.
	compactSigMagicOffset = 27
)

// FactorySECP256K1R ...
type FactorySECP256K1R struct{ Cache cache.LRU }

// NewPrivateKey implements the Factory interface
func (*FactorySECP256K1R) NewPrivateKey() (PrivateKey, error) {
	k, err := secp256k1.GeneratePrivateKey()
	return &PrivateKeySECP256K1R{sk: k}, err
}

// ToPublicKey implements the Factory interface
func (*FactorySECP256K1R) ToPublicKey(b []byte) (PublicKey, error) {
	key, err := secp256k1.ParsePubKey(b)
	return &PublicKeySECP256K1R{
		pk:    key,
		bytes: b,
	}, err
}

// ToPrivateKey implements the Factory interface
func (*FactorySECP256K1R) ToPrivateKey(b []byte) (PrivateKey, error) {
	return &PrivateKeySECP256K1R{
		sk:    secp256k1.PrivKeyFromBytes(b),
		bytes: b,
	}, nil
}

// RecoverPublicKey returns the public key from a 65 byte signature
func (f *FactorySECP256K1R) RecoverPublicKey(msg, sig []byte) (PublicKey, error) {
	return f.RecoverHashPublicKey(hashing.ComputeHash256(msg), sig)
}

// RecoverHashPublicKey returns the public key from a 65 byte signature
func (f *FactorySECP256K1R) RecoverHashPublicKey(hash, sig []byte) (PublicKey, error) {
	cacheBytes := make([]byte, len(hash)+len(sig))
	copy(cacheBytes, hash)
	copy(cacheBytes[len(hash):], sig)
	id := ids.NewID(hashing.ComputeHash256Array(cacheBytes))
	if cachedPublicKey, ok := f.Cache.Get(id); ok {
		return cachedPublicKey.(*PublicKeySECP256K1R), nil
	}

	if err := verifySECP256K1RSignatureFormat(sig); err != nil {
		return nil, err
	}

	sig, err := sigToRawSig(sig)
	if err != nil {
		return nil, err
	}

	rawPubkey, compressed, err := ecdsa.RecoverCompact(sig, hash)
	if err != nil {
		return nil, err
	}

	if compressed {
		return nil, errors.New("wasn't expecting a compresses key")
	}

	pubkey := &PublicKeySECP256K1R{pk: rawPubkey}
	f.Cache.Put(id, pubkey)
	return pubkey, nil
}

// PublicKeySECP256K1R ...
type PublicKeySECP256K1R struct {
	pk    *secp256k1.PublicKey
	addr  ids.ShortID
	bytes []byte
}

// Verify implements the PublicKey interface
func (k *PublicKeySECP256K1R) Verify(msg, sig []byte) bool {
	return k.VerifyHash(hashing.ComputeHash256(msg), sig)
}

// VerifyHash implements the PublicKey interface
func (k *PublicKeySECP256K1R) VerifyHash(hash, sig []byte) bool {
	factory := FactorySECP256K1R{}
	pk, err := factory.RecoverHashPublicKey(hash, sig)
	if err != nil {
		return false
	}
	return k.Address().Equals(pk.Address())
}

// ToECDSA returns the ecdsa representation of this public key
func (k *PublicKeySECP256K1R) ToECDSA() *stdecdsa.PublicKey {
	return k.pk.ToECDSA()
}

// Address implements the PublicKey interface
func (k *PublicKeySECP256K1R) Address() ids.ShortID {
	if k.addr.IsZero() {
		addr, err := ids.ToShortID(hashing.PubkeyBytesToAddress(k.Bytes()))
		if err != nil {
			panic(err)
		}
		k.addr = addr
	}
	return k.addr
}

// Bytes implements the PublicKey interface
func (k *PublicKeySECP256K1R) Bytes() []byte {
	if k.bytes == nil {
		k.bytes = k.pk.SerializeCompressed()
	}
	return k.bytes
}

// PrivateKeySECP256K1R ...
type PrivateKeySECP256K1R struct {
	sk    *secp256k1.PrivateKey
	pk    *PublicKeySECP256K1R
	bytes []byte
}

// PublicKey implements the PrivateKey interface
func (k *PrivateKeySECP256K1R) PublicKey() PublicKey {
	if k.pk == nil {
		k.pk = &PublicKeySECP256K1R{pk: k.sk.PubKey()}
	}
	return k.pk
}

// Sign implements the PrivateKey interface
func (k *PrivateKeySECP256K1R) Sign(msg []byte) ([]byte, error) {
	return k.SignHash(hashing.ComputeHash256(msg))
}

// SignHash implements the PrivateKey interface
func (k *PrivateKeySECP256K1R) SignHash(hash []byte) ([]byte, error) {
	sig := ecdsa.SignCompact(k.sk, hash, false) // returns [v || r || s]
	return rawSigToSig(sig)
}

// Bytes implements the PrivateKey interface
func (k *PrivateKeySECP256K1R) Bytes() []byte {
	if k.bytes == nil {
		k.bytes = k.sk.Serialize()
	}
	return k.bytes
}

// raw sig has format [v || r || s] whereas the sig has format [r || s || v]
func rawSigToSig(sig []byte) ([]byte, error) {
	if len(sig) != SECP256K1RSigLen {
		return nil, errInvalidSigLen
	}
	recCode := sig[0]
	copy(sig, sig[1:])
	sig[SECP256K1RSigLen-1] = recCode - compactSigMagicOffset
	return sig, nil
}

// sig has format [r || s || v] whereas the raw sig has format [v || r || s]
func sigToRawSig(sig []byte) ([]byte, error) {
	if len(sig) != SECP256K1RSigLen {
		return nil, errInvalidSigLen
	}
	newSig := make([]byte, SECP256K1RSigLen)
	newSig[0] = sig[SECP256K1RSigLen-1] + compactSigMagicOffset
	copy(newSig[1:], sig)
	return newSig, nil
}

// verifies the signature format in format [r || s || v]
func verifySECP256K1RSignatureFormat(sig []byte) error {
	if len(sig) != SECP256K1RSigLen {
		return errInvalidSigLen
	}

	var s secp256k1.ModNScalar
	s.SetByteSlice(sig[32:64])
	if s.IsOverHalfOrder() {
		return errMutatedSig
	}
	return nil
}

type innerSortSECP2561RSigs [][SECP256K1RSigLen]byte

func (lst innerSortSECP2561RSigs) Less(i, j int) bool { return bytes.Compare(lst[i][:], lst[j][:]) < 0 }
func (lst innerSortSECP2561RSigs) Len() int           { return len(lst) }
func (lst innerSortSECP2561RSigs) Swap(i, j int)      { lst[j], lst[i] = lst[i], lst[j] }

// SortSECP2561RSigs sorts a slice of SECP2561R signatures
func SortSECP2561RSigs(lst [][SECP256K1RSigLen]byte) { sort.Sort(innerSortSECP2561RSigs(lst)) }

// IsSortedAndUniqueSECP2561RSigs returns true if [sigs] is sorted
func IsSortedAndUniqueSECP2561RSigs(sigs [][SECP256K1RSigLen]byte) bool {
	return utils.IsSortedAndUnique(innerSortSECP2561RSigs(sigs))
}
<commit_msg>Added ToECDSA to the secp private key<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package crypto

import (
	"bytes"
	"errors"
	"sort"

	stdecdsa "crypto/ecdsa"

	"github.com/decred/dcrd/dcrec/secp256k1/v3/ecdsa"

	secp256k1 "github.com/decred/dcrd/dcrec/secp256k1/v3"

	"github.com/ava-labs/gecko/cache"
	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/utils"
	"github.com/ava-labs/gecko/utils/hashing"
)

const (
	// SECP256K1RSigLen is the number of bytes in a secp2561k recoverable
	// signature
	SECP256K1RSigLen = 65

	// SECP256K1RSKLen is the number of bytes in a secp2561k recoverable private
	// key
	SECP256K1RSKLen = 32

	// SECP256K1RPKLen is the number of bytes in a secp2561k recoverable public
	// key
	SECP256K1RPKLen = 33

	// from the decred library:
	// compactSigMagicOffset is a value used when creating the compact signature
	// recovery code inherited from Bitcoin and has no meaning, but has been
	// retained for compatibility.  For historical purposes, it was originally
	// picked to avoid a binary representation that would allow compact
	// signatures to be mistaken for other components.
	compactSigMagicOffset = 27
)

// FactorySECP256K1R ...
type FactorySECP256K1R struct{ Cache cache.LRU }

// NewPrivateKey implements the Factory interface
func (*FactorySECP256K1R) NewPrivateKey() (PrivateKey, error) {
	k, err := secp256k1.GeneratePrivateKey()
	return &PrivateKeySECP256K1R{sk: k}, err
}

// ToPublicKey implements the Factory interface
func (*FactorySECP256K1R) ToPublicKey(b []byte) (PublicKey, error) {
	key, err := secp256k1.ParsePubKey(b)
	return &PublicKeySECP256K1R{
		pk:    key,
		bytes: b,
	}, err
}

// ToPrivateKey implements the Factory interface
func (*FactorySECP256K1R) ToPrivateKey(b []byte) (PrivateKey, error) {
	return &PrivateKeySECP256K1R{
		sk:    secp256k1.PrivKeyFromBytes(b),
		bytes: b,
	}, nil
}

// RecoverPublicKey returns the public key from a 65 byte signature
func (f *FactorySECP256K1R) RecoverPublicKey(msg, sig []byte) (PublicKey, error) {
	return f.RecoverHashPublicKey(hashing.ComputeHash256(msg), sig)
}

// RecoverHashPublicKey returns the public key from a 65 byte signature
func (f *FactorySECP256K1R) RecoverHashPublicKey(hash, sig []byte) (PublicKey, error) {
	cacheBytes := make([]byte, len(hash)+len(sig))
	copy(cacheBytes, hash)
	copy(cacheBytes[len(hash):], sig)
	id := ids.NewID(hashing.ComputeHash256Array(cacheBytes))
	if cachedPublicKey, ok := f.Cache.Get(id); ok {
		return cachedPublicKey.(*PublicKeySECP256K1R), nil
	}

	if err := verifySECP256K1RSignatureFormat(sig); err != nil {
		return nil, err
	}

	sig, err := sigToRawSig(sig)
	if err != nil {
		return nil, err
	}

	rawPubkey, compressed, err := ecdsa.RecoverCompact(sig, hash)
	if err != nil {
		return nil, err
	}

	if compressed {
		return nil, errors.New("wasn't expecting a compresses key")
	}

	pubkey := &PublicKeySECP256K1R{pk: rawPubkey}
	f.Cache.Put(id, pubkey)
	return pubkey, nil
}

// PublicKeySECP256K1R ...
type PublicKeySECP256K1R struct {
	pk    *secp256k1.PublicKey
	addr  ids.ShortID
	bytes []byte
}

// Verify implements the PublicKey interface
func (k *PublicKeySECP256K1R) Verify(msg, sig []byte) bool {
	return k.VerifyHash(hashing.ComputeHash256(msg), sig)
}

// VerifyHash implements the PublicKey interface
func (k *PublicKeySECP256K1R) VerifyHash(hash, sig []byte) bool {
	factory := FactorySECP256K1R{}
	pk, err := factory.RecoverHashPublicKey(hash, sig)
	if err != nil {
		return false
	}
	return k.Address().Equals(pk.Address())
}

// ToECDSA returns the ecdsa representation of this public key
func (k *PublicKeySECP256K1R) ToECDSA() *stdecdsa.PublicKey {
	return k.pk.ToECDSA()
}

// Address implements the PublicKey interface
func (k *PublicKeySECP256K1R) Address() ids.ShortID {
	if k.addr.IsZero() {
		addr, err := ids.ToShortID(hashing.PubkeyBytesToAddress(k.Bytes()))
		if err != nil {
			panic(err)
		}
		k.addr = addr
	}
	return k.addr
}

// Bytes implements the PublicKey interface
func (k *PublicKeySECP256K1R) Bytes() []byte {
	if k.bytes == nil {
		k.bytes = k.pk.SerializeCompressed()
	}
	return k.bytes
}

// PrivateKeySECP256K1R ...
type PrivateKeySECP256K1R struct {
	sk    *secp256k1.PrivateKey
	pk    *PublicKeySECP256K1R
	bytes []byte
}

// PublicKey implements the PrivateKey interface
func (k *PrivateKeySECP256K1R) PublicKey() PublicKey {
	if k.pk == nil {
		k.pk = &PublicKeySECP256K1R{pk: k.sk.PubKey()}
	}
	return k.pk
}

// Sign implements the PrivateKey interface
func (k *PrivateKeySECP256K1R) Sign(msg []byte) ([]byte, error) {
	return k.SignHash(hashing.ComputeHash256(msg))
}

// SignHash implements the PrivateKey interface
func (k *PrivateKeySECP256K1R) SignHash(hash []byte) ([]byte, error) {
	sig := ecdsa.SignCompact(k.sk, hash, false) // returns [v || r || s]
	return rawSigToSig(sig)
}

// ToECDSA returns the ecdsa representation of this private key
func (k *PrivateKeySECP256K1R) ToECDSA() *stdecdsa.PrivateKey {
	return k.sk.ToECDSA()
}

// Bytes implements the PrivateKey interface
func (k *PrivateKeySECP256K1R) Bytes() []byte {
	if k.bytes == nil {
		k.bytes = k.sk.Serialize()
	}
	return k.bytes
}

// raw sig has format [v || r || s] whereas the sig has format [r || s || v]
func rawSigToSig(sig []byte) ([]byte, error) {
	if len(sig) != SECP256K1RSigLen {
		return nil, errInvalidSigLen
	}
	recCode := sig[0]
	copy(sig, sig[1:])
	sig[SECP256K1RSigLen-1] = recCode - compactSigMagicOffset
	return sig, nil
}

// sig has format [r || s || v] whereas the raw sig has format [v || r || s]
func sigToRawSig(sig []byte) ([]byte, error) {
	if len(sig) != SECP256K1RSigLen {
		return nil, errInvalidSigLen
	}
	newSig := make([]byte, SECP256K1RSigLen)
	newSig[0] = sig[SECP256K1RSigLen-1] + compactSigMagicOffset
	copy(newSig[1:], sig)
	return newSig, nil
}

// verifies the signature format in format [r || s || v]
func verifySECP256K1RSignatureFormat(sig []byte) error {
	if len(sig) != SECP256K1RSigLen {
		return errInvalidSigLen
	}

	var s secp256k1.ModNScalar
	s.SetByteSlice(sig[32:64])
	if s.IsOverHalfOrder() {
		return errMutatedSig
	}
	return nil
}

type innerSortSECP2561RSigs [][SECP256K1RSigLen]byte

func (lst innerSortSECP2561RSigs) Less(i, j int) bool { return bytes.Compare(lst[i][:], lst[j][:]) < 0 }
func (lst innerSortSECP2561RSigs) Len() int           { return len(lst) }
func (lst innerSortSECP2561RSigs) Swap(i, j int)      { lst[j], lst[i] = lst[i], lst[j] }

// SortSECP2561RSigs sorts a slice of SECP2561R signatures
func SortSECP2561RSigs(lst [][SECP256K1RSigLen]byte) { sort.Sort(innerSortSECP2561RSigs(lst)) }

// IsSortedAndUniqueSECP2561RSigs returns true if [sigs] is sorted
func IsSortedAndUniqueSECP2561RSigs(sigs [][SECP256K1RSigLen]byte) bool {
	return utils.IsSortedAndUnique(innerSortSECP2561RSigs(sigs))
}
<|endoftext|>


<commit_before>package utils
import (
	"fmt"
	"time"
)

func GenerateVersion() string {
	return generateDate() + "-" + GitHash();
}

func generateDate() string {
	return fmt.Sprintf("%s", time.Now().Format("20060102.150405"))
}

func GitHash() string {
	out, _ := ExecCmdGetOutput("git", "rev-parse", "--short", "HEAD")
	return out;
}
<commit_msg>add v to version generator<commit_after>package utils
import (
	"fmt"
	"time"
)

func GenerateVersion() string {
	return generateDate() + "-v" + GitHash();
}

func generateDate() string {
	return fmt.Sprintf("%s", time.Now().Format("20060102.150405"))
}

func GitHash() string {
	out, _ := ExecCmdGetOutput("git", "rev-parse", "--short", "HEAD")
	return out;
}
<|endoftext|>


<commit_before>package ginglog

import (
	"time"

	"github.com/gin-gonic/gin"
	"github.com/golang/glog"
)

func setupLogging(duration time.Duration) {
	go func() {
		for _ = range time.Tick(duration) {
			glog.Flush()
		}
	}()
}

var (
	green   = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
	white   = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
	yellow  = string([]byte{27, 91, 57, 55, 59, 52, 51, 109})
	red     = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
	blue    = string([]byte{27, 91, 57, 55, 59, 52, 52, 109})
	magenta = string([]byte{27, 91, 57, 55, 59, 52, 53, 109})
	cyan    = string([]byte{27, 91, 57, 55, 59, 52, 54, 109})
	reset   = string([]byte{27, 91, 48, 109})
)

func ErrorLogger() gin.HandlerFunc {
	return ErrorLoggerT(gin.ErrorTypeAny)
}

func ErrorLoggerT(typ gin.ErrorType) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Next()

		if !c.Writer.Written() {
			json := c.Errors.ByType(typ).JSON()
			if json != nil {
				c.JSON(-1, json)
			}
		}
	}
}

func Logger(duration time.Duration) gin.HandlerFunc {
	setupLogging(duration)
	return func(c *gin.Context) {
		t := time.Now()

		// process request
		c.Next()

		latency := time.Since(t)
		clientIP := c.ClientIP()
		method := c.Request.Method
		statusCode := c.Writer.Status()
		statusColor := colorForStatus(statusCode)
		methodColor := colorForMethod(method)

		glog.V(2).Infof("[GIN] |%s %3d %s| %12v | %s |%s  %s %-7s %s\n%s",
			statusColor, statusCode, reset,
			latency,
			clientIP,
			methodColor, reset, method,
			c.Request.URL.Path,
			c.Errors.String(),
		)

	}
}

func colorForStatus(code int) string {
	switch {
	case code >= 200 && code <= 299:
		return green
	case code >= 300 && code <= 399:
		return white
	case code >= 400 && code <= 499:
		return yellow
	default:
		return red
	}
}

func colorForMethod(method string) string {
	switch {
	case method == "GET":
		return blue
	case method == "POST":
		return cyan
	case method == "PUT":
		return yellow
	case method == "DELETE":
		return red
	case method == "PATCH":
		return green
	case method == "HEAD":
		return magenta
	case method == "OPTIONS":
		return white
	default:
		return reset
	}
}
<commit_msg>Improve error handling<commit_after>package ginglog

import (
	"time"

	"github.com/gin-gonic/gin"
	"github.com/golang/glog"
)

func setupLogging(duration time.Duration) {
	go func() {
		for _ = range time.Tick(duration) {
			glog.Flush()
		}
	}()
}

var (
	green   = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
	white   = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
	yellow  = string([]byte{27, 91, 57, 55, 59, 52, 51, 109})
	red     = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
	blue    = string([]byte{27, 91, 57, 55, 59, 52, 52, 109})
	magenta = string([]byte{27, 91, 57, 55, 59, 52, 53, 109})
	cyan    = string([]byte{27, 91, 57, 55, 59, 52, 54, 109})
	reset   = string([]byte{27, 91, 48, 109})
)

func ErrorLogger() gin.HandlerFunc {
	return ErrorLoggerT(gin.ErrorTypeAny)
}

func ErrorLoggerT(typ gin.ErrorType) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Next()

		if !c.Writer.Written() {
			json := c.Errors.ByType(typ).JSON()
			if json != nil {
				c.JSON(-1, json)
			}
		}
	}
}

func Logger(duration time.Duration) gin.HandlerFunc {
	setupLogging(duration)
	return func(c *gin.Context) {
		t := time.Now()

		// process request
		c.Next()

		latency := time.Since(t)
		clientIP := c.ClientIP()
		method := c.Request.Method
		statusCode := c.Writer.Status()
		statusColor := colorForStatus(statusCode)
		methodColor := colorForMethod(method)
		path := c.Request.URL.Path

		switch {
		case status >= 400 && status <= 499:
			{
				glog.Warningf("[GIN] |%s %3d %s| %12v | %s |%s  %s %-7s %s\n%s",
					statusColor, statusCode, reset,
					latency,
					clientIP,
					methodColor, reset, method,
					path,
					c.Errors.String(),
				)
			}
		case status == 500:
			{
				glog.Errorf("[GIN] |%s %3d %s| %12v | %s |%s  %s %-7s %s\n%s",
					statusColor, statusCode, reset,
					latency,
					clientIP,
					methodColor, reset, method,
					path,
					c.Errors.String(),
				)
			}
		default:
			glog.V(2).Infof("[GIN] |%s %3d %s| %12v | %s |%s  %s %-7s %s\n%s",
				statusColor, statusCode, reset,
				latency,
				clientIP,
				methodColor, reset, method,
				path,
				c.Errors.String(),
			)
		}

	}
}

func colorForStatus(code int) string {
	switch {
	case code >= 200 && code <= 299:
		return green
	case code >= 300 && code <= 399:
		return white
	case code >= 400 && code <= 499:
		return yellow
	default:
		return red
	}
}

func colorForMethod(method string) string {
	switch {
	case method == "GET":
		return blue
	case method == "POST":
		return cyan
	case method == "PUT":
		return yellow
	case method == "DELETE":
		return red
	case method == "PATCH":
		return green
	case method == "HEAD":
		return magenta
	case method == "OPTIONS":
		return white
	default:
		return reset
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
)

type Option struct {
	Name     string
	Question string
	Default  string
}

var options = []Option{
	{
		Name:     "prod",
		Question: "What is your production environment branch?",
		Default:  "master",
	},
	{
		Name:     "other",
		Question: "What other environment branches do you have?",
		Default:  "stage dev",
	},
	{
		Name:     "types",
		Question: "What branch types do you have?",
		Default:  "feature hotfix",
	},
}

type Config struct {
	Prod  string
	Other []string
	Types []string
}

func LoadConfig() Config {
	config = Config{}

	cfg := map[string]string{}

	for _, opt := range options {
		stdout, err := exec.Command("git", "config", "env-branch."+opt.Name).Output()
		if err != nil {
			log.Fatalf("This repo isn't git env enabled. Run 'git env init' first.")
		}
		cfg[opt.Name] = string(stdout)[:len(stdout)-1]
	}
	config.Prod = cfg["prod"]
	config.Other = strings.Split(cfg["other"], " ")
	config.Types = strings.Split(cfg["types"], " ")

	return config
}

func (c Config) IsEnv(branch string) bool {
	if branch == config.Prod {
		return true
	}
	for _, b := range config.Other {
		if branch == b {
			return true
		}
	}
	return false
}

var config Config

func main() {
	if len(os.Args) < 2 {
		help()
	}

	switch os.Args[1] {
	case "init":
		cmdInit()
	case "branch":
		config = LoadConfig()
		cmdBranch(os.Args[2:])
	default:
		help()
	}
}

// Commands

func cmdInit() {
	values := map[string]string{}
	reader := bufio.NewReader(os.Stdin)

	for _, opt := range options {
		fmt.Printf("%s [%s] ", opt.Question, opt.Default)
		value, err := reader.ReadString('\n')
		if err != nil {
			panic(err)
		}
		if value == "\n" {
			values[opt.Name] = opt.Default
		} else {
			values[opt.Name] = value[:len(value)-1]
		}
	}

	for k, v := range values {
		err := exec.Command("git", "config", "--local", "--replace-all", "env-branch."+k, v).Run()
		if err != nil {
			panic(err)
		}
	}

	fmt.Println("You're ready to go.")
}

func cmdBranch(args []string) {
	if len(args) < 1 {
		help()
	}

	switch args[0] {
	case "start":
		if len(args) < 2 {
			help()
		}
		newBranch := args[1]

		gitCommand("checkout", config.Prod)
		gitCommand("pull", "--rebase", "origin", config.Prod)
		gitCommand("checkout", "-b", newBranch)

	case "deploy":

		if len(args) < 2 {
			help()
		}

		deployEnv := args[1]
		var feature string
		var err error

		if len(args) > 2 {
			feature = args[2]
		} else {
			feature, err = getCurrentBranch()
			if err != nil {
				panic(err)
			}
		}

		if !config.IsEnv(deployEnv) {
			log.Fatalf("Branch %s is not an env branch. Can't merge a feature into it.\n", deployEnv)
		}

		if config.IsEnv(feature) {
			log.Fatalf("Branch %s is an env branch. Can't merge an env branch into another env branch.\n", deployEnv)
		}

		gitCommand("checkout", feature)
		gitCommand("pull", "--rebase", "origin", config.Prod)
		gitCommand("checkout", deployEnv)
		gitCommand("pull", "--rebase", "origin", deployEnv)
		gitCommand("merge", feature)
		gitCommand("push", "origin", deployEnv)

	default:
		help()
	}
}

// Everything else

func help() {
	fmt.Println("TODO: implement help")
	os.Exit(1)
}

func gitCommand(args ...string) {
	fmt.Printf("+ git %s\n", strings.Join(args, " "))
	cmd := exec.Command("git", args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	err := cmd.Run()

	if err != nil {
		log.Fatalf("Failed executing command: git %s\n", strings.Join(args, " "))
	}
}

func getCurrentBranch() (string, error) {
	stdout, err := exec.Command("git", "branch").Output()
	if err != nil {
		return "", err
	}

	lines := strings.Split(string(stdout), "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "* ") {
			items := strings.Split(line, " ")

			//TODO
			// return error if the branch is an environment

			return items[1], nil
		}
	}

	return "", errors.New("could not detect current branch")
}
<commit_msg>git env branch start/deploy => git env start/deploy<commit_after>package main

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
)

type Option struct {
	Name     string
	Question string
	Default  string
}

var options = []Option{
	{
		Name:     "prod",
		Question: "What is your production environment branch?",
		Default:  "master",
	},
	{
		Name:     "other",
		Question: "What other environment branches do you have?",
		Default:  "stage dev",
	},
	{
		Name:     "types",
		Question: "What branch types do you have?",
		Default:  "feature hotfix",
	},
}

type Config struct {
	Prod  string
	Other []string
	Types []string
}

func LoadConfig() Config {
	config = Config{}

	cfg := map[string]string{}

	for _, opt := range options {
		stdout, err := exec.Command("git", "config", "env-branch."+opt.Name).Output()
		if err != nil {
			log.Fatalf("This repo isn't git env enabled. Run 'git env init' first.")
		}
		cfg[opt.Name] = string(stdout)[:len(stdout)-1]
	}
	config.Prod = cfg["prod"]
	config.Other = strings.Split(cfg["other"], " ")
	config.Types = strings.Split(cfg["types"], " ")

	return config
}

func (c Config) IsEnv(branch string) bool {
	if branch == config.Prod {
		return true
	}
	for _, b := range config.Other {
		if branch == b {
			return true
		}
	}
	return false
}

var config Config

func main() {
	if len(os.Args) < 2 {
		help()
	}

	switch os.Args[1] {
	case "init":
		cmdInit()
	case "start":
		config = LoadConfig()
		cmdStart(os.Args[2:])
	case "deploy":
		config = LoadConfig()
		cmdDeploy(os.Args[2:])
	default:
		help()
	}
}

// Commands

func cmdInit() {
	values := map[string]string{}
	reader := bufio.NewReader(os.Stdin)

	for _, opt := range options {
		fmt.Printf("%s [%s] ", opt.Question, opt.Default)
		value, err := reader.ReadString('\n')
		if err != nil {
			panic(err)
		}
		if value == "\n" {
			values[opt.Name] = opt.Default
		} else {
			values[opt.Name] = value[:len(value)-1]
		}
	}

	for k, v := range values {
		err := exec.Command("git", "config", "--local", "--replace-all", "env-branch."+k, v).Run()
		if err != nil {
			panic(err)
		}
	}

	fmt.Println("You're ready to go.")
}

func cmdStart(args []string) {
	if len(args) < 1 {
		help()
	}

	newBranch := args[0]

	gitCommand("checkout", config.Prod)
	gitCommand("pull", "--rebase", "origin", config.Prod)
	gitCommand("checkout", "-b", newBranch)
}

func cmdDeploy(args []string) {
	if len(args) < 1 {
		help()
	}

	deployEnv := args[0]
	var feature string
	var err error

	if len(args) > 1 {
		feature = args[1]
	} else {
		feature, err = getCurrentBranch()
		if err != nil {
			panic(err)
		}
	}

	if !config.IsEnv(deployEnv) {
		log.Fatalf("Branch %s is not an env branch. Can't merge a feature into it.\n", deployEnv)
	}

	if config.IsEnv(feature) {
		log.Fatalf("Branch %s is an env branch. Can't merge an env branch into another env branch.\n", deployEnv)
	}

	gitCommand("checkout", feature)
	gitCommand("pull", "--rebase", "origin", config.Prod)
	gitCommand("checkout", deployEnv)
	gitCommand("pull", "--rebase", "origin", deployEnv)
	gitCommand("merge", feature)
	gitCommand("push", "origin", deployEnv)
}

// Everything else

func help() {
	fmt.Println("TODO: implement help")
	os.Exit(1)
}

func gitCommand(args ...string) {
	fmt.Printf("+ git %s\n", strings.Join(args, " "))
	cmd := exec.Command("git", args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	err := cmd.Run()

	if err != nil {
		log.Fatalf("Failed executing command: git %s\n", strings.Join(args, " "))
	}
}

func getCurrentBranch() (string, error) {
	stdout, err := exec.Command("git", "branch").Output()
	if err != nil {
		return "", err
	}

	lines := strings.Split(string(stdout), "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "* ") {
			items := strings.Split(line, " ")

			//TODO
			// return error if the branch is an environment

			return items[1], nil
		}
	}

	return "", errors.New("could not detect current branch")
}
<|endoftext|>


<commit_before>// Copyright 2019 Istio 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 viperconfig

import (
	"fmt"
	"os"
	"reflect"
	"strings"

	"github.com/spf13/viper"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// AddConfigFlag appends a persistent flag for retrieving Viper config, as well as an initializer
// for reading that file
func AddConfigFlag(rootCmd *cobra.Command, viper *viper.Viper) {
	var cfgFile string
	rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "", "Config file containing args")

	cobra.OnInitialize(func() {
		if len(cfgFile) > 0 {
			viper.SetConfigFile(cfgFile)
			err := viper.ReadInConfig() // Find and read the config file
			if err != nil {             // Handle errors reading the config file
				_, _ = os.Stderr.WriteString(fmt.Errorf("fatal error in config file: %s", err).Error())
				os.Exit(1)
			}
		}
	})
}

// ProcessViperConfig retrieves Viper values for each Cobra Val Flag
func ProcessViperConfig(cmd *cobra.Command, viper *viper.Viper) {
	viper.SetTypeByDefaultValue(true)
	cmd.Flags().VisitAll(func(f *pflag.Flag) {
		if reflect.TypeOf(viper.Get(f.Name)).Kind() == reflect.Slice {
			// Viper cannot convert slices to strings, so this is our workaround.
			_ = f.Value.Set(strings.Join(viper.GetStringSlice(f.Name), ","))
		} else {
			_ = f.Value.Set(viper.GetString(f.Name))
		}
	})
}

// ViperizeRootCmd takes a root command, and ensures that all flags of all sub commands
// are bound to viper, with one master config flag accepting a config file.
// At runtime, it will then assign all cobra variables to have their viper values.
func ViperizeRootCmd(cmd *cobra.Command, viper *viper.Viper) {
	AddConfigFlag(cmd, viper)
	subCommandPreRun(cmd, viper)
}

func subCommandPreRun(cmd *cobra.Command, viper *viper.Viper) {
	original := cmd.PreRun
	cmd.PreRun = func(cmd *cobra.Command, args []string) {
		if original != nil {
			original(cmd, args)
		}
		_ = viper.BindPFlags(cmd.Flags())
		ProcessViperConfig(cmd, viper)
	}

	for _, c := range cmd.Commands() {
		subCommandPreRun(c, viper)
	}
}

//ViperizeRootCmdDefault calls ViperizeRootCmd using viper.GetViper()
func ViperizeRootCmdDefault(cmd *cobra.Command) {
	ViperizeRootCmd(cmd, viper.GetViper())
}
<commit_msg>Remove -c shorthand from viperize (#50)<commit_after>// Copyright 2019 Istio 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 viperconfig

import (
	"fmt"
	"os"
	"reflect"
	"strings"

	"github.com/spf13/viper"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// AddConfigFlag appends a persistent flag for retrieving Viper config, as well as an initializer
// for reading that file
func AddConfigFlag(rootCmd *cobra.Command, viper *viper.Viper) {
	var cfgFile string
	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "Config file containing args")

	cobra.OnInitialize(func() {
		if len(cfgFile) > 0 {
			viper.SetConfigFile(cfgFile)
			err := viper.ReadInConfig() // Find and read the config file
			if err != nil {             // Handle errors reading the config file
				_, _ = os.Stderr.WriteString(fmt.Errorf("fatal error in config file: %s", err).Error())
				os.Exit(1)
			}
		}
	})
}

// ProcessViperConfig retrieves Viper values for each Cobra Val Flag
func ProcessViperConfig(cmd *cobra.Command, viper *viper.Viper) {
	viper.SetTypeByDefaultValue(true)
	cmd.Flags().VisitAll(func(f *pflag.Flag) {
		if reflect.TypeOf(viper.Get(f.Name)).Kind() == reflect.Slice {
			// Viper cannot convert slices to strings, so this is our workaround.
			_ = f.Value.Set(strings.Join(viper.GetStringSlice(f.Name), ","))
		} else {
			_ = f.Value.Set(viper.GetString(f.Name))
		}
	})
}

// ViperizeRootCmd takes a root command, and ensures that all flags of all sub commands
// are bound to viper, with one master config flag accepting a config file.
// At runtime, it will then assign all cobra variables to have their viper values.
func ViperizeRootCmd(cmd *cobra.Command, viper *viper.Viper) {
	AddConfigFlag(cmd, viper)
	subCommandPreRun(cmd, viper)
}

func subCommandPreRun(cmd *cobra.Command, viper *viper.Viper) {
	original := cmd.PreRun
	cmd.PreRun = func(cmd *cobra.Command, args []string) {
		if original != nil {
			original(cmd, args)
		}
		_ = viper.BindPFlags(cmd.Flags())
		ProcessViperConfig(cmd, viper)
	}

	for _, c := range cmd.Commands() {
		subCommandPreRun(c, viper)
	}
}

//ViperizeRootCmdDefault calls ViperizeRootCmd using viper.GetViper()
func ViperizeRootCmdDefault(cmd *cobra.Command) {
	ViperizeRootCmd(cmd, viper.GetViper())
}
<|endoftext|>


<commit_before>package git

import (
	"log"
	"os"
	"os/exec"
)

// Init initializes a git repo in the given path.
// If the path does not already exists, it will be created.
func Init(path string) {
	if err := os.MkdirAll(path, os.ModeDir); err != nil {
		log.Fatal(err)
	}
	err := exec.Command("git", "init", path).Run()
	if err != nil {
		log.Fatal(err)
	}
}

// Add wraps the git add call and will change into the path of the git repo
// and add the file given by name.
func Add(path, filename string) {
	os.Chdir(path)
	err := exec.Command("git", "add", filename).Run()
	if err != nil {
		log.Fatal(err)
	}
}

// Commit will change into the path of the git repo and execute git commit.
// In addition a message and date for the commit are specified.
func Commit(path, message, date string) {
	os.Chdir(path)
	err := exec.Command("git", "commit", "-m", message, "--date", date).Run()
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Fix rights for dir creation.<commit_after>package git

import (
	"log"
	"os"
	"os/exec"
)

// permission to create a directory and read/write in it
const PERM = 0755

// Init initializes a git repo in the given path.
// If the path does not already exists, it will be created.
func Init(path string) {
	if err := os.MkdirAll(path, PERM); err != nil {
		log.Fatal(err)
	}
	err := exec.Command("git", "init", path).Run()
	if err != nil {
		log.Fatal(err)
	}
}

// Add wraps the git add call and will change into the path of the git repo
// and add the file given by name.
func Add(path, filename string) {
	os.Chdir(path)
	err := exec.Command("git", "add", filename).Run()
	if err != nil {
		log.Fatal(err)
	}
}

// Commit will change into the path of the git repo and execute git commit.
// In addition a message and date for the commit are specified.
func Commit(path, message, date string) {
	os.Chdir(path)
	err := exec.Command("git", "commit", "-m", message, "--date", date).Run()
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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 service

import (
	"reflect"
	"testing"

	"k8s.io/kubernetes/pkg/api/v1"
)

func buildServiceStatus(ingresses [][]string) v1.LoadBalancerStatus {
	status := v1.LoadBalancerStatus{
		Ingress: []v1.LoadBalancerIngress{},
	}
	for _, element := range ingresses {
		ingress := v1.LoadBalancerIngress{IP: element[0], Hostname: element[1]}
		status.Ingress = append(status.Ingress, ingress)
	}
	return status
}

func TestProcessServiceUpdate(t *testing.T) {
	cc := clusterClientCache{
		clientMap: make(map[string]*clusterCache),
	}
	tests := []struct {
		name             string
		cachedService    *cachedService
		service          *v1.Service
		clusterName      string
		expectNeedUpdate bool
		expectStatus     v1.LoadBalancerStatus
	}{
		{
			"no-cache",
			&cachedService{
				lastState:        &v1.Service{},
				serviceStatusMap: make(map[string]v1.LoadBalancerStatus),
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
			"foo",
			true,
			buildServiceStatus([][]string{{"ip1", ""}}),
		},
		{
			"same-ingress",
			&cachedService{
				lastState: &v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo1": {Ingress: []v1.LoadBalancerIngress{{"ip1", ""}}},
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
			"foo1",
			false,
			buildServiceStatus([][]string{{"ip1", ""}}),
		},
		{
			"diff-cluster",
			&cachedService{
				lastState: &v1.Service{
					ObjectMeta: v1.ObjectMeta{Name: "bar1"},
				},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo2": {Ingress: []v1.LoadBalancerIngress{{"ip1", ""}}},
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
			"foo1",
			true,
			buildServiceStatus([][]string{{"ip1", ""}}),
		},
		{
			"diff-ingress",
			&cachedService{
				lastState: &v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip4", ""}, {"ip1", ""}, {"ip2", ""}})}},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo1": buildServiceStatus([][]string{{"ip4", ""}, {"ip1", ""}, {"ip2", ""}}),
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip2", ""}, {"ip3", ""}, {"ip5", ""}})}},
			"foo1",
			true,
			buildServiceStatus([][]string{{"ip2", ""}, {"ip3", ""}, {"ip5", ""}}),
		},
	}
	for _, test := range tests {
		result := cc.processServiceUpdate(test.cachedService, test.service, test.clusterName)
		if test.expectNeedUpdate != result {
			t.Errorf("Test failed for %s, expected %v, saw %v", test.name, test.expectNeedUpdate, result)
		}
		if !reflect.DeepEqual(test.expectStatus, test.cachedService.lastState.Status.LoadBalancer) {
			t.Errorf("Test failed for %s, expected %v, saw %v", test.name, test.expectStatus, test.cachedService.lastState.Status.LoadBalancer)
		}
	}
}

func TestProcessServiceDeletion(t *testing.T) {
	cc := clusterClientCache{
		clientMap: make(map[string]*clusterCache),
	}
	tests := []struct {
		name             string
		cachedService    *cachedService
		service          *v1.Service
		clusterName      string
		expectNeedUpdate bool
		expectStatus     v1.LoadBalancerStatus
	}{
		{
			"same-ingress",
			&cachedService{
				lastState: &v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo1": {Ingress: []v1.LoadBalancerIngress{{"ip1", ""}}},
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
			"foo1",
			true,
			buildServiceStatus([][]string{}),
		},
		{
			"diff-ingress",
			&cachedService{
				lastState: &v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip4", ""}, {"ip1", ""}, {"ip2", ""}, {"ip3", ""}, {"ip5", ""}, {"ip6", ""}, {"ip8", ""}})}},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo1": buildServiceStatus([][]string{{"ip1", ""}, {"ip2", ""}, {"ip3", ""}}),
					"foo2": buildServiceStatus([][]string{{"ip5", ""}, {"ip6", ""}, {"ip8", ""}}),
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}, {"ip2", ""}, {"ip3", ""}})}},
			"foo1",
			true,
			buildServiceStatus([][]string{{"ip4", ""}, {"ip5", ""}, {"ip6", ""}, {"ip8", ""}}),
		},
	}
	for _, test := range tests {
		result := cc.processServiceDeletion(test.cachedService, test.clusterName)
		if test.expectNeedUpdate != result {
			t.Errorf("Test failed for %s, expected %v, saw %v", test.name, test.expectNeedUpdate, result)
		}
		if !reflect.DeepEqual(test.expectStatus, test.cachedService.lastState.Status.LoadBalancer) {
			t.Errorf("Test failed for %s, expected %+v, saw %+v", test.name, test.expectStatus, test.cachedService.lastState.Status.LoadBalancer)
		}
	}
}
<commit_msg>fix go vet errors<commit_after>/*
Copyright 2016 The Kubernetes 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 service

import (
	"reflect"
	"testing"

	"k8s.io/kubernetes/pkg/api/v1"
)

func buildServiceStatus(ingresses [][]string) v1.LoadBalancerStatus {
	status := v1.LoadBalancerStatus{
		Ingress: []v1.LoadBalancerIngress{},
	}
	for _, element := range ingresses {
		ingress := v1.LoadBalancerIngress{IP: element[0], Hostname: element[1]}
		status.Ingress = append(status.Ingress, ingress)
	}
	return status
}

func TestProcessServiceUpdate(t *testing.T) {
	cc := clusterClientCache{
		clientMap: make(map[string]*clusterCache),
	}
	tests := []struct {
		name             string
		cachedService    *cachedService
		service          *v1.Service
		clusterName      string
		expectNeedUpdate bool
		expectStatus     v1.LoadBalancerStatus
	}{
		{
			"no-cache",
			&cachedService{
				lastState:        &v1.Service{},
				serviceStatusMap: make(map[string]v1.LoadBalancerStatus),
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
			"foo",
			true,
			buildServiceStatus([][]string{{"ip1", ""}}),
		},
		{
			"same-ingress",
			&cachedService{
				lastState: &v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo1": {Ingress: []v1.LoadBalancerIngress{{IP: "ip1", Hostname: ""}}},
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
			"foo1",
			false,
			buildServiceStatus([][]string{{"ip1", ""}}),
		},
		{
			"diff-cluster",
			&cachedService{
				lastState: &v1.Service{
					ObjectMeta: v1.ObjectMeta{Name: "bar1"},
				},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo2": {Ingress: []v1.LoadBalancerIngress{{IP: "ip1", Hostname: ""}}},
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
			"foo1",
			true,
			buildServiceStatus([][]string{{"ip1", ""}}),
		},
		{
			"diff-ingress",
			&cachedService{
				lastState: &v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip4", ""}, {"ip1", ""}, {"ip2", ""}})}},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo1": buildServiceStatus([][]string{{"ip4", ""}, {"ip1", ""}, {"ip2", ""}}),
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip2", ""}, {"ip3", ""}, {"ip5", ""}})}},
			"foo1",
			true,
			buildServiceStatus([][]string{{"ip2", ""}, {"ip3", ""}, {"ip5", ""}}),
		},
	}
	for _, test := range tests {
		result := cc.processServiceUpdate(test.cachedService, test.service, test.clusterName)
		if test.expectNeedUpdate != result {
			t.Errorf("Test failed for %s, expected %v, saw %v", test.name, test.expectNeedUpdate, result)
		}
		if !reflect.DeepEqual(test.expectStatus, test.cachedService.lastState.Status.LoadBalancer) {
			t.Errorf("Test failed for %s, expected %v, saw %v", test.name, test.expectStatus, test.cachedService.lastState.Status.LoadBalancer)
		}
	}
}

func TestProcessServiceDeletion(t *testing.T) {
	cc := clusterClientCache{
		clientMap: make(map[string]*clusterCache),
	}
	tests := []struct {
		name             string
		cachedService    *cachedService
		service          *v1.Service
		clusterName      string
		expectNeedUpdate bool
		expectStatus     v1.LoadBalancerStatus
	}{
		{
			"same-ingress",
			&cachedService{
				lastState: &v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo1": {Ingress: []v1.LoadBalancerIngress{{IP: "ip1", Hostname: ""}}},
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}})}},
			"foo1",
			true,
			buildServiceStatus([][]string{}),
		},
		{
			"diff-ingress",
			&cachedService{
				lastState: &v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip4", ""}, {"ip1", ""}, {"ip2", ""}, {"ip3", ""}, {"ip5", ""}, {"ip6", ""}, {"ip8", ""}})}},
				serviceStatusMap: map[string]v1.LoadBalancerStatus{
					"foo1": buildServiceStatus([][]string{{"ip1", ""}, {"ip2", ""}, {"ip3", ""}}),
					"foo2": buildServiceStatus([][]string{{"ip5", ""}, {"ip6", ""}, {"ip8", ""}}),
				},
			},
			&v1.Service{Status: v1.ServiceStatus{LoadBalancer: buildServiceStatus([][]string{{"ip1", ""}, {"ip2", ""}, {"ip3", ""}})}},
			"foo1",
			true,
			buildServiceStatus([][]string{{"ip4", ""}, {"ip5", ""}, {"ip6", ""}, {"ip8", ""}}),
		},
	}
	for _, test := range tests {
		result := cc.processServiceDeletion(test.cachedService, test.clusterName)
		if test.expectNeedUpdate != result {
			t.Errorf("Test failed for %s, expected %v, saw %v", test.name, test.expectNeedUpdate, result)
		}
		if !reflect.DeepEqual(test.expectStatus, test.cachedService.lastState.Status.LoadBalancer) {
			t.Errorf("Test failed for %s, expected %+v, saw %+v", test.name, test.expectStatus, test.cachedService.lastState.Status.LoadBalancer)
		}
	}
}
<|endoftext|>