<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"mime"
	"net/http"
	"net/http/httputil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/taskcluster/httpbackoff"
	"github.com/taskcluster/taskcluster-client-go/queue"
)

type (
	Artifact interface {
		Name() string
		ProcessResponse() error
		ResponseObject() interface{}
	}

	BaseArtifact struct {
		CanonicalPath string
		Expires       queue.Time
	}

	S3Artifact struct {
		BaseArtifact
		MimeType           string
		S3ArtifactRequest  queue.S3ArtifactRequest
		S3ArtifactResponse queue.S3ArtifactResponse
	}

	AzureArtifact struct {
		BaseArtifact
		MimeType string
	}

	RedirectArtifact struct {
		BaseArtifact
		MimeType string
		URL      string
	}

	ErrorArtifact struct {
		BaseArtifact
		Message string
		Reason  string
	}
)

func (base BaseArtifact) Name() string {
	return base.CanonicalPath
}

func (artifact ErrorArtifact) ProcessResponse() error {
	// TODO: process error response
	return nil
}

func (artifact S3Artifact) ResponseObject() interface{} {
	return new(queue.S3ArtifactResponse)
}

func (artifact ErrorArtifact) ResponseObject() interface{} {
	return new(queue.ErrorArtifactResponse)
}

func (artifact S3Artifact) ProcessResponse() error {
	httpClient := &http.Client{}
	httpCall := func() (*http.Response, error, error) {
		// instead of using fileReader, read it into memory and then use a
		// bytes.Reader since then http.NewRequest will properly set
		// Content-Length header for us, which is needed by the API we call
		fileReader, err := os.Open(filepath.Join(TaskUser.HomeDir, artifact.CanonicalPath))
		requestPayload, err := ioutil.ReadAll(fileReader)
		if err != nil {
			return nil, nil, err
		}
		defer fileReader.Close()
		bytesReader := bytes.NewReader(requestPayload)
		// http.NewRequest automatically sets Content-Length correctly for bytes.Reader
		httpRequest, err := http.NewRequest("PUT", artifact.S3ArtifactResponse.PutUrl, bytesReader)
		if err != nil {
			return nil, nil, err
		}
		debug("MimeType in put request: %v", artifact.MimeType)
		httpRequest.Header.Set("Content-Type", artifact.MimeType)
		// request body could be a) binary and b) massive, so don't show it...
		requestFull, dumpError := httputil.DumpRequestOut(httpRequest, false)
		if dumpError != nil {
			debug("Could not dump request, never mind...")
		} else {
			debug("Request")
			debug(string(requestFull))
		}
		putResp, err := httpClient.Do(httpRequest)
		return putResp, err, nil
	}
	putResp, putAttempts, err := httpbackoff.Retry(httpCall)
	debug("%v put requests issued to %v", putAttempts, artifact.S3ArtifactResponse.PutUrl)
	respBody, dumpError := httputil.DumpResponse(putResp, true)
	if dumpError != nil {
		debug("Could not dump response output, never mind...")
	} else {
		debug("Response")
		debug(string(respBody))
	}
	return err
}

// Returns the artifacts as listed in the payload of the task (note this does
// not include log files)
func (task *TaskRun) PayloadArtifacts() []Artifact {
	artifacts := make([]Artifact, 0)
	debug("Artifacts:")
	for _, artifact := range task.Payload.Artifacts {
		base := BaseArtifact{
			CanonicalPath: canonicalPath(artifact.Path),
			Expires:       artifact.Expires,
		}
		// first check file exists!
		switch artifact.Type {
		case "file":
			artifacts = append(artifacts, resolve(base))
		case "directory":
			walkFn := func(path string, info os.FileInfo, err error) error {
				if !info.IsDir() {
					relativePath, err := filepath.Rel(TaskUser.HomeDir, path)
					if err != nil {
						debug("WIERD ERROR - skipping file: %s", err)
						return nil
					}
					b := BaseArtifact{
						CanonicalPath: relativePath,
						Expires:       artifact.Expires,
					}
					artifacts = append(artifacts, resolve(b))
				}
				return nil
			}
			filepath.Walk(filepath.Join(TaskUser.HomeDir, base.CanonicalPath), walkFn)
		}
	}
	return artifacts
}

// Pass in a BaseArtifact and it will return either an S3 Artifact if file
// exists and is readable, or an ErrorArtifact if not
func resolve(base BaseArtifact) Artifact {
	fileReader, err := os.Open(filepath.Join(TaskUser.HomeDir, base.CanonicalPath))
	defer fileReader.Close()
	if err != nil {
		// cannot read file, create an error artifact
		return ErrorArtifact{
			BaseArtifact: base,
			Message:      "",
			Reason:       "",
		}
	}
	mimeType := mime.TypeByExtension(filepath.Ext(base.CanonicalPath))
	// check we have a mime type!
	if mimeType == "" {
		// application/octet-stream is the mime type for "unknown"
		mimeType = "application/octet-stream"
	}
	return S3Artifact{
		BaseArtifact: base,
		MimeType:     mimeType,
	}
}

// The Queue expects paths to use a forward slash, so let's make sure we have a
// way to generate a path in this format
func canonicalPath(path string) string {
	if os.PathSeparator == '/' {
		return path
	}
	return strings.Replace(path, string(os.PathSeparator), "/", -1)
}

func (task *TaskRun) uploadLog(logFile string) error {
	// logs expire after one year...
	logExpiry := queue.Time(time.Now().AddDate(1, 0, 0))
	log := S3Artifact{
		BaseArtifact: BaseArtifact{
			CanonicalPath: logFile,
			Expires:       logExpiry,
		},
		MimeType: "text/plain",
	}
	return task.uploadArtifact(log)
}

func (task *TaskRun) uploadArtifact(artifact Artifact) error {
	task.Artifacts = append(task.Artifacts, artifact)
	payload, err := json.Marshal(artifact)
	if err != nil {
		return err
	}
	par := queue.PostArtifactRequest(json.RawMessage(payload))
	parsp, callSummary := Queue.CreateArtifact(
		task.TaskId,
		strconv.Itoa(int(task.RunId)),
		artifact.Name(),
		&par,
	)
	if callSummary.Error != nil {
		debug("Could not upload artifact: %v", artifact)
		debug("%v", callSummary)
		debug("%v", parsp)
		debug("Request Headers")
		callSummary.HttpRequest.Header.Write(os.Stdout)
		debug("Request Body")
		debug(callSummary.HttpRequestBody)
		debug("Response Headers")
		callSummary.HttpResponse.Header.Write(os.Stdout)
		debug("Response Body")
		debug(callSummary.HttpResponseBody)
		return callSummary.Error
	}
	debug("Response body RAW")
	debug(callSummary.HttpResponseBody)
	debug("Response body INTERPRETED")
	debug(string(*parsp))
	// unmarshal response into object
	resp := artifact.ResponseObject()
	err = json.Unmarshal(json.RawMessage(*parsp), resp)
	if err != nil {
		return err
	}
	err = artifact.ProcessResponse()
	return err
}
<commit_msg>fixed 'reason' field in ErrorArtifactRequest<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"net/http/httputil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/taskcluster/httpbackoff"
	"github.com/taskcluster/taskcluster-client-go/queue"
)

type (
	Artifact interface {
		Name() string
		ProcessResponse() error
		ResponseObject() interface{}
	}

	BaseArtifact struct {
		CanonicalPath string
		Expires       queue.Time
	}

	S3Artifact struct {
		BaseArtifact
		MimeType           string
		S3ArtifactRequest  queue.S3ArtifactRequest
		S3ArtifactResponse queue.S3ArtifactResponse
	}

	AzureArtifact struct {
		BaseArtifact
		MimeType string
	}

	RedirectArtifact struct {
		BaseArtifact
		MimeType string
		URL      string
	}

	ErrorArtifact struct {
		BaseArtifact
		Message string
		Reason  string
	}
)

func (base BaseArtifact) Name() string {
	return base.CanonicalPath
}

func (artifact ErrorArtifact) ProcessResponse() error {
	// TODO: process error response
	return nil
}

func (artifact S3Artifact) ResponseObject() interface{} {
	return new(queue.S3ArtifactResponse)
}

func (artifact ErrorArtifact) ResponseObject() interface{} {
	return new(queue.ErrorArtifactResponse)
}

func (artifact S3Artifact) ProcessResponse() error {
	httpClient := &http.Client{}
	httpCall := func() (*http.Response, error, error) {
		// instead of using fileReader, read it into memory and then use a
		// bytes.Reader since then http.NewRequest will properly set
		// Content-Length header for us, which is needed by the API we call
		fileReader, err := os.Open(filepath.Join(TaskUser.HomeDir, artifact.CanonicalPath))
		requestPayload, err := ioutil.ReadAll(fileReader)
		if err != nil {
			return nil, nil, err
		}
		defer fileReader.Close()
		bytesReader := bytes.NewReader(requestPayload)
		// http.NewRequest automatically sets Content-Length correctly for bytes.Reader
		httpRequest, err := http.NewRequest("PUT", artifact.S3ArtifactResponse.PutUrl, bytesReader)
		if err != nil {
			return nil, nil, err
		}
		debug("MimeType in put request: %v", artifact.MimeType)
		httpRequest.Header.Set("Content-Type", artifact.MimeType)
		// request body could be a) binary and b) massive, so don't show it...
		requestFull, dumpError := httputil.DumpRequestOut(httpRequest, false)
		if dumpError != nil {
			debug("Could not dump request, never mind...")
		} else {
			debug("Request")
			debug(string(requestFull))
		}
		putResp, err := httpClient.Do(httpRequest)
		return putResp, err, nil
	}
	putResp, putAttempts, err := httpbackoff.Retry(httpCall)
	debug("%v put requests issued to %v", putAttempts, artifact.S3ArtifactResponse.PutUrl)
	respBody, dumpError := httputil.DumpResponse(putResp, true)
	if dumpError != nil {
		debug("Could not dump response output, never mind...")
	} else {
		debug("Response")
		debug(string(respBody))
	}
	return err
}

// Returns the artifacts as listed in the payload of the task (note this does
// not include log files)
func (task *TaskRun) PayloadArtifacts() []Artifact {
	artifacts := make([]Artifact, 0)
	debug("Artifacts:")
	for _, artifact := range task.Payload.Artifacts {
		base := BaseArtifact{
			CanonicalPath: canonicalPath(artifact.Path),
			Expires:       artifact.Expires,
		}
		// first check file exists!
		switch artifact.Type {
		case "file":
			artifacts = append(artifacts, resolve(base))
		case "directory":
			walkFn := func(path string, info os.FileInfo, err error) error {
				if !info.IsDir() {
					relativePath, err := filepath.Rel(TaskUser.HomeDir, path)
					if err != nil {
						debug("WIERD ERROR - skipping file: %s", err)
						return nil
					}
					b := BaseArtifact{
						CanonicalPath: relativePath,
						Expires:       artifact.Expires,
					}
					artifacts = append(artifacts, resolve(b))
				}
				return nil
			}
			filepath.Walk(filepath.Join(TaskUser.HomeDir, base.CanonicalPath), walkFn)
		}
	}
	return artifacts
}

// Pass in a BaseArtifact and it will return either an S3 Artifact if file
// exists and is readable, or an ErrorArtifact if not
func resolve(base BaseArtifact) Artifact {
	fileReader, err := os.Open(filepath.Join(TaskUser.HomeDir, base.CanonicalPath))
	defer fileReader.Close()
	if err != nil {
		// cannot read file, create an error artifact
		return ErrorArtifact{
			BaseArtifact: base,
			Message:      fmt.Sprintf("Could not read file '%s'", fileReader.Name()),
			// TODO: need to also handle "invalid-resource-on-worker"
			// TODO: need to also handle "too-large-file-on-worker"
			Reason: "file-missing-on-worker",
		}
	}
	mimeType := mime.TypeByExtension(filepath.Ext(base.CanonicalPath))
	// check we have a mime type!
	if mimeType == "" {
		// application/octet-stream is the mime type for "unknown"
		mimeType = "application/octet-stream"
	}
	return S3Artifact{
		BaseArtifact: base,
		MimeType:     mimeType,
	}
}

// The Queue expects paths to use a forward slash, so let's make sure we have a
// way to generate a path in this format
func canonicalPath(path string) string {
	if os.PathSeparator == '/' {
		return path
	}
	return strings.Replace(path, string(os.PathSeparator), "/", -1)
}

func (task *TaskRun) uploadLog(logFile string) error {
	// logs expire after one year...
	logExpiry := queue.Time(time.Now().AddDate(1, 0, 0))
	log := S3Artifact{
		BaseArtifact: BaseArtifact{
			CanonicalPath: logFile,
			Expires:       logExpiry,
		},
		MimeType: "text/plain",
	}
	return task.uploadArtifact(log)
}

func (task *TaskRun) uploadArtifact(artifact Artifact) error {
	task.Artifacts = append(task.Artifacts, artifact)
	payload, err := json.Marshal(artifact)
	if err != nil {
		return err
	}
	par := queue.PostArtifactRequest(json.RawMessage(payload))
	parsp, callSummary := Queue.CreateArtifact(
		task.TaskId,
		strconv.Itoa(int(task.RunId)),
		artifact.Name(),
		&par,
	)
	if callSummary.Error != nil {
		debug("Could not upload artifact: %v", artifact)
		debug("%v", callSummary)
		debug("%v", parsp)
		debug("Request Headers")
		callSummary.HttpRequest.Header.Write(os.Stdout)
		debug("Request Body")
		debug(callSummary.HttpRequestBody)
		debug("Response Headers")
		callSummary.HttpResponse.Header.Write(os.Stdout)
		debug("Response Body")
		debug(callSummary.HttpResponseBody)
		return callSummary.Error
	}
	debug("Response body RAW")
	debug(callSummary.HttpResponseBody)
	debug("Response body INTERPRETED")
	debug(string(*parsp))
	// unmarshal response into object
	resp := artifact.ResponseObject()
	err = json.Unmarshal(json.RawMessage(*parsp), resp)
	if err != nil {
		return err
	}
	err = artifact.ProcessResponse()
	return err
}
<|endoftext|>


<commit_before>package mains

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"../dos"
	. "../ifdbg"
	"../interpreter"
	"../lua"
	"../shell"
)

func versionOrStamp() string {
	if Version != "" {
		return Version
	} else {
		return "v" + Stamp
	}
}

func loadBundleScript1(it *shell.Cmd, L lua.Lua, path string) error {
	if DBG {
		println("load cached ", path)
	}
	bin, err := Asset(path)
	if err != nil {
		return err
	}
	err = L.LoadBufferX(path, bin, "t")
	if err != nil {
		return err
	}
	err = NyagosCallLua(L, it, 0, 0)
	if err != nil {
		return err
	}
	return nil
}

type InterpreterT interface {
	Interpret(string) (int, error)
}

func loadScripts(it *shell.Cmd, L lua.Lua) error {
	exeName, exeNameErr := os.Executable()
	if exeNameErr != nil {
		fmt.Fprintln(os.Stderr, exeNameErr)
	}
	exeFolder := filepath.Dir(exeName)

	if !silentmode {
		fmt.Printf("Nihongo Yet Another GOing Shell %s-%s by %s & %s\n",
			versionOrStamp(),
			runtime.GOARCH,
			runtime.Version(),
			"Lua 5.3")
		fmt.Println("(c) 2014-2017 NYAOS.ORG <http://www.nyaos.org>")
	}

	nyagos_d := filepath.Join(exeFolder, "nyagos.d")
	nyagos_d_fd, nyagos_d_err := os.Open(nyagos_d)
	if nyagos_d_err == nil {
		defer nyagos_d_fd.Close()
		finfos, err := nyagos_d_fd.Readdir(-1)
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
		} else {
			for _, finfo1 := range finfos {
				name1 := finfo1.Name()
				if !strings.HasSuffix(strings.ToLower(name1), ".lua") {
					continue
				}
				relpath := "nyagos.d/" + name1
				asset1, assetErr := AssetInfo(relpath)
				if assetErr == nil && asset1.Size() == finfo1.Size() && !asset1.ModTime().Truncate(time.Second).Before(finfo1.ModTime().Truncate(time.Second)) {
					if err := loadBundleScript1(it, L, relpath); err != nil {
						fmt.Fprintf(os.Stderr, "cached %s: %s\n", relpath, err)
					}
				} else {
					path1 := filepath.Join(nyagos_d, name1)
					if DBG {
						println("load real ", path1)
					}
					if err := L.Source(path1); err != nil {
						fmt.Fprintf(os.Stderr, "%s: %s\n", name1, err.Error())
					}
				}
			}
		}
	} else if assertdir, err := AssetDir("nyagos.d"); err == nil {
		// nyagos.d/ not found.
		for _, name1 := range assertdir {
			if !strings.HasSuffix(strings.ToLower(name1), ".lua") {
				continue
			}
			relpath := "nyagos.d/" + name1
			if err1 := loadBundleScript1(it, L, relpath); err1 != nil {
				fmt.Fprintf(os.Stderr, "bundled %s: %s\n", relpath, err1.Error())
			}
		}
	}
	barNyagos(it, exeFolder, L)
	if err := dotNyagos(it, L); err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
	}
	barNyagos(it, dos.GetHome(), L)
	return nil
}

func dotNyagos(it *shell.Cmd, L lua.Lua) error {
	dot_nyagos := filepath.Join(dos.GetHome(), ".nyagos")
	dotStat, dotErr := os.Stat(dot_nyagos)
	if dotErr != nil {
		return nil
	}
	cachePath := filepath.Join(AppDataDir(), "dotnyagos.luac")
	cacheStat, cacheErr := os.Stat(cachePath)
	if cacheErr == nil && !dotStat.ModTime().After(cacheStat.ModTime()) {
		if DBG {
			println("load cached ", cachePath)
		}
		if _, err := L.LoadFile(cachePath, "b"); err == nil {
			NyagosCallLua(L, it, 0, 0)
			return nil
		}
	}
	if DBG {
		println("load real ", dot_nyagos)
	}
	if _, err := L.LoadFile(dot_nyagos, "bt"); err != nil {
		return err
	}
	chank := L.Dump()
	if err := NyagosCallLua(L, it, 0, 0); err != nil {
		return err
	}
	w, w_err := os.Create(cachePath)
	if w_err != nil {
		return w_err
	}
	w.Write(chank)
	w.Close()
	return nil
}

func barNyagos(it InterpreterT, folder string, L lua.Lua) {
	bar_nyagos := filepath.Join(folder, "_nyagos")
	fd, fd_err := os.Open(bar_nyagos)
	if fd_err != nil {
		return
	}
	defer fd.Close()
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		text := scanner.Text()
		text = doLuaFilter(L, text)
		_, err := it.Interpret(text)
		if err != nil {
			fmt.Fprint(os.Stderr, err.Error())
		}
	}
}
<commit_msg>Fix #215: resolve conflict. (wrong import)<commit_after>package mains

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"../dos"
	. "../ifdbg"
	"../lua"
	"../shell"
)

func versionOrStamp() string {
	if Version != "" {
		return Version
	} else {
		return "v" + Stamp
	}
}

func loadBundleScript1(it *shell.Cmd, L lua.Lua, path string) error {
	if DBG {
		println("load cached ", path)
	}
	bin, err := Asset(path)
	if err != nil {
		return err
	}
	err = L.LoadBufferX(path, bin, "t")
	if err != nil {
		return err
	}
	err = NyagosCallLua(L, it, 0, 0)
	if err != nil {
		return err
	}
	return nil
}

type InterpreterT interface {
	Interpret(string) (int, error)
}

func loadScripts(it *shell.Cmd, L lua.Lua) error {
	exeName, exeNameErr := os.Executable()
	if exeNameErr != nil {
		fmt.Fprintln(os.Stderr, exeNameErr)
	}
	exeFolder := filepath.Dir(exeName)

	if !silentmode {
		fmt.Printf("Nihongo Yet Another GOing Shell %s-%s by %s & %s\n",
			versionOrStamp(),
			runtime.GOARCH,
			runtime.Version(),
			"Lua 5.3")
		fmt.Println("(c) 2014-2017 NYAOS.ORG <http://www.nyaos.org>")
	}

	nyagos_d := filepath.Join(exeFolder, "nyagos.d")
	nyagos_d_fd, nyagos_d_err := os.Open(nyagos_d)
	if nyagos_d_err == nil {
		defer nyagos_d_fd.Close()
		finfos, err := nyagos_d_fd.Readdir(-1)
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
		} else {
			for _, finfo1 := range finfos {
				name1 := finfo1.Name()
				if !strings.HasSuffix(strings.ToLower(name1), ".lua") {
					continue
				}
				relpath := "nyagos.d/" + name1
				asset1, assetErr := AssetInfo(relpath)
				if assetErr == nil && asset1.Size() == finfo1.Size() && !asset1.ModTime().Truncate(time.Second).Before(finfo1.ModTime().Truncate(time.Second)) {
					if err := loadBundleScript1(it, L, relpath); err != nil {
						fmt.Fprintf(os.Stderr, "cached %s: %s\n", relpath, err)
					}
				} else {
					path1 := filepath.Join(nyagos_d, name1)
					if DBG {
						println("load real ", path1)
					}
					if err := L.Source(path1); err != nil {
						fmt.Fprintf(os.Stderr, "%s: %s\n", name1, err.Error())
					}
				}
			}
		}
	} else if assertdir, err := AssetDir("nyagos.d"); err == nil {
		// nyagos.d/ not found.
		for _, name1 := range assertdir {
			if !strings.HasSuffix(strings.ToLower(name1), ".lua") {
				continue
			}
			relpath := "nyagos.d/" + name1
			if err1 := loadBundleScript1(it, L, relpath); err1 != nil {
				fmt.Fprintf(os.Stderr, "bundled %s: %s\n", relpath, err1.Error())
			}
		}
	}
	barNyagos(it, exeFolder, L)
	if err := dotNyagos(it, L); err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
	}
	barNyagos(it, dos.GetHome(), L)
	return nil
}

func dotNyagos(it *shell.Cmd, L lua.Lua) error {
	dot_nyagos := filepath.Join(dos.GetHome(), ".nyagos")
	dotStat, dotErr := os.Stat(dot_nyagos)
	if dotErr != nil {
		return nil
	}
	cachePath := filepath.Join(AppDataDir(), "dotnyagos.luac")
	cacheStat, cacheErr := os.Stat(cachePath)
	if cacheErr == nil && !dotStat.ModTime().After(cacheStat.ModTime()) {
		if DBG {
			println("load cached ", cachePath)
		}
		if _, err := L.LoadFile(cachePath, "b"); err == nil {
			NyagosCallLua(L, it, 0, 0)
			return nil
		}
	}
	if DBG {
		println("load real ", dot_nyagos)
	}
	if _, err := L.LoadFile(dot_nyagos, "bt"); err != nil {
		return err
	}
	chank := L.Dump()
	if err := NyagosCallLua(L, it, 0, 0); err != nil {
		return err
	}
	w, w_err := os.Create(cachePath)
	if w_err != nil {
		return w_err
	}
	w.Write(chank)
	w.Close()
	return nil
}

func barNyagos(it InterpreterT, folder string, L lua.Lua) {
	bar_nyagos := filepath.Join(folder, "_nyagos")
	fd, fd_err := os.Open(bar_nyagos)
	if fd_err != nil {
		return
	}
	defer fd.Close()
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		text := scanner.Text()
		text = doLuaFilter(L, text)
		_, err := it.Interpret(text)
		if err != nil {
			fmt.Fprint(os.Stderr, err.Error())
		}
	}
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package manager

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/ernestio/ernest-cli/helper"
	"github.com/fatih/color"
)

// ********************* Login *******************

// Login does a login action against the api
func (m *Manager) Login(username string, password string) (token string, err error) {
	var t Token

	f := url.Values{}
	f.Add("username", username)
	f.Add("password", password)

	url := m.URL + "/auth"
	req, err := http.NewRequest("POST", url, strings.NewReader(f.Encode()))
	req.Form = f
	req.PostForm = f
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("User-Agent", "Ernest/"+m.Version)

	resp, err := m.client().Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		color.Red(err.Error())
	}

	if resp.StatusCode != 200 {
		e := helper.ResponseMessage(body)
		return "", errors.New(e.Message)
	}

	err = json.Unmarshal(body, &t)
	if err != nil {
		color.Red(err.Error())
	}

	token = t.Token

	return token, nil
}
<commit_msg>Change authentication from Forms to JSON payload<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package manager

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"

	"github.com/ernestio/ernest-cli/helper"
	"github.com/fatih/color"
)

// ********************* Login *******************

// Login does a login action against the api
func (m *Manager) Login(username string, password string) (token string, err error) {
	var t Token

	url := m.URL + "/auth"
	body := []byte(`{"username": "` + username + `", "password": "` + password + `"}`)
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(body))
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("User-Agent", "Ernest/"+m.Version)

	resp, err := m.client().Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	body, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		color.Red(err.Error())
	}

	if resp.StatusCode != 200 {
		e := helper.ResponseMessage(body)
		return "", errors.New(e.Message)
	}

	err = json.Unmarshal(body, &t)
	if err != nil {
		color.Red(err.Error())
	}

	token = t.Token

	return token, nil
}
<|endoftext|>


<commit_before>/*
 * portal - marshal
 *
 * a library that implements an algorithm for doing consumer coordination within Kafka, rather
 * than using Zookeeper or another external system.
 *
 */

package marshal

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dropbox/kafka/proto"
)

// Admins will wait a total of consumerReleaseClaimWaitTime
// for paused consumers to release their claims on partitions.
const (
	consumerReleaseClaimWaitSleep = time.Duration(5 * time.Second)
	consumerReleaseClaimWaitTime  = 15 * time.Minute
)

// Admin is used to pause a consumer group and set what position it reads from
// for certain partitions.
type Admin interface {
	// SetConsumerGroupPosition resets consumers to read starting from
	// offsets on each topic, partition pair in positions.
	SetConsumerGroupPosition(groupID string, offsets map[string]map[int]int64) error
}

type consumerGroupAdmin struct {
	clientID     string
	groupID      string
	marshaler    *Marshaler
	pauseTimeout time.Duration

	// claimHealth is 0 if any of our successfully-claimed claims fail to heartbeat.
	claimHealth *int32
	// The lock protects the structs below.
	lock *sync.RWMutex
	// claims are partitions that we've successfully claimed after they've been released,
	// that we'd like to reset the offsets for.
	claims []claimAttempt
	// releaseGroupPartitions keeps track of which Marshal partitions
	// we've produced ReleaseGroup  messages to.
	releaseGroupPartitions []int32
}

// claimAttempt represents a topic, partition we'd like to reset the offset of.
type claimAttempt struct {
	topic  string
	partID int
	// What we'd like to set the offset of a particular partition to be.
	newOffset int64
	// What the current offset of a particular partition is.
	currentOffset int64
}

// addClaimAttempt adds a successfully-claimed partition to our Admin.
func (a *consumerGroupAdmin) addClaimAttempt(topic string,
	partID int, currentOffset, newOffset int64) {

	a.lock.Lock()
	defer a.lock.Unlock()

	c := claimAttempt{
		topic:         topic,
		partID:        partID,
		newOffset:     newOffset,
		currentOffset: currentOffset}
	a.claims = append(a.claims, c)
}

// claimHealth returns whether or not any of the admin's claims have failed to heartbeat.
func (a *consumerGroupAdmin) claimsHealthy() bool {
	return atomic.LoadInt32(a.claimHealth) == 0
}

// NewAdmin returns a new Admin struct bound to a Marshaler. The Marshaler should not have
// any consumers associated with it.
func (m *Marshaler) NewAdmin(groupID string, pauseTimeout time.Duration) (Admin, error) {
	m.lock.RLock()
	defer m.lock.RUnlock()

	if len(m.consumers) != 0 {
		return nil, fmt.Errorf(
			"The Marshaler instance bound to an Admin should not have any consumers.")
	}
	return &consumerGroupAdmin{
		clientID:     m.clientID,
		groupID:      groupID,
		marshaler:    m,
		pauseTimeout: pauseTimeout,

		claimHealth: new(int32),
		lock:        &sync.RWMutex{},
	}, nil
}

// release releases an Admin's claim on a partition. Optionally resets the offset on the partition.
func (a consumerGroupAdmin) release(topic string, partID int, offset int64) bool {
	if err := a.marshaler.ReleasePartition(topic, partID, offset); err != nil {
		log.Errorf("[%s:%d] Admin failed to release partition with offset %d: %s",
			topic, partID, offset, err)
		return false
	}

	return true
}

// releaseClaims releases all claims the Admin has, optionally resetting their offsets.
func (a *consumerGroupAdmin) releaseClaims(resetOffset bool) error {
	a.lock.RLock()
	defer a.lock.RUnlock()

	if !resetOffset {
		log.Infof("Admin releasing claims without resetting offsets.")
	}

	fail := make(chan bool)
	defer close(fail)
	var wg sync.WaitGroup
	for _, claim := range a.claims {
		wg.Add(1)

		releaseOffset := claim.currentOffset
		if resetOffset {
			releaseOffset = claim.newOffset
		}

		go func(t string, p int, offset int64) {
			if ok := a.release(t, p, offset); !ok {
				fail <- true
			}
			wg.Done()
		}(claim.topic, claim.partID, releaseOffset)
	}

	// Wait on all workers to reset their respective Kafka offset.
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-fail:
		return fmt.Errorf("Admin failed to reset Kafka offset!")
	case <-done:
		return nil
	}
}

// heartbeatLoop hearbeats as if we had a claim to this partition and were simply
// not reading past where the previous owner had left off.
func (a *consumerGroupAdmin) heartbeatLoop(
	topic string, partID int, lastOffset int64, stopChan chan struct{}) {

	for {
		select {
		// Stop claimHealth either when all topic, partitions have been successfully claimed,
		// or the Admin has failed to do so and needs to abort.
		case <-stopChan:
			return
		default:
			// If we fail to heartbeat, record this in claimHealth.
			// The Admin will take care of cleaning up other claims.
			if err := a.marshaler.Heartbeat(topic, partID, lastOffset); err != nil {
				log.Errorf("[%s:%d] Admin failed to heartbeat. It is now unhealthy "+
					"and will not reset offsets.", topic, partID)
				atomic.StoreInt32(a.claimHealth, 0)
				return
			}
			time.Sleep(<-a.marshaler.cluster.jitters)
		}
	}
}

// claimAndHeartbeat attempts to claim a partition released by a paused consumer.
// It heartbeats the previous offset.
func (a *consumerGroupAdmin) claimAndHeartbeat(topic string,
	partID int, newOffset int64, stopHeartbeat chan struct{}) bool {

	// Get current offsets, which we will try to keep claimHealth.
	partitionClaim := a.marshaler.GetLastPartitionClaim(topic, partID)

	// Next, try to claim the partition.
	if !a.marshaler.ClaimPartition(topic, partID) {
		log.Errorf("[%s:%d] Admin couldn't claim partition to set Kafka offset",
			topic, partID)
		return false
	}

	// Continuously heartbeat the last offsets.
	a.addClaimAttempt(topic, partID, partitionClaim.CurrentOffset, newOffset)
	go a.heartbeatLoop(topic, partID, partitionClaim.CurrentOffset, stopHeartbeat)
	return true
}

// constructReleaseGroupMessage returns a ReleaseGroup message to write to the Marshal topic.
func (a *consumerGroupAdmin) constructReleaseGroupMessage() *msgReleaseGroup {
	now := time.Now()
	base := &msgBase{
		Time:       int(now.Unix()),
		InstanceID: a.marshaler.instanceID,
		ClientID:   a.clientID,
		GroupID:    a.groupID,
	}
	return &msgReleaseGroup{
		msgBase:       *base,
		MsgExpireTime: int(now.Add(a.pauseTimeout).Unix()),
	}
}

// sendReleaseGroupMessage sends a ReleaseGroup message for a consumer group reading froma given topic.
func (a *consumerGroupAdmin) sendReleaseGroupMessage(topicName string, partID int) error {
	a.lock.Lock()
	defer a.lock.Unlock()
	topic := a.marshaler.cluster.getPartitionState(a.groupID, topicName, partID)

	for _, partition := range a.releaseGroupPartitions {
		if int32(topic.claimPartition) == partition {
			return nil
		}
	}

	a.releaseGroupPartitions = append(a.releaseGroupPartitions, int32(topic.claimPartition))
	rg := a.constructReleaseGroupMessage()
	_, err := a.marshaler.cluster.producer.Produce(MarshalTopic,
		int32(topic.claimPartition), &proto.Message{Value: []byte(rg.Encode())})
	return err
}

// pauseGroupAndWaitForRelease is called for every partition we'd like to change the offset for.
// It first sends a ReleaseGroup message to Marshal, then waits for it to be released,
// then attempts to claim it.
func (a *consumerGroupAdmin) pauseGroupAndWaitForRelease(topicName string, partID int) bool {
	if err := a.sendReleaseGroupMessage(topicName, partID); err != nil {
		log.Errorf("[%s:%d] Admin failed to produce ReleaseMessage group to Kafka: %s",
			topicName, partID, err)
		return false
	}

	// Wait for the paused consumer to release its claim.
	tick := time.NewTicker(consumerReleaseClaimWaitSleep)
	defer tick.Stop()

	select {
	case <-tick.C:
		if cl := a.marshaler.GetPartitionClaim(topicName, partID); cl.LastHeartbeat == 0 {
			break
		}
	case <-time.After(consumerReleaseClaimWaitTime):
		return false
	}
	return true
}

// SetConsumerGroupPosition sets where the consumer group identified by groupID
// should start reading from for given partitions.
func (a *consumerGroupAdmin) SetConsumerGroupPosition(groupID string,
	offsets map[string]map[int]int64) error {

	log.Infof("Admin %s going to pause consumer group %s", a.clientID, groupID)
	var wg sync.WaitGroup
	// Send out a ReleaseGroup message to Marshal for each partition we want to set the position for,
	// then wait for all the partitions to be released.
	fail := make(chan bool)
	defer close(fail)
	for topicName, partitionOffsets := range offsets {
		for partID, _ := range partitionOffsets {
			wg.Add(1)
			go func(topicName string, partID int) {
				if ok := a.pauseGroupAndWaitForRelease(topicName, partID); !ok {
					fail <- true
				}
				wg.Done()
			}(topicName, partID)
		}
	}

	// Wait on all partitions to be released, or one failure.
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-fail:
		return fmt.Errorf("Consumer group %s has not been reset", groupID)
	case <-done:
		break
	}

	// Attempt to claim the now-released partitions, then heartbeat old offsets after a successful claim.
	log.Infof("Admin now claiming released partitions.")
	claimFailures := make(chan bool)
	defer close(claimFailures)

	var claimsWg sync.WaitGroup
	// stopHeartbeating channel instructs all successfully claimed and heartbeating claims to stop.
	stopHeartbeating := make(chan struct{})
	for topicName, partitionOffsets := range offsets {
		for partID, offset := range partitionOffsets {
			claimsWg.Add(1)

			go func(topicName string, partID int, offset int64) {
				if ok := a.claimAndHeartbeat(topicName, partID, offset, stopHeartbeating); !ok {
					claimFailures <- true
				}
				claimsWg.Done()
			}(topicName, partID, offset)
		}
	}

	// Wait on attempts to claim partitions.
	claimsDone := make(chan struct{})
	go func() {
		claimsWg.Wait()
		close(claimsDone)
	}()

	select {
	case <-claimFailures:
		log.Errorf("Couldn't claim a partition -- admin failed to reset consumer group position! " +
			"Now releasing all existing claims without resetting offsets.")
		close(stopHeartbeating)
		return a.releaseClaims(false)
	case <-claimsDone:
		close(stopHeartbeating)
		// Release claims and reset offsets, if all claims have been successfully heartbeating.
		// If not, we'll release claims and not reset offsets.
		return a.releaseClaims(a.claimsHealthy())
	}
}
<commit_msg>Return error when SetConsumerGroupPosition fails.<commit_after>/*
 * portal - marshal
 *
 * a library that implements an algorithm for doing consumer coordination within Kafka, rather
 * than using Zookeeper or another external system.
 *
 */

package marshal

import (
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dropbox/kafka/proto"
)

// Admins will wait a total of consumerReleaseClaimWaitTime
// for paused consumers to release their claims on partitions.
const (
	consumerReleaseClaimWaitSleep = time.Duration(5 * time.Second)
	consumerReleaseClaimWaitTime  = 15 * time.Minute
)

// Admin is used to pause a consumer group and set what position it reads from
// for certain partitions.
type Admin interface {
	// SetConsumerGroupPosition resets consumers to read starting from
	// offsets on each topic, partition pair in positions.
	SetConsumerGroupPosition(groupID string, offsets map[string]map[int]int64) error
}

type consumerGroupAdmin struct {
	clientID     string
	groupID      string
	marshaler    *Marshaler
	pauseTimeout time.Duration

	// claimHealth is 0 if any of our successfully-claimed claims fail to heartbeat.
	claimHealth *int32
	// The lock protects the structs below.
	lock *sync.RWMutex
	// claims are partitions that we've successfully claimed after they've been released,
	// that we'd like to reset the offsets for.
	claims []claimAttempt
	// releaseGroupPartitions keeps track of which Marshal partitions
	// we've produced ReleaseGroup  messages to.
	releaseGroupPartitions []int32
}

// claimAttempt represents a topic, partition we'd like to reset the offset of.
type claimAttempt struct {
	topic  string
	partID int
	// What we'd like to set the offset of a particular partition to be.
	newOffset int64
	// What the current offset of a particular partition is.
	currentOffset int64
}

// addClaimAttempt adds a successfully-claimed partition to our Admin.
func (a *consumerGroupAdmin) addClaimAttempt(topic string,
	partID int, currentOffset, newOffset int64) {

	a.lock.Lock()
	defer a.lock.Unlock()

	c := claimAttempt{
		topic:         topic,
		partID:        partID,
		newOffset:     newOffset,
		currentOffset: currentOffset}
	a.claims = append(a.claims, c)
}

// claimHealth returns whether or not any of the admin's claims have failed to heartbeat.
func (a *consumerGroupAdmin) claimsHealthy() bool {
	return atomic.LoadInt32(a.claimHealth) == 0
}

// NewAdmin returns a new Admin struct bound to a Marshaler. The Marshaler should not have
// any consumers associated with it.
func (m *Marshaler) NewAdmin(groupID string, pauseTimeout time.Duration) (Admin, error) {
	m.lock.RLock()
	defer m.lock.RUnlock()

	if len(m.consumers) != 0 {
		return nil, fmt.Errorf(
			"The Marshaler instance bound to an Admin should not have any consumers.")
	}
	return &consumerGroupAdmin{
		clientID:     m.clientID,
		groupID:      groupID,
		marshaler:    m,
		pauseTimeout: pauseTimeout,

		claimHealth: new(int32),
		lock:        &sync.RWMutex{},
	}, nil
}

// release releases an Admin's claim on a partition. Optionally resets the offset on the partition.
func (a consumerGroupAdmin) release(topic string, partID int, offset int64) bool {
	if err := a.marshaler.ReleasePartition(topic, partID, offset); err != nil {
		log.Errorf("[%s:%d] Admin failed to release partition with offset %d: %s",
			topic, partID, offset, err)
		return false
	}

	return true
}

// releaseClaims releases all claims the Admin has, optionally resetting their offsets.
func (a *consumerGroupAdmin) releaseClaims(resetOffset bool) error {
	a.lock.RLock()
	defer a.lock.RUnlock()

	if !resetOffset {
		log.Infof("Admin releasing claims without resetting offsets.")
	}

	fail := make(chan bool)
	defer close(fail)
	var wg sync.WaitGroup
	for _, claim := range a.claims {
		wg.Add(1)

		releaseOffset := claim.currentOffset
		if resetOffset {
			releaseOffset = claim.newOffset
		}

		go func(t string, p int, offset int64) {
			if ok := a.release(t, p, offset); !ok {
				fail <- true
			}
			wg.Done()
		}(claim.topic, claim.partID, releaseOffset)
	}

	// Wait on all workers to reset their respective Kafka offset.
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-fail:
		return fmt.Errorf("Admin failed to reset Kafka offset!")
	case <-done:
		return nil
	}
}

// heartbeatLoop hearbeats as if we had a claim to this partition and were simply
// not reading past where the previous owner had left off.
func (a *consumerGroupAdmin) heartbeatLoop(
	topic string, partID int, lastOffset int64, stopChan chan struct{}) {

	for {
		select {
		// Stop claimHealth either when all topic, partitions have been successfully claimed,
		// or the Admin has failed to do so and needs to abort.
		case <-stopChan:
			return
		default:
			// If we fail to heartbeat, record this in claimHealth.
			// The Admin will take care of cleaning up other claims.
			if err := a.marshaler.Heartbeat(topic, partID, lastOffset); err != nil {
				log.Errorf("[%s:%d] Admin failed to heartbeat. It is now unhealthy "+
					"and will not reset offsets.", topic, partID)
				atomic.StoreInt32(a.claimHealth, 0)
				return
			}
			time.Sleep(<-a.marshaler.cluster.jitters)
		}
	}
}

// claimAndHeartbeat attempts to claim a partition released by a paused consumer.
// It heartbeats the previous offset.
func (a *consumerGroupAdmin) claimAndHeartbeat(topic string,
	partID int, newOffset int64, stopHeartbeat chan struct{}) bool {

	// Get current offsets, which we will try to keep claimHealth.
	partitionClaim := a.marshaler.GetLastPartitionClaim(topic, partID)

	// Next, try to claim the partition.
	if !a.marshaler.ClaimPartition(topic, partID) {
		log.Errorf("[%s:%d] Admin couldn't claim partition to set Kafka offset",
			topic, partID)
		return false
	}

	// Continuously heartbeat the last offsets.
	a.addClaimAttempt(topic, partID, partitionClaim.CurrentOffset, newOffset)
	go a.heartbeatLoop(topic, partID, partitionClaim.CurrentOffset, stopHeartbeat)
	return true
}

// constructReleaseGroupMessage returns a ReleaseGroup message to write to the Marshal topic.
func (a *consumerGroupAdmin) constructReleaseGroupMessage() *msgReleaseGroup {
	now := time.Now()
	base := &msgBase{
		Time:       int(now.Unix()),
		InstanceID: a.marshaler.instanceID,
		ClientID:   a.clientID,
		GroupID:    a.groupID,
	}
	return &msgReleaseGroup{
		msgBase:       *base,
		MsgExpireTime: int(now.Add(a.pauseTimeout).Unix()),
	}
}

// sendReleaseGroupMessage sends a ReleaseGroup message for a consumer group reading froma given topic.
func (a *consumerGroupAdmin) sendReleaseGroupMessage(topicName string, partID int) error {
	a.lock.Lock()
	defer a.lock.Unlock()
	topic := a.marshaler.cluster.getPartitionState(a.groupID, topicName, partID)

	for _, partition := range a.releaseGroupPartitions {
		if int32(topic.claimPartition) == partition {
			return nil
		}
	}

	a.releaseGroupPartitions = append(a.releaseGroupPartitions, int32(topic.claimPartition))
	rg := a.constructReleaseGroupMessage()
	_, err := a.marshaler.cluster.producer.Produce(MarshalTopic,
		int32(topic.claimPartition), &proto.Message{Value: []byte(rg.Encode())})
	return err
}

// pauseGroupAndWaitForRelease is called for every partition we'd like to change the offset for.
// It first sends a ReleaseGroup message to Marshal, then waits for it to be released,
// then attempts to claim it.
func (a *consumerGroupAdmin) pauseGroupAndWaitForRelease(topicName string, partID int) bool {
	if err := a.sendReleaseGroupMessage(topicName, partID); err != nil {
		log.Errorf("[%s:%d] Admin failed to produce ReleaseMessage group to Kafka: %s",
			topicName, partID, err)
		return false
	}

	// Wait for the paused consumer to release its claim.
	tick := time.NewTicker(consumerReleaseClaimWaitSleep)
	defer tick.Stop()

	select {
	case <-tick.C:
		if cl := a.marshaler.GetPartitionClaim(topicName, partID); cl.LastHeartbeat == 0 {
			break
		}
	case <-time.After(consumerReleaseClaimWaitTime):
		return false
	}
	return true
}

// SetConsumerGroupPosition sets where the consumer group identified by groupID
// should start reading from for given partitions.
func (a *consumerGroupAdmin) SetConsumerGroupPosition(groupID string,
	offsets map[string]map[int]int64) error {

	log.Infof("Admin %s going to pause consumer group %s", a.clientID, groupID)
	var wg sync.WaitGroup
	// Send out a ReleaseGroup message to Marshal for each partition we want to set the position for,
	// then wait for all the partitions to be released.
	fail := make(chan bool)
	defer close(fail)
	for topicName, partitionOffsets := range offsets {
		for partID, _ := range partitionOffsets {
			wg.Add(1)
			go func(topicName string, partID int) {
				if ok := a.pauseGroupAndWaitForRelease(topicName, partID); !ok {
					fail <- true
				}
				wg.Done()
			}(topicName, partID)
		}
	}

	// Wait on all partitions to be released, or one failure.
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-fail:
		return fmt.Errorf("Consumer group %s has not been reset", groupID)
	case <-done:
		break
	}

	// Attempt to claim the now-released partitions, then heartbeat old offsets after a successful claim.
	log.Infof("Admin now claiming released partitions.")
	claimFailures := make(chan bool)
	defer close(claimFailures)

	var claimsWg sync.WaitGroup
	// stopHeartbeating channel instructs all successfully claimed and heartbeating claims to stop.
	stopHeartbeating := make(chan struct{})
	for topicName, partitionOffsets := range offsets {
		for partID, offset := range partitionOffsets {
			claimsWg.Add(1)

			go func(topicName string, partID int, offset int64) {
				if ok := a.claimAndHeartbeat(topicName, partID, offset, stopHeartbeating); !ok {
					claimFailures <- true
				}
				claimsWg.Done()
			}(topicName, partID, offset)
		}
	}

	// Wait on attempts to claim partitions.
	claimsDone := make(chan struct{})
	go func() {
		claimsWg.Wait()
		close(claimsDone)
	}()

	select {
	case <-claimFailures:
		err := errors.New("Couldn't claim a partition -- admin failed to reset consumer group position! " +
			"Now releasing all existing claims without resetting offsets.")
		close(stopHeartbeating)
		a.releaseClaims(false)
		return err
	case <-claimsDone:
		close(stopHeartbeating)
		// Release claims and reset offsets, if all claims have been successfully heartbeating.
		// If not, we'll release claims and not reset offsets.
		return a.releaseClaims(a.claimsHealthy())
	}
}
<|endoftext|>


<commit_before>package master

import (
	"bytes"
	"encoding/json"
	"fmt"
	"math"
	"net/http"
	"time"

	"github.com/aaronang/cong-the-ripper/lib"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

type slave struct {
	tasks    []*lib.Task
	maxTasks int
	// TODO others
}

type scheduler interface {
	schedule(map[string]slave) string
}

type Master struct {
	instances      map[string]slave
	jobs           map[int]*job
	jobsChan       chan lib.Job
	heartbeatChan  chan lib.Heartbeat
	statusChan     chan chan string // dummy
	newTasks       []*lib.Task
	scheduledTasks []*lib.Task
	controllerChan chan string // dummy
	scheduleChan   chan bool   // channel to instruct the main loop to schedule tasks
}

func Init() Master {
	// TODO initialise Master correctly
	return Master{}
}

func (m *Master) Run() {
	http.HandleFunc(lib.JobsCreatePath, m.jobsHandler)
	http.HandleFunc(lib.HeartbeatPath, m.heartbeatHandler)
	http.HandleFunc(lib.StatusPath, m.statusHandler)

	go http.ListenAndServe(lib.Port, nil)
	go func() {
		// TODO Test how this performs when a lot of tasks get submitted.
		time.Sleep(time.Duration(100/len(m.newTasks)) * time.Millisecond)
		m.scheduleChan <- true
	}()

	for {
		select {
		case <-m.scheduleChan:
			// we shedule the tasks when something is in this channel
			// give the controller new data
			// (controller runs in the background and manages the number of instances)
			// call load balancer function to schedule the tasks
			// move tasks from `newTasks` to `scheduledTasks`
			if m.slotsAvailable() {
				if tidx := m.getTaskToSchedule(); tidx != -1 {
					m.scheduleTask(tidx)
				}
			}
		case job := <-m.jobsChan:
			// split the job into tasks
			// update `jobs` and `newTasks`
			_ = job
		case beat := <-m.heartbeatChan:
			// update task statuses
			// check whether a job has completed all its tasks
			_ = beat
		case c := <-m.statusChan:
			// status handler gives us a channel,
			// we write the status into the channel and the the handler "serves" the result
			_ = c
		}
	}
}

func (m *Master) jobsHandler(w http.ResponseWriter, r *http.Request) {
	var j lib.Job
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&j); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	m.jobsChan <- j
}

func (m *Master) heartbeatHandler(w http.ResponseWriter, r *http.Request) {
	var beat lib.Heartbeat
	// TODO parse json and sends the results directly to the  main loop
	m.heartbeatChan <- beat
}

func (m *Master) statusHandler(w http.ResponseWriter, r *http.Request) {
	resultsChan := make(chan string)
	m.statusChan <- resultsChan
	<-resultsChan
	// TODO read the results and serve status page
}

// CreateSlaves creates a new slave instance.
func CreateSlaves(svc *ec2.EC2, count int64) ([]*ec2.Instance, error) {
	params := &ec2.RunInstancesInput{
		ImageId:      aws.String(lib.SlaveImage),
		InstanceType: aws.String(lib.SlaveType),
		MinCount:     aws.Int64(count),
		MaxCount:     aws.Int64(count),
		IamInstanceProfile: &ec2.IamInstanceProfileSpecification{
			Arn: aws.String(lib.SlaveARN),
		},
	}
	resp, err := svc.RunInstances(params)
	return resp.Instances, err
}

// TerminateSlaves terminates a slave instance.
func TerminateSlaves(svc *ec2.EC2, instances []*ec2.Instance) (*ec2.TerminateInstancesOutput, error) {
	params := &ec2.TerminateInstancesInput{
		InstanceIds: instanceIds(instances),
	}
	return svc.TerminateInstances(params)
}

// SendTask sends a task to a slave instance.
func SendTask(t *lib.Task, ip string) (*http.Response, error) {
	url := lib.Protocol + ip + lib.Port + lib.TasksCreatePath
	body, err := t.ToJSON()
	if err != nil {
		panic(err)
	}
	return http.Post(url, lib.BodyType, bytes.NewBuffer(body))
}

func instanceIds(instances []*ec2.Instance) []*string {
	instanceIds := make([]*string, len(instances))
	for i, instance := range instances {
		instanceIds[i] = instance.InstanceId
	}
	return instanceIds
}

func (m *Master) getTaskToSchedule() int {
	for idx, t := range m.newTasks {
		if !m.jobs[t.JobID].reachedMaxTasks() {
			return idx
		}
	}
	return -1
}

func (m *Master) slotsAvailable() bool {
	for _, i := range m.instances {
		if len(i.tasks) < i.maxTasks {
			return true
		}
	}
	return false
}

func (m *Master) scheduleTask(tidx int) {
	minResources := math.MaxInt64
	var slaveIP string
	for k, v := range m.instances {
		if len(v.tasks) < minResources {
			slaveIP = k
		}
	}
	if _, err := SendTask(m.newTasks[tidx], slaveIP); err != nil {
		fmt.Println("Sending task to slave did not execute correctly.")
	} else {
		m.scheduledTasks = append(m.scheduledTasks, m.newTasks[tidx])
		m.newTasks = append(m.newTasks[:tidx], m.newTasks[tidx+1:]...)
	}
}
<commit_msg>PID controller WIP<commit_after>package master

import (
	"bytes"
	"encoding/json"
	"fmt"
	"math"
	"net/http"
	"time"

	"github.com/aaronang/cong-the-ripper/lib"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

type slave struct {
	tasks    []*lib.Task
	maxSlots int
}

type scheduler interface {
	schedule(map[string]slave) string
}

type Master struct {
	instances        map[string]slave
	jobs             map[int]*job
	jobsChan         chan lib.Job
	heartbeatChan    chan lib.Heartbeat
	statusChan       chan chan string // dummy
	newTasks         []*lib.Task
	scheduledTasks   []*lib.Task
	controllerTicker *time.Ticker
	scheduleChan     chan bool // channel to instruct the main loop to schedule tasks
	controller       controller
}

type controller struct {
	dt       time.Duration
	kp       float64
	kd       float64
	ki       float64
	prevErr  float64
	integral float64
}

func Init() Master {
	// TODO initialise Master correctly
	return Master{}
}

func (m *Master) Run() {
	http.HandleFunc(lib.JobsCreatePath, m.jobsHandler)
	http.HandleFunc(lib.HeartbeatPath, m.heartbeatHandler)
	http.HandleFunc(lib.StatusPath, m.statusHandler)

	go http.ListenAndServe(lib.Port, nil)
	go func() {
		// TODO test how this performs when a lot of tasks get submitted.
		time.Sleep(time.Duration(100/len(m.newTasks)) * time.Millisecond)
		m.scheduleChan <- true
	}()

	m.controllerTicker = time.NewTicker(m.controller.dt)

	for {
		select {
		case <-m.controllerTicker.C:
			// run one iteration of the controller
			m.runController()
		case <-m.scheduleChan:
			// we shedule the tasks when something is in this channel
			// give the controller new data
			// (controller runs in the background and manages the number of instances)
			// call load balancer function to schedule the tasks
			// move tasks from `newTasks` to `scheduledTasks`
			if m.slotsAvailable() {
				if tIdx := m.getTaskToSchedule(); tIdx != -1 {
					m.scheduleTask(tIdx)
				}
			}
		case job := <-m.jobsChan:
			// split the job into tasks
			// update `jobs` and `newTasks`
			_ = job
		case beat := <-m.heartbeatChan:
			// update task statuses
			// check whether a job has completed all its tasks
			_ = beat
		case c := <-m.statusChan:
			// status handler gives us a channel,
			// we write the status into the channel and the the handler "serves" the result
			_ = c
		}
	}
}

func (m *Master) jobsHandler(w http.ResponseWriter, r *http.Request) {
	var j lib.Job
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&j); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	m.jobsChan <- j
}

func (m *Master) heartbeatHandler(w http.ResponseWriter, r *http.Request) {
	var beat lib.Heartbeat
	// TODO parse json and sends the results directly to the  main loop
	m.heartbeatChan <- beat
}

func (m *Master) statusHandler(w http.ResponseWriter, r *http.Request) {
	resultsChan := make(chan string)
	m.statusChan <- resultsChan
	<-resultsChan
	// TODO read the results and serve status page
}

// CreateSlaves creates a new slave instance.
func CreateSlaves(svc *ec2.EC2, count int64) ([]*ec2.Instance, error) {
	params := &ec2.RunInstancesInput{
		ImageId:      aws.String(lib.SlaveImage),
		InstanceType: aws.String(lib.SlaveType),
		MinCount:     aws.Int64(count),
		MaxCount:     aws.Int64(count),
		IamInstanceProfile: &ec2.IamInstanceProfileSpecification{
			Arn: aws.String(lib.SlaveARN),
		},
	}
	resp, err := svc.RunInstances(params)
	return resp.Instances, err
}

// TerminateSlaves terminates a slave instance.
func TerminateSlaves(svc *ec2.EC2, instances []*ec2.Instance) (*ec2.TerminateInstancesOutput, error) {
	params := &ec2.TerminateInstancesInput{
		InstanceIds: instanceIds(instances),
	}
	return svc.TerminateInstances(params)
}

// SendTask sends a task to a slave instance.
func SendTask(t *lib.Task, ip string) (*http.Response, error) {
	url := lib.Protocol + ip + lib.Port + lib.TasksCreatePath
	body, err := t.ToJSON()
	if err != nil {
		panic(err)
	}
	return http.Post(url, lib.BodyType, bytes.NewBuffer(body))
}

func instanceIds(instances []*ec2.Instance) []*string {
	instanceIds := make([]*string, len(instances))
	for i, instance := range instances {
		instanceIds[i] = instance.InstanceId
	}
	return instanceIds
}

func (m *Master) getTaskToSchedule() int {
	for idx, t := range m.newTasks {
		if !m.jobs[t.JobID].reachedMaxTasks() {
			return idx
		}
	}
	return -1
}

func (m *Master) slotsAvailable() bool {
	for _, i := range m.instances {
		if len(i.tasks) < i.maxSlots {
			return true
		}
	}
	return false
}

func (m *Master) scheduleTask(tIdx int) {
	minResources := math.MaxInt64
	var slaveIP string
	for k, v := range m.instances {
		if len(v.tasks) < minResources {
			slaveIP = k
		}
	}
	// NOTE: if SendTask takes too long then it may block the main loop
	if _, err := SendTask(m.newTasks[tIdx], slaveIP); err != nil {
		fmt.Println("Sending task to slave did not execute correctly.")
	} else {
		m.scheduledTasks = append(m.scheduledTasks, m.newTasks[tIdx])
		m.newTasks = append(m.newTasks[:tIdx], m.newTasks[tIdx+1:]...)
	}
}

func (m *Master) countTotalSlots() int {
	cnt := 0
	for _, i := range m.instances {
		cnt += i.maxSlots
	}
	return cnt
}

func (m *Master) maxSlots() int {
	// TODO
	return 20 * 2
}

func (m *Master) countRequiredSlots() int {
	cnt := len(m.scheduledTasks)
	cnt += len(m.newTasks)
	if cnt > m.maxSlots() {
		return m.maxSlots()
	}
	return cnt
}

// runController runs one iteration
func (m *Master) runController() float64 {
	err := float64(m.countRequiredSlots() - m.countTotalSlots())

	dt := m.controller.dt.Seconds()
	m.controller.integral = m.controller.integral + err*dt
	derivative := (err - m.controller.prevErr) / dt
	output := m.controller.kp*err +
		m.controller.ki*m.controller.integral +
		m.controller.kd*derivative
	m.controller.prevErr = err

	return output
}

func (m *Master) killSlaves() {

}
<|endoftext|>


<commit_before>// impl generates method stubs for implementing an interface.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"

	"golang.org/x/tools/imports"
)

const usage = `impl [-dir directory] <recv> <iface>

impl generates method stubs for recv to implement iface.

Examples:

impl 'f *File' io.Reader
impl Murmur hash.Hash
impl -dir $GOPATH/src/github.com/josharian/impl Murmur hash.Hash

Don't forget the single quotes around the receiver type
to prevent shell globbing.
`

var (
	flagSrcDir = flag.String("dir", "", "package source directory, useful for vendored code")
)

// findInterface returns the import path and identifier of an interface.
// For example, given "http.ResponseWriter", findInterface returns
// "net/http", "ResponseWriter".
// If a fully qualified interface is given, such as "net/http.ResponseWriter",
// it simply parses the input.
func findInterface(iface string, srcDir string) (path string, id string, err error) {
	if len(strings.Fields(iface)) != 1 {
		return "", "", fmt.Errorf("couldn't parse interface: %s", iface)
	}

	srcPath := filepath.Join(srcDir, "__go_impl__.go")

	if slash := strings.LastIndex(iface, "/"); slash > -1 {
		// package path provided
		dot := strings.LastIndex(iface, ".")
		// make sure iface does not end with "/" (e.g. reject net/http/)
		if slash+1 == len(iface) {
			return "", "", fmt.Errorf("interface name cannot end with a '/' character: %s", iface)
		}
		// make sure iface does not end with "." (e.g. reject net/http.)
		if dot+1 == len(iface) {
			return "", "", fmt.Errorf("interface name cannot end with a '.' character: %s", iface)
		}
		// make sure iface has exactly one "." after "/" (e.g. reject net/http/httputil)
		if strings.Count(iface[slash:], ".") != 1 {
			return "", "", fmt.Errorf("invalid interface name: %s", iface)
		}
		return iface[:dot], iface[dot+1:], nil
	}

	src := []byte("package hack\n" + "var i " + iface)
	// If we couldn't determine the import path, goimports will
	// auto fix the import path.
	imp, err := imports.Process(srcPath, src, nil)
	if err != nil {
		return "", "", fmt.Errorf("couldn't parse interface: %s", iface)
	}

	// imp should now contain an appropriate import.
	// Parse out the import and the identifier.
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, srcPath, imp, 0)
	if err != nil {
		panic(err)
	}
	if len(f.Imports) == 0 {
		return "", "", fmt.Errorf("unrecognized interface: %s", iface)
	}
	raw := f.Imports[0].Path.Value   // "io"
	path, err = strconv.Unquote(raw) // io
	if err != nil {
		panic(err)
	}
	decl := f.Decls[1].(*ast.GenDecl)      // var i io.Reader
	spec := decl.Specs[0].(*ast.ValueSpec) // i io.Reader
	sel := spec.Type.(*ast.SelectorExpr)   // io.Reader
	id = sel.Sel.Name                      // Reader
	return path, id, nil
}

// Pkg is a parsed build.Package.
type Pkg struct {
	*build.Package
	*token.FileSet
}

// Spec is ast.TypeSpec with the associated comment map.
type Spec struct {
	*ast.TypeSpec
	ast.CommentMap
}

// typeSpec locates the *ast.TypeSpec for type id in the import path.
func typeSpec(path string, id string, srcDir string) (Pkg, Spec, error) {
	pkg, err := build.Import(path, srcDir, 0)
	if err != nil {
		return Pkg{}, Spec{}, fmt.Errorf("couldn't find package %s: %v", path, err)
	}

	fset := token.NewFileSet() // share one fset across the whole package
	for _, file := range pkg.GoFiles {
		f, err := parser.ParseFile(fset, filepath.Join(pkg.Dir, file), nil, parser.ParseComments)
		if err != nil {
			continue
		}

		cmap := ast.NewCommentMap(fset, f, f.Comments)

		for _, decl := range f.Decls {
			decl, ok := decl.(*ast.GenDecl)
			if !ok || decl.Tok != token.TYPE {
				continue
			}
			for _, spec := range decl.Specs {
				spec := spec.(*ast.TypeSpec)
				if spec.Name.Name != id {
					continue
				}
				p := Pkg{Package: pkg, FileSet: fset}
				s := Spec{TypeSpec: spec, CommentMap: cmap.Filter(decl)}
				return p, s, nil
			}
		}
	}
	return Pkg{}, Spec{}, fmt.Errorf("type %s not found in %s", id, path)
}

// gofmt pretty-prints e.
func (p Pkg) gofmt(e ast.Expr) string {
	var buf bytes.Buffer
	printer.Fprint(&buf, p.FileSet, e)
	return buf.String()
}

// fullType returns the fully qualified type of e.
// Examples, assuming package net/http:
// 	fullType(int) => "int"
// 	fullType(Handler) => "http.Handler"
// 	fullType(io.Reader) => "io.Reader"
// 	fullType(*Request) => "*http.Request"
func (p Pkg) fullType(e ast.Expr) string {
	ast.Inspect(e, func(n ast.Node) bool {
		switch n := n.(type) {
		case *ast.Ident:
			// Using typeSpec instead of IsExported here would be
			// more accurate, but it'd be crazy expensive, and if
			// the type isn't exported, there's no point trying
			// to implement it anyway.
			if n.IsExported() {
				n.Name = p.Package.Name + "." + n.Name
			}
		case *ast.SelectorExpr:
			return false
		}
		return true
	})
	return p.gofmt(e)
}

func (p Pkg) params(field *ast.Field) []Param {
	var params []Param
	typ := p.fullType(field.Type)
	for _, name := range field.Names {
		params = append(params, Param{Name: name.Name, Type: typ})
	}
	// Handle anonymous params
	if len(params) == 0 {
		params = []Param{Param{Type: typ}}
	}
	return params
}

// Method represents a method signature.
type Method struct {
	Recv string
	Func
}

// Func represents a function signature.
type Func struct {
	Name     string
	Params   []Param
	Res      []Param
	Comments string
}

// Param represents a parameter in a function or method signature.
type Param struct {
	Name string
	Type string
}

func (p Pkg) funcsig(f *ast.Field, cmap ast.CommentMap) Func {
	fn := Func{Name: f.Names[0].Name}
	typ := f.Type.(*ast.FuncType)
	if typ.Params != nil {
		for _, field := range typ.Params.List {
			for _, param := range p.params(field) {
				// only for method parameters:
				// assign a blank identifier "_" to an anonymous parameter
				if param.Name == "" {
					param.Name = "_"
				}
				fn.Params = append(fn.Params, param)
			}
		}
	}
	if typ.Results != nil {
		for _, field := range typ.Results.List {
			fn.Res = append(fn.Res, p.params(field)...)
		}
	}
	if commentsBefore(f, cmap.Comments()) {
		fn.Comments = flattenCommentMap(cmap)
	}
	return fn
}

// The error interface is built-in.
var errorInterface = []Func{{
	Name: "Error",
	Res:  []Param{{Type: "string"}},
}}

// funcs returns the set of methods required to implement iface.
// It is called funcs rather than methods because the
// function descriptions are functions; there is no receiver.
func funcs(iface string, srcDir string) ([]Func, error) {
	// Special case for the built-in error interface.
	if iface == "error" {
		return errorInterface, nil
	}

	// Locate the interface.
	path, id, err := findInterface(iface, srcDir)
	if err != nil {
		return nil, err
	}

	// Parse the package and find the interface declaration.
	p, spec, err := typeSpec(path, id, srcDir)
	if err != nil {
		return nil, fmt.Errorf("interface %s not found: %s", iface, err)
	}
	idecl, ok := spec.Type.(*ast.InterfaceType)
	if !ok {
		return nil, fmt.Errorf("not an interface: %s", iface)
	}

	if idecl.Methods == nil {
		return nil, fmt.Errorf("empty interface: %s", iface)
	}

	var fns []Func
	for _, fndecl := range idecl.Methods.List {
		if len(fndecl.Names) == 0 {
			// Embedded interface: recurse
			embedded, err := funcs(p.fullType(fndecl.Type), srcDir)
			if err != nil {
				return nil, err
			}
			fns = append(fns, embedded...)
			continue
		}

		fn := p.funcsig(fndecl, spec.CommentMap.Filter(fndecl))
		fns = append(fns, fn)
	}
	return fns, nil
}

const stub = "{{if .Comments}}{{.Comments}}{{end}}" +
	"func ({{.Recv}}) {{.Name}}" +
	"({{range .Params}}{{.Name}} {{.Type}}, {{end}})" +
	"({{range .Res}}{{.Name}} {{.Type}}, {{end}})" +
	"{\n" + "panic(\"not implemented\") // TODO: Implement" + "}\n\n"

var tmpl = template.Must(template.New("test").Parse(stub))

// genStubs prints nicely formatted method stubs
// for fns using receiver expression recv.
// If recv is not a valid receiver expression,
// genStubs will panic.
func genStubs(recv string, fns []Func) []byte {
	var buf bytes.Buffer
	for _, fn := range fns {
		meth := Method{Recv: recv, Func: fn}
		tmpl.Execute(&buf, meth)
	}

	pretty, err := format.Source(buf.Bytes())
	if err != nil {
		panic(err)
	}
	return pretty
}

// validReceiver reports whether recv is a valid receiver expression.
func validReceiver(recv string) bool {
	if recv == "" {
		// The parse will parse empty receivers, but we don't want to accept them,
		// since it won't generate a usable code snippet.
		return false
	}
	fset := token.NewFileSet()
	_, err := parser.ParseFile(fset, "", "package hack\nfunc ("+recv+") Foo()", 0)
	return err == nil
}

// commentsBefore reports whether commentGroups precedes a field.
func commentsBefore(field *ast.Field, cg []*ast.CommentGroup) bool {
	if len(cg) > 0 {
		return cg[0].Pos() < field.Pos()
	}
	return false
}

// flattenCommentMap flattens the comment map to a string.
// This function must be used at the point when m is expected to have a single
// element.
func flattenCommentMap(m ast.CommentMap) string {
	if len(m) != 1 {
		panic("flattenCommentMap expects comment map of length 1")
	}
	var result strings.Builder
	for _, cgs := range m {
		for _, cg := range cgs {
			for _, c := range cg.List {
				result.WriteString(c.Text)
				// add an end-of-line character if this is '//'-style comment
				if c.Text[1] == '/' {
					result.WriteString("\n")
				}
			}
		}
	}

	// for '/*'-style comments, make sure to append EOL character to the comment
	// block
	if s := result.String(); !strings.HasSuffix(s, "\n") {
		result.WriteString("\n")
	}

	return result.String()
}

func main() {
	flag.Parse()

	if len(flag.Args()) < 2 {
		fmt.Fprint(os.Stderr, usage)
		os.Exit(2)
	}

	recv, iface := flag.Arg(0), flag.Arg(1)
	if !validReceiver(recv) {
		fatal(fmt.Sprintf("invalid receiver: %q", recv))
	}

	if *flagSrcDir == "" {
		if dir, err := os.Getwd(); err == nil {
			*flagSrcDir = dir
		}
	}

	fns, err := funcs(iface, *flagSrcDir)
	if err != nil {
		fatal(err)
	}

	src := genStubs(recv, fns)
	fmt.Print(string(src))
}

func fatal(msg interface{}) {
	fmt.Fprintln(os.Stderr, msg)
	os.Exit(1)
}
<commit_msg>fix comment<commit_after>// impl generates method stubs for implementing an interface.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"

	"golang.org/x/tools/imports"
)

const usage = `impl [-dir directory] <recv> <iface>

impl generates method stubs for recv to implement iface.

Examples:

impl 'f *File' io.Reader
impl Murmur hash.Hash
impl -dir $GOPATH/src/github.com/josharian/impl Murmur hash.Hash

Don't forget the single quotes around the receiver type
to prevent shell globbing.
`

var (
	flagSrcDir = flag.String("dir", "", "package source directory, useful for vendored code")
)

// findInterface returns the import path and identifier of an interface.
// For example, given "http.ResponseWriter", findInterface returns
// "net/http", "ResponseWriter".
// If a fully qualified interface is given, such as "net/http.ResponseWriter",
// it simply parses the input.
func findInterface(iface string, srcDir string) (path string, id string, err error) {
	if len(strings.Fields(iface)) != 1 {
		return "", "", fmt.Errorf("couldn't parse interface: %s", iface)
	}

	srcPath := filepath.Join(srcDir, "__go_impl__.go")

	if slash := strings.LastIndex(iface, "/"); slash > -1 {
		// package path provided
		dot := strings.LastIndex(iface, ".")
		// make sure iface does not end with "/" (e.g. reject net/http/)
		if slash+1 == len(iface) {
			return "", "", fmt.Errorf("interface name cannot end with a '/' character: %s", iface)
		}
		// make sure iface does not end with "." (e.g. reject net/http.)
		if dot+1 == len(iface) {
			return "", "", fmt.Errorf("interface name cannot end with a '.' character: %s", iface)
		}
		// make sure iface has exactly one "." after "/" (e.g. reject net/http/httputil)
		if strings.Count(iface[slash:], ".") != 1 {
			return "", "", fmt.Errorf("invalid interface name: %s", iface)
		}
		return iface[:dot], iface[dot+1:], nil
	}

	src := []byte("package hack\n" + "var i " + iface)
	// If we couldn't determine the import path, goimports will
	// auto fix the import path.
	imp, err := imports.Process(srcPath, src, nil)
	if err != nil {
		return "", "", fmt.Errorf("couldn't parse interface: %s", iface)
	}

	// imp should now contain an appropriate import.
	// Parse out the import and the identifier.
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, srcPath, imp, 0)
	if err != nil {
		panic(err)
	}
	if len(f.Imports) == 0 {
		return "", "", fmt.Errorf("unrecognized interface: %s", iface)
	}
	raw := f.Imports[0].Path.Value   // "io"
	path, err = strconv.Unquote(raw) // io
	if err != nil {
		panic(err)
	}
	decl := f.Decls[1].(*ast.GenDecl)      // var i io.Reader
	spec := decl.Specs[0].(*ast.ValueSpec) // i io.Reader
	sel := spec.Type.(*ast.SelectorExpr)   // io.Reader
	id = sel.Sel.Name                      // Reader
	return path, id, nil
}

// Pkg is a parsed build.Package.
type Pkg struct {
	*build.Package
	*token.FileSet
}

// Spec is ast.TypeSpec with the associated comment map.
type Spec struct {
	*ast.TypeSpec
	ast.CommentMap
}

// typeSpec locates the *ast.TypeSpec for type id in the import path.
func typeSpec(path string, id string, srcDir string) (Pkg, Spec, error) {
	pkg, err := build.Import(path, srcDir, 0)
	if err != nil {
		return Pkg{}, Spec{}, fmt.Errorf("couldn't find package %s: %v", path, err)
	}

	fset := token.NewFileSet() // share one fset across the whole package
	for _, file := range pkg.GoFiles {
		f, err := parser.ParseFile(fset, filepath.Join(pkg.Dir, file), nil, parser.ParseComments)
		if err != nil {
			continue
		}

		cmap := ast.NewCommentMap(fset, f, f.Comments)

		for _, decl := range f.Decls {
			decl, ok := decl.(*ast.GenDecl)
			if !ok || decl.Tok != token.TYPE {
				continue
			}
			for _, spec := range decl.Specs {
				spec := spec.(*ast.TypeSpec)
				if spec.Name.Name != id {
					continue
				}
				p := Pkg{Package: pkg, FileSet: fset}
				s := Spec{TypeSpec: spec, CommentMap: cmap.Filter(decl)}
				return p, s, nil
			}
		}
	}
	return Pkg{}, Spec{}, fmt.Errorf("type %s not found in %s", id, path)
}

// gofmt pretty-prints e.
func (p Pkg) gofmt(e ast.Expr) string {
	var buf bytes.Buffer
	printer.Fprint(&buf, p.FileSet, e)
	return buf.String()
}

// fullType returns the fully qualified type of e.
// Examples, assuming package net/http:
// 	fullType(int) => "int"
// 	fullType(Handler) => "http.Handler"
// 	fullType(io.Reader) => "io.Reader"
// 	fullType(*Request) => "*http.Request"
func (p Pkg) fullType(e ast.Expr) string {
	ast.Inspect(e, func(n ast.Node) bool {
		switch n := n.(type) {
		case *ast.Ident:
			// Using typeSpec instead of IsExported here would be
			// more accurate, but it'd be crazy expensive, and if
			// the type isn't exported, there's no point trying
			// to implement it anyway.
			if n.IsExported() {
				n.Name = p.Package.Name + "." + n.Name
			}
		case *ast.SelectorExpr:
			return false
		}
		return true
	})
	return p.gofmt(e)
}

func (p Pkg) params(field *ast.Field) []Param {
	var params []Param
	typ := p.fullType(field.Type)
	for _, name := range field.Names {
		params = append(params, Param{Name: name.Name, Type: typ})
	}
	// Handle anonymous params
	if len(params) == 0 {
		params = []Param{Param{Type: typ}}
	}
	return params
}

// Method represents a method signature.
type Method struct {
	Recv string
	Func
}

// Func represents a function signature.
type Func struct {
	Name     string
	Params   []Param
	Res      []Param
	Comments string
}

// Param represents a parameter in a function or method signature.
type Param struct {
	Name string
	Type string
}

func (p Pkg) funcsig(f *ast.Field, cmap ast.CommentMap) Func {
	fn := Func{Name: f.Names[0].Name}
	typ := f.Type.(*ast.FuncType)
	if typ.Params != nil {
		for _, field := range typ.Params.List {
			for _, param := range p.params(field) {
				// only for method parameters:
				// assign a blank identifier "_" to an anonymous parameter
				if param.Name == "" {
					param.Name = "_"
				}
				fn.Params = append(fn.Params, param)
			}
		}
	}
	if typ.Results != nil {
		for _, field := range typ.Results.List {
			fn.Res = append(fn.Res, p.params(field)...)
		}
	}
	if commentsBefore(f, cmap.Comments()) {
		fn.Comments = flattenCommentMap(cmap)
	}
	return fn
}

// The error interface is built-in.
var errorInterface = []Func{{
	Name: "Error",
	Res:  []Param{{Type: "string"}},
}}

// funcs returns the set of methods required to implement iface.
// It is called funcs rather than methods because the
// function descriptions are functions; there is no receiver.
func funcs(iface string, srcDir string) ([]Func, error) {
	// Special case for the built-in error interface.
	if iface == "error" {
		return errorInterface, nil
	}

	// Locate the interface.
	path, id, err := findInterface(iface, srcDir)
	if err != nil {
		return nil, err
	}

	// Parse the package and find the interface declaration.
	p, spec, err := typeSpec(path, id, srcDir)
	if err != nil {
		return nil, fmt.Errorf("interface %s not found: %s", iface, err)
	}
	idecl, ok := spec.Type.(*ast.InterfaceType)
	if !ok {
		return nil, fmt.Errorf("not an interface: %s", iface)
	}

	if idecl.Methods == nil {
		return nil, fmt.Errorf("empty interface: %s", iface)
	}

	var fns []Func
	for _, fndecl := range idecl.Methods.List {
		if len(fndecl.Names) == 0 {
			// Embedded interface: recurse
			embedded, err := funcs(p.fullType(fndecl.Type), srcDir)
			if err != nil {
				return nil, err
			}
			fns = append(fns, embedded...)
			continue
		}

		fn := p.funcsig(fndecl, spec.CommentMap.Filter(fndecl))
		fns = append(fns, fn)
	}
	return fns, nil
}

const stub = "{{if .Comments}}{{.Comments}}{{end}}" +
	"func ({{.Recv}}) {{.Name}}" +
	"({{range .Params}}{{.Name}} {{.Type}}, {{end}})" +
	"({{range .Res}}{{.Name}} {{.Type}}, {{end}})" +
	"{\n" + "panic(\"not implemented\") // TODO: Implement" + "\n}\n\n"

var tmpl = template.Must(template.New("test").Parse(stub))

// genStubs prints nicely formatted method stubs
// for fns using receiver expression recv.
// If recv is not a valid receiver expression,
// genStubs will panic.
func genStubs(recv string, fns []Func) []byte {
	var buf bytes.Buffer
	for _, fn := range fns {
		meth := Method{Recv: recv, Func: fn}
		tmpl.Execute(&buf, meth)
	}

	pretty, err := format.Source(buf.Bytes())
	if err != nil {
		panic(err)
	}
	return pretty
}

// validReceiver reports whether recv is a valid receiver expression.
func validReceiver(recv string) bool {
	if recv == "" {
		// The parse will parse empty receivers, but we don't want to accept them,
		// since it won't generate a usable code snippet.
		return false
	}
	fset := token.NewFileSet()
	_, err := parser.ParseFile(fset, "", "package hack\nfunc ("+recv+") Foo()", 0)
	return err == nil
}

// commentsBefore reports whether commentGroups precedes a field.
func commentsBefore(field *ast.Field, cg []*ast.CommentGroup) bool {
	if len(cg) > 0 {
		return cg[0].Pos() < field.Pos()
	}
	return false
}

// flattenCommentMap flattens the comment map to a string.
// This function must be used at the point when m is expected to have a single
// element.
func flattenCommentMap(m ast.CommentMap) string {
	if len(m) != 1 {
		panic("flattenCommentMap expects comment map of length 1")
	}
	var result strings.Builder
	for _, cgs := range m {
		for _, cg := range cgs {
			for _, c := range cg.List {
				result.WriteString(c.Text)
				// add an end-of-line character if this is '//'-style comment
				if c.Text[1] == '/' {
					result.WriteString("\n")
				}
			}
		}
	}

	// for '/*'-style comments, make sure to append EOL character to the comment
	// block
	if s := result.String(); !strings.HasSuffix(s, "\n") {
		result.WriteString("\n")
	}

	return result.String()
}

func main() {
	flag.Parse()

	if len(flag.Args()) < 2 {
		fmt.Fprint(os.Stderr, usage)
		os.Exit(2)
	}

	recv, iface := flag.Arg(0), flag.Arg(1)
	if !validReceiver(recv) {
		fatal(fmt.Sprintf("invalid receiver: %q", recv))
	}

	if *flagSrcDir == "" {
		if dir, err := os.Getwd(); err == nil {
			*flagSrcDir = dir
		}
	}

	fns, err := funcs(iface, *flagSrcDir)
	if err != nil {
		fatal(err)
	}

	src := genStubs(recv, fns)
	fmt.Print(string(src))
}

func fatal(msg interface{}) {
	fmt.Fprintln(os.Stderr, msg)
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

//import "gopkg.in/yaml.v2"
import "github.com/spf13/cobra"
import "github.com/peterh/liner"
import "fmt"
import "os"

func init() {
	initCmd := &cobra.Command{
		Use:   "init",
		Short: "initializes a config for your function",
		Run:   initPhage,
	}

	cmds = append(cmds, initCmd)
}

type prompt struct {
	text           string
	def            string
	stringStore    **string
	stringSetStore *[]*string
}

// helps you build a config file
func initPhage(c *cobra.Command, _ []string) {
	l := liner.NewLiner()

	fmt.Println(`
		HELLO AND WELCOME

		This command will help you set up your code for deployment to lambda!
		Please answer the prompts as they appear below:
	`)

	//reqMsg := "Sorry, that field is required. Try again."

	cfg := new(Config)
	wd, _ := os.Getwd()
	st, _ := os.Stat(wd)

	prompts := []prompt{
		prompt{
			"Enter a project name",
			st.Name(),
			&cfg.Name,
			nil,
		},
	}

	for _, p := range prompts {
		text := p.text
		if p.def != "" {
			text += " [" + p.def + "]"
		}

		text += ": "
		if s, err := l.Prompt(text); err == nil {
			if p.stringStore != nil {
				*p.stringStore = &s
			}
		} else if err == liner.ErrPromptAborted {
			fmt.Println("Aborted")
			return
		} else {
			fmt.Println("Error reading line: ", err)
			return
		}
	}

	l.Close()
}
<commit_msg>add more prompts for init<commit_after>package main

import "gopkg.in/yaml.v2"
import "github.com/spf13/cobra"
import "github.com/peterh/liner"
import "strconv"
import "strings"
import "fmt"
import "io/ioutil"
import "os"

func init() {
	initCmd := &cobra.Command{
		Use:   "init",
		Short: "initializes a config for your function",
		Run:   initPhage,
	}

	cmds = append(cmds, initCmd)
}

type prompt struct {
	text           string
	def            string
	stringStore    **string
	stringSetStore *[]*string
	intStore       **int64
}

// helps you build a config file
func initPhage(c *cobra.Command, _ []string) {
	l := liner.NewLiner()

	fmt.Println(`
		HELLO AND WELCOME

		This command will help you set up your code for deployment to lambda!
		Please answer the prompts as they appear below:
	`)

	//reqMsg := "Sorry, that field is required. Try again."

	cfg := new(Config)
	cfg.IamRole = new(IamRole)
	cfg.Location = new(Location)
	wd, _ := os.Getwd()
	st, _ := os.Stat(wd)

	prompts := []prompt{
		prompt{
			"Enter a project name",
			st.Name(),
			&cfg.Name,
			nil,
			nil,
		},
		prompt{
			"Enter a project description if you'd like",
			"",
			&cfg.Description,
			nil,
			nil,
		},
		prompt{
			"Enter an archive name if you'd like",
			st.Name() + ".zip",
			&cfg.Archive,
			nil,
			nil,
		},
		prompt{
			"What runtime are you using: nodejs, java8, or python 2.7?",
			"nodejs",
			&cfg.Runtime,
			nil,
			nil,
		},

		prompt{
			"Enter an entry point or handler name",
			"index.handler",
			&cfg.EntryPoint,
			nil,
			nil,
		},

		prompt{
			"Enter memory size",
			"128",
			nil,
			nil,
			&cfg.MemorySize,
		},
		prompt{
			"Enter timeout",
			"5",
			nil,
			nil,
			&cfg.Timeout,
		},
		prompt{
			"Enter AWS regions where this function will run",
			"us-east-1",
			nil,
			&cfg.Regions,
			nil,
		},
		prompt{
			"Enter IAM role name",
			"us-east-1",
			&cfg.IamRole.Name,
			nil,
			nil,
		},
	}

	for _, cPrompt := range prompts {
		p := cPrompt
		text := p.text
		if p.def != "" {
			text += " [" + p.def + "]"
		}

		text += ": "
		if s, err := l.Prompt(text); err == nil {
			input := s
			hasInput := input != ""
			if p.stringStore != nil {
				if hasInput {
					*p.stringStore = &input
				} else {
					*p.stringStore = &p.def
				}

			} else if p.stringSetStore != nil {
				var splitMe string
				if hasInput {
					splitMe = input
				} else {
					splitMe = p.def
				}

				spl := strings.Split(splitMe, ",")
				pspl := make([]*string, len(spl))
				for i, v := range spl {
					// we need to set the value
					// in a variable local to this block
					// because the pointed-to value in
					// `v` will change on the next
					// loop iteration
					realVal := v
					pspl[i] = &realVal
				}

				*p.stringSetStore = pspl

			} else if p.intStore != nil {
				var tParse string
				if hasInput {
					tParse = input
				} else {
					tParse = p.def
				}

				i, _ := strconv.ParseInt(tParse, 10, 64)
				*p.intStore = &i
			}
		} else if err == liner.ErrPromptAborted {
			fmt.Println("Aborted")
			return
		} else {
			fmt.Println("Error reading line: ", err)
			return
		}
	}

	l.Close()

	d, err := yaml.Marshal(cfg)
	if err != nil {
		panic(err)
	}
	ioutil.WriteFile("l-p.yml", d, os.FileMode(0644))
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/gramework/gramework"
)

func setup(app *gramework.App) {
	app.TLSEmails = []string{
		"k@gramework.win",
	}

	app.EnableFirewall = true
	app.Settings.Firewall = gramework.FirewallSettings{
		MaxReqPerMin: 1 << 28,
		BlockTimeout: int64(15 * time.Second),
	}

	app.GET("/*any", func(ctx *gramework.Context) {
		if ctx.RouteArg("any") == "/f777d0332159.html" {
			ctx.Writef("ec3f5942dd6a")
			return
		}
		app.Forbidden(ctx)
	})
}
<commit_msg>block: log blocks<commit_after>package main

import (
	"time"

	"github.com/apex/log"

	"github.com/gramework/gramework"
)

func setup(app *gramework.App) {
	app.TLSEmails = []string{
		"k@gramework.win",
	}

	app.EnableFirewall = true
	app.Settings.Firewall = gramework.FirewallSettings{
		MaxReqPerMin: 1 << 28,
		BlockTimeout: int64(15 * time.Second),
	}

	app.GET("/*any", func(ctx *gramework.Context) {
		if ctx.RouteArg("any") == "/f777d0332159.html" {
			ctx.Writef("ec3f5942dd6a")
			return
		}
		app.Forbidden(ctx)
		app.Logger.WithFields(log.Fields{
			"ip":  ctx.RemoteIP().String(),
			"url": ctx.URI().String(),
		}).Infof("blocked")
	})
}
<|endoftext|>


<commit_before>// Copyright ©2012 Dan Kortschak <dan.kortschak@adelaide.edu.au>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package boom

import ()

type BAMFile struct {
	*samFile
}

func OpenBAM(filename string) (b *BAMFile, err error) {
	sf, err := samOpen(filename, "rb", nil)
	if err != nil {
		return
	}
	return &BAMFile{sf}, nil
}

var bWModes = [2]string{"wb", "wbu"}

func CreateBAM(filename string, ref *Header, comp bool) (b *BAMFile, err error) {
	var mode string
	if comp {
		mode = bWModes[0]
	} else {
		mode = bWModes[1]
	}
	sf, err := samOpen(filename, mode, ref.bamHeader)
	if err != nil {
		return
	}
	return &BAMFile{sf}, nil
}

func (self *BAMFile) Close() error {
	return self.samClose()
}

func (self *BAMFile) Read() (r *Record, n int, err error) {
	n, br, err := self.samRead()
	r = &Record{bamRecord: br}
	return
}

func (self *BAMFile) Write(r *Record) (n int, err error) {
	return self.samWrite(r.bamRecord)
}

// GetTargetID returns the tid corresponding to the string chr and true if a match is present.
// If no matching tid is found -1 and false are returned.
func (self *BAMFile) GetTargetID(chr string) (id int, ok bool) {
	id = self.header().bamGetTid(chr)
	if id < 0 {
		return
	}
	ok = true

	return
}

func (self *BAMFile) ReferenceNames() []string {
	return self.header().targetNames()
}

func (self *BAMFile) ReferenceLengths() []uint32 {
	return self.header().targetLengths()
}

func (self *BAMFile) Text() string {
	return self.header().text()
}

// A FetchFn is called on each Record found by Fetch.
type FetchFn func(*Record)

// Fetch calls fn on all BAM records within the interval [beg, end) of the reference sequence
// identified by chr. Note that beg >= 0 || beg = 0.
func (self *BAMFile) Fetch(i *Index, tid int, beg, end int, fn FetchFn) (ret int, err error) {
	f := func(b *bamRecord) {
		fn(&Record{bamRecord: b})
	}

	return self.bamFetch(i.bamIndex, tid, beg, end, f)
}
<commit_msg>Unnecessary import line removed<commit_after>// Copyright ©2012 Dan Kortschak <dan.kortschak@adelaide.edu.au>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package boom

type BAMFile struct {
	*samFile
}

func OpenBAM(filename string) (b *BAMFile, err error) {
	sf, err := samOpen(filename, "rb", nil)
	if err != nil {
		return
	}
	return &BAMFile{sf}, nil
}

var bWModes = [2]string{"wb", "wbu"}

func CreateBAM(filename string, ref *Header, comp bool) (b *BAMFile, err error) {
	var mode string
	if comp {
		mode = bWModes[0]
	} else {
		mode = bWModes[1]
	}
	sf, err := samOpen(filename, mode, ref.bamHeader)
	if err != nil {
		return
	}
	return &BAMFile{sf}, nil
}

func (self *BAMFile) Close() error {
	return self.samClose()
}

func (self *BAMFile) Read() (r *Record, n int, err error) {
	n, br, err := self.samRead()
	r = &Record{bamRecord: br}
	return
}

func (self *BAMFile) Write(r *Record) (n int, err error) {
	return self.samWrite(r.bamRecord)
}

// GetTargetID returns the tid corresponding to the string chr and true if a match is present.
// If no matching tid is found -1 and false are returned.
func (self *BAMFile) GetTargetID(chr string) (id int, ok bool) {
	id = self.header().bamGetTid(chr)
	if id < 0 {
		return
	}
	ok = true

	return
}

func (self *BAMFile) ReferenceNames() []string {
	return self.header().targetNames()
}

func (self *BAMFile) ReferenceLengths() []uint32 {
	return self.header().targetLengths()
}

func (self *BAMFile) Text() string {
	return self.header().text()
}

// A FetchFn is called on each Record found by Fetch.
type FetchFn func(*Record)

// Fetch calls fn on all BAM records within the interval [beg, end) of the reference sequence
// identified by chr. Note that beg >= 0 || beg = 0.
func (self *BAMFile) Fetch(i *Index, tid int, beg, end int, fn FetchFn) (ret int, err error) {
	f := func(b *bamRecord) {
		fn(&Record{bamRecord: b})
	}

	return self.bamFetch(i.bamIndex, tid, beg, end, f)
}
<|endoftext|>