<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/russross/blackfriday"
)

var defaultLayout = []byte("{{Body}}")

func main() {
	var dest = ":8080"

	if len(os.Args) > 2 {
		log.Fatal("usage: md2html [dest]")
	}
	if len(os.Args) == 2 {
		dest = os.Args[1]
	}

	if !strings.Contains(dest, ":") {
		convert(dest)
		os.Exit(0)
	}

	serve(dest)
}

func serve(addr string) {
	fmt.Printf("serving at: http://localhost%s\n", addr)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		path := "." + r.URL.Path
		if filepath.Ext(path) != "" {
			http.ServeFile(w, r, path)
			return
		}
		b, err := render(path + ".md")
		if os.IsNotExist(err) {
			http.NotFound(w, r)
		} else if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		w.Write(b)
	})
	log.Fatal(http.ListenAndServe(addr, nil))
}

func convert(dest string) {
	fmt.Printf("Converting markdown to: %s\n", dest)
	err := filepath.Walk(".", func(path string, f os.FileInfo, err error) error {
		if f.IsDir() {
			return nil
		}
		match, err := filepath.Match("*.md", f.Name())
		printe(err)
		if !match {
			b, err := ioutil.ReadFile(path)
			if err != nil {
				return err
			}
			err = os.MkdirAll(filepath.Dir(filepath.Join(dest, path)), 0777)
			if err != nil {
				return err
			}
			return ioutil.WriteFile(filepath.Join(dest, path), b, 0644)
		}
		b, err := render(path)
		printe(err)
		path = dest + "/" + path
		printe(os.MkdirAll(filepath.Dir(path), 0777))
		printe(ioutil.WriteFile(strings.TrimSuffix(path, ".md"), b, 0644))
		fmt.Printf("converted: %s\n", path)
		return nil
	})
	printe(err)
}

func printe(err error) {
	if err != nil {
		fmt.Println(err)
	}
}

func render(f string) ([]byte, error) {
	src, err := ioutil.ReadFile(f)
	if err != nil {
		return nil, err
	}

	src = interpolateCode(src, path.Dir(f))

	layout, err := layout(f)
	if err != nil {
		return nil, err
	}

	return bytes.Replace(layout, defaultLayout, markdown(src), 1), nil
}

// Returns the contents of a layout.html file
// starting in the directory of p and ending at the command's
// working directory.
// If no layout.html file is found defaultLayout is returned.
func layout(p string) ([]byte, error) {
	// Don't search for layouts beyond the working dir
	wd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	for {
		p = path.Dir(p)
		l, err := ioutil.ReadFile(p + "/layout.html")
		if err == nil {
			return l, nil
		}
		if !os.IsNotExist(err) {
			return nil, err
		}

		if wd == p {
			break
		}
	}

	return defaultLayout, nil
}

func interpolateCode(md []byte, hostPath string) []byte {
	const pat = `$code `
	w := new(bytes.Buffer)
	scanner := bufio.NewScanner(bytes.NewBuffer(md))
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, pat) {
			var snippath, snippet string
			fields := strings.Fields(line)
			if len(fields) >= 2 {
				snippath = fields[1]
			}
			if len(fields) >= 3 {
				snippet = fields[2]
			}

			if path.IsAbs(snippath) {
				snippath = path.Join(os.Getenv("CHAIN"), snippath)
			} else {
				snippath = path.Join(hostPath, snippath)
			}

			writeCode(w, snippath, snippet)
			continue
		}
		fmt.Fprintln(w, line)
	}
	return w.Bytes()
}

func writeCode(w io.Writer, path, snippet string) {
	b, err := ioutil.ReadFile(path)
	if err != nil {
		fmt.Fprintf(w, "`Unable to read source file: %s`\n", path)
		return
	}
	//If the snippet is unset, print the whole file --omitting snippet definitions.
	if snippet == "" {
		fmt.Fprintln(w, "```")
		for _, line := range strings.Split(string(b), "\n") {
			if strings.Contains(line, "snippet") {
				continue
			}
			fmt.Fprintln(w, line)
		}
		fmt.Fprintln(w, "```")
		return
	}

	if !bytes.Contains(b, []byte("snippet "+snippet)) {
		fmt.Fprintf(w, "`Snippet %q is not in %q`\n", snippet, path)
		return
	}
	found := false
	for _, line := range strings.Split(string(b), "\n") {
		if strings.Contains(line, "snippet "+snippet) {
			fmt.Fprintln(w, "```")
			found = true
			continue
		}
		if found && strings.Contains(line, "endsnippet") {
			fmt.Fprintln(w, "```")
			found = false
			continue
		}
		if found {
			fmt.Fprintln(w, line)
		}
	}
}

func markdown(source []byte) []byte {
	htmlFlags := 0
	htmlFlags |= blackfriday.HTML_USE_XHTML
	htmlFlags |= blackfriday.HTML_USE_SMARTYPANTS
	htmlFlags |= blackfriday.HTML_SMARTYPANTS_FRACTIONS
	htmlFlags |= blackfriday.HTML_SMARTYPANTS_LATEX_DASHES
	htmlFlags |= blackfriday.HTML_SMARTYPANTS_DASHES
	htmlFlags |= blackfriday.HTML_FOOTNOTE_RETURN_LINKS
	renderer := blackfriday.HtmlRenderer(htmlFlags, "", "")

	extensions := 0
	extensions |= blackfriday.EXTENSION_FOOTNOTES
	extensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	extensions |= blackfriday.EXTENSION_TABLES
	extensions |= blackfriday.EXTENSION_FENCED_CODE
	extensions |= blackfriday.EXTENSION_AUTOLINK
	extensions |= blackfriday.EXTENSION_STRIKETHROUGH
	extensions |= blackfriday.EXTENSION_SPACE_HEADERS
	extensions |= blackfriday.EXTENSION_HEADER_IDS
	extensions |= blackfriday.EXTENSION_LAX_HTML_BLOCKS
	extensions |= blackfriday.EXTENSION_AUTO_HEADER_IDS

	return blackfriday.Markdown(source, renderer, extensions)
}
<commit_msg>cmd/md2html: only display relative indentation for interpolated code<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/russross/blackfriday"
)

var defaultLayout = []byte("{{Body}}")

func main() {
	var dest = ":8080"

	if len(os.Args) > 2 {
		log.Fatal("usage: md2html [dest]")
	}
	if len(os.Args) == 2 {
		dest = os.Args[1]
	}

	if !strings.Contains(dest, ":") {
		convert(dest)
		os.Exit(0)
	}

	serve(dest)
}

func serve(addr string) {
	fmt.Printf("serving at: http://localhost%s\n", addr)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		path := "." + r.URL.Path
		if filepath.Ext(path) != "" {
			http.ServeFile(w, r, path)
			return
		}
		b, err := render(path + ".md")
		if os.IsNotExist(err) {
			http.NotFound(w, r)
		} else if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		w.Write(b)
	})
	log.Fatal(http.ListenAndServe(addr, nil))
}

func convert(dest string) {
	fmt.Printf("Converting markdown to: %s\n", dest)
	err := filepath.Walk(".", func(path string, f os.FileInfo, err error) error {
		if f.IsDir() {
			return nil
		}
		match, err := filepath.Match("*.md", f.Name())
		printe(err)
		if !match {
			b, err := ioutil.ReadFile(path)
			if err != nil {
				return err
			}
			err = os.MkdirAll(filepath.Dir(filepath.Join(dest, path)), 0777)
			if err != nil {
				return err
			}
			return ioutil.WriteFile(filepath.Join(dest, path), b, 0644)
		}
		b, err := render(path)
		printe(err)
		path = dest + "/" + path
		printe(os.MkdirAll(filepath.Dir(path), 0777))
		printe(ioutil.WriteFile(strings.TrimSuffix(path, ".md"), b, 0644))
		fmt.Printf("converted: %s\n", path)
		return nil
	})
	printe(err)
}

func printe(err error) {
	if err != nil {
		fmt.Println(err)
	}
}

func render(f string) ([]byte, error) {
	src, err := ioutil.ReadFile(f)
	if err != nil {
		return nil, err
	}

	src = interpolateCode(src, path.Dir(f))

	layout, err := layout(f)
	if err != nil {
		return nil, err
	}

	return bytes.Replace(layout, defaultLayout, markdown(src), 1), nil
}

// Returns the contents of a layout.html file
// starting in the directory of p and ending at the command's
// working directory.
// If no layout.html file is found defaultLayout is returned.
func layout(p string) ([]byte, error) {
	// Don't search for layouts beyond the working dir
	wd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	for {
		p = path.Dir(p)
		l, err := ioutil.ReadFile(p + "/layout.html")
		if err == nil {
			return l, nil
		}
		if !os.IsNotExist(err) {
			return nil, err
		}

		if wd == p {
			break
		}
	}

	return defaultLayout, nil
}

func interpolateCode(md []byte, hostPath string) []byte {
	const pat = `$code `
	w := new(bytes.Buffer)
	scanner := bufio.NewScanner(bytes.NewBuffer(md))
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, pat) {
			var snippath, snippet string
			fields := strings.Fields(line)
			if len(fields) >= 2 {
				snippath = fields[1]
			}
			if len(fields) >= 3 {
				snippet = fields[2]
			}

			if path.IsAbs(snippath) {
				snippath = path.Join(os.Getenv("CHAIN"), snippath)
			} else {
				snippath = path.Join(hostPath, snippath)
			}

			writeCode(w, snippath, snippet)
			continue
		}
		fmt.Fprintln(w, line)
	}
	return w.Bytes()
}

func writeCode(w io.Writer, path, snippet string) {
	s, err := readSnippet(path, snippet)
	if err != nil {
		s = err.Error()
	} else {
		s = removeCommonIndent(s)
	}

	fmt.Fprintln(w, "```\n"+s+"\n```")
}

func readSnippet(path, snippet string) (string, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("Unable to read source file: %s\n", path)
	}

	src := string(b)

	// If the snippet is unset, return everything, omitting snippet definitions.
	if snippet == "" {
		var res string
		for _, line := range strings.SplitAfter(src, "\n") {
			if strings.Contains(line, "snippet") {
				continue
			}
			res += line
		}
		return res, nil
	}

	if !strings.Contains(src, "snippet "+snippet) {
		return "", fmt.Errorf("Snippet %q is not in %q", snippet, path)
	}

	var (
		res   []string
		found bool
	)

	for _, line := range strings.Split(src, "\n") {
		if !found && strings.Contains(line, "snippet "+snippet) {
			found = true
			continue
		}

		if found && strings.Contains(line, "endsnippet") {
			break
		}

		if found {
			res = append(res, line)
		}
	}

	return strings.Join(res, "\n"), nil
}

func removeCommonIndent(s string) string {
	lines := strings.Split(s, "\n")
	ci := commonIndent(lines)

	for i, line := range lines {
		lines[i] = strings.TrimPrefix(line, ci)
	}

	return strings.Join(lines, "\n")
}

// commonIndent returns the longest indentation shared by all non-empty lines
// in the input.
func commonIndent(lines []string) string {
	var (
		indent = regexp.MustCompile(`^[ \t]+`)
		res    string
		resSet bool
	)

	for _, line := range lines {
		// blank lines should not influence the common indentation
		if strings.TrimSpace(line) == "" {
			continue
		}

		if !resSet {
			res = indent.FindString(line)
			resSet = true
			continue
		}

		res = commonPrefix(res, indent.FindString(line))
		if res == "" {
			break
		}
	}

	return res
}

// commonPrefix returns the longest shared prefix between two strings. It
// assumes the characters can be represented by a single byte.
func commonPrefix(a, b string) string {
	var res []byte
	for i := range a {
		if i >= len(b) || a[i] != b[i] {
			break
		}
		res = append(res, a[i])
	}
	return string(res)
}

func markdown(source []byte) []byte {
	htmlFlags := 0
	htmlFlags |= blackfriday.HTML_USE_XHTML
	htmlFlags |= blackfriday.HTML_USE_SMARTYPANTS
	htmlFlags |= blackfriday.HTML_SMARTYPANTS_FRACTIONS
	htmlFlags |= blackfriday.HTML_SMARTYPANTS_LATEX_DASHES
	htmlFlags |= blackfriday.HTML_SMARTYPANTS_DASHES
	htmlFlags |= blackfriday.HTML_FOOTNOTE_RETURN_LINKS
	renderer := blackfriday.HtmlRenderer(htmlFlags, "", "")

	extensions := 0
	extensions |= blackfriday.EXTENSION_FOOTNOTES
	extensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	extensions |= blackfriday.EXTENSION_TABLES
	extensions |= blackfriday.EXTENSION_FENCED_CODE
	extensions |= blackfriday.EXTENSION_AUTOLINK
	extensions |= blackfriday.EXTENSION_STRIKETHROUGH
	extensions |= blackfriday.EXTENSION_SPACE_HEADERS
	extensions |= blackfriday.EXTENSION_HEADER_IDS
	extensions |= blackfriday.EXTENSION_LAX_HTML_BLOCKS
	extensions |= blackfriday.EXTENSION_AUTO_HEADER_IDS

	return blackfriday.Markdown(source, renderer, extensions)
}
<|endoftext|>


<commit_before>/*package memo is the worst code I have ever written in my life.*/
package memo

import (
	"encoding/binary"
	"fmt"
	"os"
	"path"
	"sort"

	"github.com/phil-mansfield/shellfish/cmd/env"

	"github.com/phil-mansfield/shellfish/cmd/halo"
	"github.com/phil-mansfield/shellfish/io"
)

const (
	rockstarMemoDir       = "rockstar"
	rockstarMemoFile      = "halo_%d.dat"
	rockstarShortMemoFile = "halo_short_%d.dat"
	rockstarShortMemoNum  = 10 * 1000

	headerMemoFile = "hd_snap%d.dat"
)

// ReadSortedRockstarIDs returns a slice of IDs corresponding to the highest
// values of some quantity in a particular snapshot. maxID is the number of
// halos to return.
func ReadSortedRockstarIDs(
	snap, maxID int, valName string, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment,
) ([]int, error) {
	hds, _, err := ReadHeaders(snap, buf, e)
	if err != nil {
		return nil, err
	}
	cosmo := &hds[0].Cosmo

	dir := path.Join(e.MemoDir, rockstarMemoDir)
	if _, err := os.Stat(dir); err != nil {
		err = os.Mkdir(dir, 0777)
		if err != nil {
			return nil, err
		}
	}

	var (
		ids []int
		vals [][]float64
		ms  []float64
	)

	if maxID >= rockstarShortMemoNum || maxID == -1 {
		file := path.Join(dir, fmt.Sprintf(rockstarMemoFile, snap))
		ids, vals, err = readRockstar(
			file, []string{valName}, -1, snap, nil, vars, buf, e, cosmo,
		)
		if err != nil {
			return nil, err
		}
		ms = vals[0]
	} else {
		file := path.Join(dir, fmt.Sprintf(rockstarShortMemoFile, snap))
		ids, vals, err = readRockstar(
			file, []string{valName}, rockstarShortMemoNum,
			snap, nil, vars, buf, e, cosmo,
		)
		if err != nil {
			return nil, err
		}
		ms = vals[0]
	}

	if len(ids) < maxID {
		return nil, fmt.Errorf(
			"ID %d too large for snapshot %d", maxID, snap,
		)
	}

	sortRockstar(ids, ms)
	if maxID == -1 {
		return ids, nil
	}
	return ids[:maxID+1], nil
}

type massSet struct {
	ids []int
	ms  []float64
}

func (set massSet) Len() int { return len(set.ids) }

// We're reverse sorting.
func (set massSet) Less(i, j int) bool { return set.ms[i] > set.ms[j] }
func (set massSet) Swap(i, j int) {
	set.ms[i], set.ms[j] = set.ms[j], set.ms[i]
	set.ids[i], set.ids[j] = set.ids[j], set.ids[i]
}

func sortRockstar(ids []int, ms []float64) {
	set := massSet{ids, ms}
	sort.Sort(set)
}

// This function does fairly large heap allocations even when it doesn't need
// to. Consider passing it a buffer.
func ReadRockstar(
	snap int, valNames []string, ids []int, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment,
) (outIDs []int, vals [][]float64, err error) {
	hds, _, err := ReadHeaders(snap, buf, e)
	if err != nil {
		return nil, nil, err
	}
	cosmo := &hds[0].Cosmo

	// Find binFile.
	dir := path.Join(e.MemoDir, rockstarMemoDir)
	if _, err := os.Stat(dir); err != nil {
		err = os.Mkdir(dir, 0777)
		if err != nil {
			return nil, nil, err
		}
	}

	binFile := path.Join(dir, fmt.Sprintf(rockstarMemoFile, snap))
	shortBinFile := path.Join(dir, fmt.Sprintf(rockstarShortMemoFile, snap))

	// This wastes a read the first time it's called. You need to decide if you
	// care. (Answer: probably not.)
	outIDs, vals, err = readRockstar(
		shortBinFile, valNames, rockstarShortMemoNum, snap,
		ids, vars, buf, e, cosmo,
	)
	// TODO: Fix error handling here.
	// (Confession: I have no idea what this comment means: I wrote it half a
	// year ago and didn't make an Issue about it.)
	if err == nil {
		return outIDs, vals, err
	}
	outIDs, vals, err = readRockstar(
		binFile, valNames, -1, snap, ids, vars, buf, e, cosmo,
	)

	if err != nil {
		return nil, nil, err
	}
	return outIDs, vals, nil
}

func readRockstar(
	binFile string, valNames []string, n, snap int, ids []int,
	vars *halo.VarColumns, buf io.VectorBuffer, e *env.Environment,
	cosmo *io.CosmologyHeader,
) (outIDs []int, vals [][]float64, err error) {
	hds, _, err := ReadHeaders(snap, buf, e)
	if err != nil {
		return nil, nil, err
	}
	hd := &hds[0]

	// If binFile doesn't exist, create it.
	if _, err := os.Stat(binFile); err != nil {
		if n == -1 {
			err = halo.RockstarConvert(
				e.HaloCatalog(snap), binFile, vars, &hd.Cosmo,
			)
			if err != nil {
				return nil, nil, err
			}
		} else {
			err = halo.RockstarConvertTopN(
				e.HaloCatalog(snap), binFile, n, vars, &hd.Cosmo,
			)
			if err != nil {
				return nil, nil, err
			}
		}
	}

	rids, rawCols, err := halo.ReadBinaryRockstar(binFile, vars)
	if err != nil {
		return nil, nil, err
	}

	rvals := make([][]float64, len(valNames))
	for i := range valNames {
		rvals[i] = vars.GetColumn(rawCols, valNames[i], cosmo)
	}

	// Select out only the IDs we want.
	if ids == nil {
		return rids, rvals, nil
	}
	vals = make([][]float64, len(rvals))

	for i := range vals {
		vals[i] = make([]float64, len(ids))
	}

	f := NewIntFinder(rids)
	for i, id := range ids {
		line, ok := f.Find(id)
		if !ok {
			err = fmt.Errorf("Could not find ID %d", id)
			return nil, nil, err
		}
		for vi := range vals {
			vals[vi][i] = rvals[vi][line]
		}
	}

	return ids, vals, nil
}

// A quick generic wrapper for doing those one-to-one mappings I need to do so
// often. Written like this so the backend can be swapped out easily.
type IntFinder struct {
	m map[int]int
}

// NewIntFinder creates a new IntFinder struct for a given slice of Rockstar
// IDs.
func NewIntFinder(rids []int) IntFinder {
	f := IntFinder{}
	f.m = make(map[int]int)
	for i, rid := range rids {
		f.m[rid] = i
	}
	return f
}

// Find returns the index which the given ID corresponds to and true if the
// ID is in the finder. Otherwise, false is returned.
func (f IntFinder) Find(rid int) (int, bool) {
	line, ok := f.m[rid]
	return line, ok
}

func readUnmemoizedHeaders(
	snap int, buf io.VectorBuffer, e *env.Environment,
) ([]io.Header, []string, error) {
	fmt.Println("Buffer:", buf)
	fmt.Println("BlockNum:", e.Blocks())
	files := make([]string, e.Blocks())
	hds := make([]io.Header, e.Blocks())

	for i := range files {
		fmt.Println("Block:", i)
		files[i] = e.ParticleCatalog(snap, i)
		fmt.Println("File:", files[i])
		err := buf.ReadHeader(files[i], &hds[i])
		if err != nil {
			return nil, nil, err
		}
	}
	return hds, files, nil
}

// ReadHeaders returns all the segment headers and segment file names for all
// the segments at a given snapshot.
func ReadHeaders(
	snap int, buf io.VectorBuffer, e *env.Environment,
) ([]io.Header, []string, error) {
	if _, err := os.Stat(e.MemoDir); err != nil {
		return nil, nil, err
	}
	memoFile := path.Join(e.MemoDir, fmt.Sprintf(headerMemoFile, snap))

	if _, err := os.Stat(memoFile); err != nil {
		// File not written yet.
		hds, files, err := readUnmemoizedHeaders(snap, buf, e)
		if err != nil {
			return nil, nil, err
		}

		f, err := os.Create(memoFile)
		if err != nil {
			return nil, nil, err
		}
		defer f.Close()

		binary.Write(f, binary.LittleEndian, hds)

		return hds, files, nil
	} else {
		// File exists: read from it instead.

		f, err := os.Open(memoFile)
		if err != nil {
			return nil, nil, err
		}
		defer f.Close()

		hds := make([]io.Header, e.Blocks())
		binary.Read(f, binary.LittleEndian, hds)
		files := make([]string, e.Blocks())
		for i := range files {
			files[i] = e.ParticleCatalog(snap, i)
		}

		return hds, files, nil
	}
}
<commit_msg>Removed old logging statements.<commit_after>/*package memo is the worst code I have ever written in my life.*/
package memo

import (
	"encoding/binary"
	"fmt"
	"os"
	"path"
	"sort"

	"github.com/phil-mansfield/shellfish/cmd/env"

	"github.com/phil-mansfield/shellfish/cmd/halo"
	"github.com/phil-mansfield/shellfish/io"
)

const (
	rockstarMemoDir       = "rockstar"
	rockstarMemoFile      = "halo_%d.dat"
	rockstarShortMemoFile = "halo_short_%d.dat"
	rockstarShortMemoNum  = 10 * 1000

	headerMemoFile = "hd_snap%d.dat"
)

// ReadSortedRockstarIDs returns a slice of IDs corresponding to the highest
// values of some quantity in a particular snapshot. maxID is the number of
// halos to return.
func ReadSortedRockstarIDs(
	snap, maxID int, valName string, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment,
) ([]int, error) {
	hds, _, err := ReadHeaders(snap, buf, e)
	if err != nil {
		return nil, err
	}
	cosmo := &hds[0].Cosmo

	dir := path.Join(e.MemoDir, rockstarMemoDir)
	if _, err := os.Stat(dir); err != nil {
		err = os.Mkdir(dir, 0777)
		if err != nil {
			return nil, err
		}
	}

	var (
		ids []int
		vals [][]float64
		ms  []float64
	)

	if maxID >= rockstarShortMemoNum || maxID == -1 {
		file := path.Join(dir, fmt.Sprintf(rockstarMemoFile, snap))
		ids, vals, err = readRockstar(
			file, []string{valName}, -1, snap, nil, vars, buf, e, cosmo,
		)
		if err != nil {
			return nil, err
		}
		ms = vals[0]
	} else {
		file := path.Join(dir, fmt.Sprintf(rockstarShortMemoFile, snap))
		ids, vals, err = readRockstar(
			file, []string{valName}, rockstarShortMemoNum,
			snap, nil, vars, buf, e, cosmo,
		)
		if err != nil {
			return nil, err
		}
		ms = vals[0]
	}

	if len(ids) < maxID {
		return nil, fmt.Errorf(
			"ID %d too large for snapshot %d", maxID, snap,
		)
	}

	sortRockstar(ids, ms)
	if maxID == -1 {
		return ids, nil
	}
	return ids[:maxID+1], nil
}

type massSet struct {
	ids []int
	ms  []float64
}

func (set massSet) Len() int { return len(set.ids) }

// We're reverse sorting.
func (set massSet) Less(i, j int) bool { return set.ms[i] > set.ms[j] }
func (set massSet) Swap(i, j int) {
	set.ms[i], set.ms[j] = set.ms[j], set.ms[i]
	set.ids[i], set.ids[j] = set.ids[j], set.ids[i]
}

func sortRockstar(ids []int, ms []float64) {
	set := massSet{ids, ms}
	sort.Sort(set)
}

// This function does fairly large heap allocations even when it doesn't need
// to. Consider passing it a buffer.
func ReadRockstar(
	snap int, valNames []string, ids []int, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment,
) (outIDs []int, vals [][]float64, err error) {
	hds, _, err := ReadHeaders(snap, buf, e)
	if err != nil {
		return nil, nil, err
	}
	cosmo := &hds[0].Cosmo

	// Find binFile.
	dir := path.Join(e.MemoDir, rockstarMemoDir)
	if _, err := os.Stat(dir); err != nil {
		err = os.Mkdir(dir, 0777)
		if err != nil {
			return nil, nil, err
		}
	}

	binFile := path.Join(dir, fmt.Sprintf(rockstarMemoFile, snap))
	shortBinFile := path.Join(dir, fmt.Sprintf(rockstarShortMemoFile, snap))

	// This wastes a read the first time it's called. You need to decide if you
	// care. (Answer: probably not.)
	outIDs, vals, err = readRockstar(
		shortBinFile, valNames, rockstarShortMemoNum, snap,
		ids, vars, buf, e, cosmo,
	)
	// TODO: Fix error handling here.
	// (Confession: I have no idea what this comment means: I wrote it half a
	// year ago and didn't make an Issue about it.)
	if err == nil {
		return outIDs, vals, err
	}
	outIDs, vals, err = readRockstar(
		binFile, valNames, -1, snap, ids, vars, buf, e, cosmo,
	)

	if err != nil {
		return nil, nil, err
	}
	return outIDs, vals, nil
}

func readRockstar(
	binFile string, valNames []string, n, snap int, ids []int,
	vars *halo.VarColumns, buf io.VectorBuffer, e *env.Environment,
	cosmo *io.CosmologyHeader,
) (outIDs []int, vals [][]float64, err error) {
	hds, _, err := ReadHeaders(snap, buf, e)
	if err != nil {
		return nil, nil, err
	}
	hd := &hds[0]

	// If binFile doesn't exist, create it.
	if _, err := os.Stat(binFile); err != nil {
		if n == -1 {
			err = halo.RockstarConvert(
				e.HaloCatalog(snap), binFile, vars, &hd.Cosmo,
			)
			if err != nil {
				return nil, nil, err
			}
		} else {
			err = halo.RockstarConvertTopN(
				e.HaloCatalog(snap), binFile, n, vars, &hd.Cosmo,
			)
			if err != nil {
				return nil, nil, err
			}
		}
	}

	rids, rawCols, err := halo.ReadBinaryRockstar(binFile, vars)
	if err != nil {
		return nil, nil, err
	}

	rvals := make([][]float64, len(valNames))
	for i := range valNames {
		rvals[i] = vars.GetColumn(rawCols, valNames[i], cosmo)
	}

	// Select out only the IDs we want.
	if ids == nil {
		return rids, rvals, nil
	}
	vals = make([][]float64, len(rvals))

	for i := range vals {
		vals[i] = make([]float64, len(ids))
	}

	f := NewIntFinder(rids)
	for i, id := range ids {
		line, ok := f.Find(id)
		if !ok {
			err = fmt.Errorf("Could not find ID %d", id)
			return nil, nil, err
		}
		for vi := range vals {
			vals[vi][i] = rvals[vi][line]
		}
	}

	return ids, vals, nil
}

// A quick generic wrapper for doing those one-to-one mappings I need to do so
// often. Written like this so the backend can be swapped out easily.
type IntFinder struct {
	m map[int]int
}

// NewIntFinder creates a new IntFinder struct for a given slice of Rockstar
// IDs.
func NewIntFinder(rids []int) IntFinder {
	f := IntFinder{}
	f.m = make(map[int]int)
	for i, rid := range rids {
		f.m[rid] = i
	}
	return f
}

// Find returns the index which the given ID corresponds to and true if the
// ID is in the finder. Otherwise, false is returned.
func (f IntFinder) Find(rid int) (int, bool) {
	line, ok := f.m[rid]
	return line, ok
}

func readUnmemoizedHeaders(
	snap int, buf io.VectorBuffer, e *env.Environment,
) ([]io.Header, []string, error) {
	files := make([]string, e.Blocks())
	hds := make([]io.Header, e.Blocks())

	for i := range files {
		files[i] = e.ParticleCatalog(snap, i)
		err := buf.ReadHeader(files[i], &hds[i])
		if err != nil {
			return nil, nil, err
		}
	}
	return hds, files, nil
}

// ReadHeaders returns all the segment headers and segment file names for all
// the segments at a given snapshot.
func ReadHeaders(
	snap int, buf io.VectorBuffer, e *env.Environment,
) ([]io.Header, []string, error) {
	if _, err := os.Stat(e.MemoDir); err != nil {
		return nil, nil, err
	}
	memoFile := path.Join(e.MemoDir, fmt.Sprintf(headerMemoFile, snap))

	if _, err := os.Stat(memoFile); err != nil {
		// File not written yet.
		hds, files, err := readUnmemoizedHeaders(snap, buf, e)
		if err != nil {
			return nil, nil, err
		}

		f, err := os.Create(memoFile)
		if err != nil {
			return nil, nil, err
		}
		defer f.Close()

		binary.Write(f, binary.LittleEndian, hds)

		return hds, files, nil
	} else {
		// File exists: read from it instead.

		f, err := os.Open(memoFile)
		if err != nil {
			return nil, nil, err
		}
		defer f.Close()

		hds := make([]io.Header, e.Blocks())
		binary.Read(f, binary.LittleEndian, hds)
		files := make([]string, e.Blocks())
		for i := range files {
			files[i] = e.ParticleCatalog(snap, i)
		}

		return hds, files, nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/nightlyone/lockfile"
	"testing"
)

func TestLocking(t *testing.T) {
	monitoringEvent = "TestLocking"
	lf, err := createLock(false)
	if err != nil {
		t.Fatal("environment problem: ", err)
		return
	}
	t.Log("Got lockfile")
	if lf2, err := createLock(false); err == nil {
		t.Error("got lockfile, but expected %v", lockfile.ErrBusy)
		lf2.Unlock()
	} else if err != lockfile.ErrBusy {
		t.Error("bad error got %v, want %v", err, lockfile.ErrBusy)
	} else {
		t.Log("got expected %v", err)
	}
	lf.Unlock()

	lf3, err := createLock(false)
	if err != nil {
		t.Error("unexpected error %v", err)
	} else {
		t.Log("got lockfile", err)
		lf3.Unlock()
	}
}
<commit_msg>Fix lock test to use formatting<commit_after>package main

import (
	"github.com/nightlyone/lockfile"
	"testing"
)

func TestLocking(t *testing.T) {
	monitoringEvent = "TestLocking"
	lf, err := createLock(false)
	if err != nil {
		t.Fatal("environment problem: ", err)
		return
	}
	t.Log("Got lockfile")
	if lf2, err := createLock(false); err == nil {
		t.Errorf("got lockfile, but expected '%v'", lockfile.ErrBusy)
		lf2.Unlock()
	} else if err != lockfile.ErrBusy {
		t.Errorf("bad error got '%v', want '%v'", err, lockfile.ErrBusy)
	} else {
		t.Logf("got expected '%v'", err)
	}
	lf.Unlock()

	lf3, err := createLock(false)
	if err != nil {
		t.Errorf("unexpected error %v", err)
	} else {
		t.Log("got lockfile", lf3)
		lf3.Unlock()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"

	"path/filepath"

	executorinit "code.cloudfoundry.org/executor/initializer"
)

var gardenNetwork = flag.String(
	"gardenNetwork",
	executorinit.DefaultConfiguration.GardenNetwork,
	"network mode for garden server (tcp, unix)",
)

var gardenAddr = flag.String(
	"gardenAddr",
	executorinit.DefaultConfiguration.GardenAddr,
	"network address for garden server",
)

var memoryMBFlag = flag.String(
	"memoryMB",
	executorinit.DefaultConfiguration.MemoryMB,
	"the amount of memory the executor has available in megabytes",
)

var diskMBFlag = flag.String(
	"diskMB",
	executorinit.DefaultConfiguration.DiskMB,
	"the amount of disk the executor has available in megabytes",
)

var tempDir = flag.String(
	"tempDir",
	executorinit.DefaultConfiguration.TempDir,
	"location to store temporary assets",
)

var reservedExpirationTime = flag.Duration(
	"reservedExpirationTime",
	executorinit.DefaultConfiguration.ReservedExpirationTime,
	"amount of time during which a container can remain in the allocated state",
)

var containerReapInterval = flag.Duration(
	"containerReapInterval",
	executorinit.DefaultConfiguration.ContainerReapInterval,
	"interval at which the executor reaps extra/missing containers",
)

var containerInodeLimit = flag.Uint64(
	"containerInodeLimit",
	executorinit.DefaultConfiguration.ContainerInodeLimit,
	"max number of inodes per container",
)

var containerMaxCpuShares = flag.Uint64(
	"containerMaxCpuShares",
	executorinit.DefaultConfiguration.ContainerMaxCpuShares,
	"cpu shares allocatable to a container",
)

var cachePath = flag.String(
	"cachePath",
	executorinit.DefaultConfiguration.CachePath,
	"location to cache assets",
)

var maxCacheSizeInBytes = flag.Uint64(
	"maxCacheSizeInBytes",
	executorinit.DefaultConfiguration.MaxCacheSizeInBytes,
	"maximum size of the cache (in bytes) - you should include a healthy amount of overhead",
)

var skipCertVerify = flag.Bool(
	"skipCertVerify",
	executorinit.DefaultConfiguration.SkipCertVerify,
	"skip SSL certificate verification",
)

var pathToCACertsForDownloads = flag.String(
	"caCertsForDownloads",
	"",
	"path to CA certificate bundle to be used when downloading assets",
)

var healthyMonitoringInterval = flag.Duration(
	"healthyMonitoringInterval",
	executorinit.DefaultConfiguration.HealthyMonitoringInterval,
	"interval on which to check healthy containers",
)

var unhealthyMonitoringInterval = flag.Duration(
	"unhealthyMonitoringInterval",
	executorinit.DefaultConfiguration.UnhealthyMonitoringInterval,
	"interval on which to check unhealthy containers",
)

var exportNetworkEnvVars = flag.Bool(
	"exportNetworkEnvVars",
	executorinit.DefaultConfiguration.ExportNetworkEnvVars,
	"export network environment variables into container (e.g. CF_INSTANCE_IP, CF_INSTANCE_PORT)",
)

var containerOwnerName = flag.String(
	"containerOwnerName",
	executorinit.DefaultConfiguration.ContainerOwnerName,
	"owner name with which to tag containers",
)

var createWorkPoolSize = flag.Int(
	"createWorkPoolSize",
	executorinit.DefaultConfiguration.CreateWorkPoolSize,
	"Number of concurrent create operations in garden",
)

var deleteWorkPoolSize = flag.Int(
	"deleteWorkPoolSize",
	executorinit.DefaultConfiguration.DeleteWorkPoolSize,
	"Number of concurrent delete operations in garden",
)

var readWorkPoolSize = flag.Int(
	"readWorkPoolSize",
	executorinit.DefaultConfiguration.ReadWorkPoolSize,
	"Number of concurrent read operations in garden",
)

var metricsWorkPoolSize = flag.Int(
	"metricsWorkPoolSize",
	executorinit.DefaultConfiguration.MetricsWorkPoolSize,
	"Number of concurrent metrics operations in garden",
)

var healthCheckWorkPoolSize = flag.Int(
	"healthCheckWorkPoolSize",
	executorinit.DefaultConfiguration.HealthCheckWorkPoolSize,
	"Number of concurrent ping operations in garden",
)

var maxConcurrentDownloads = flag.Int(
	"maxConcurrentDownloads",
	executorinit.DefaultConfiguration.MaxConcurrentDownloads,
	"Number of concurrent download steps",
)

var gardenHealthcheckInterval = flag.Duration(
	"gardenHealthcheckInterval",
	executorinit.DefaultConfiguration.GardenHealthcheckInterval,
	"Frequency for healthchecking garden",
)

var gardenHealthcheckEmissionInterval = flag.Duration(
	"gardenHealthcheckEmissionInterval",
	executorinit.DefaultConfiguration.GardenHealthcheckEmissionInterval,
	"Frequency for emitting UnhealthyCell metric",
)

var gardenHealthcheckTimeout = flag.Duration(
	"gardenHealthcheckTimeout",
	executorinit.DefaultConfiguration.GardenHealthcheckTimeout,
	"Maximum allowed time for garden healthcheck",
)

var gardenHealthcheckCommandRetryPause = flag.Duration(
	"gardenHealthcheckCommandRetryPause",
	executorinit.DefaultConfiguration.GardenHealthcheckCommandRetryPause,
	"Time to wait between retrying garden commands",
)

var gardenHealthcheckProcessPath = flag.String(
	"gardenHealthcheckProcessPath",
	executorinit.DefaultConfiguration.GardenHealthcheckProcessPath,
	"Path of the command to run to perform a container healthcheck",
)

var gardenHealthcheckProcessUser = flag.String(
	"gardenHealthcheckProcessUser",
	executorinit.DefaultConfiguration.GardenHealthcheckProcessUser,
	"User to use while performing a container healthcheck",
)

var gardenHealthcheckProcessDir = flag.String(
	"gardenHealthcheckProcessDir",
	executorinit.DefaultConfiguration.GardenHealthcheckProcessDir,
	"Directory to run the healthcheck process from",
)

var postSetupHook = flag.String(
	"postSetupHook",
	"",
	"Set of commands to run after the setup action",
)

var postSetupUser = flag.String(
	"postSetupUser",
	"root",
	"User to run the post setup hook",
)

var trustedSystemCertificatesPath = flag.String(
	"trustedSystemCertificatesPath",
	"",
	"path to directory containing trusted system ca certs.",
)

var volmanDriverPaths = flag.String(
	"volmanDriverPaths",
	"",
	"path to directories that the volume manager uses to discover drivers.  May be an OS-specific path-separated set of paths; e.g. /path/to/a:/path/to/b",
)

func executorConfig(caCertsForDownloads []byte, gardenHealthcheckRootFS string, gardenHealthcheckArgs, gardenHealthcheckEnv []string) executorinit.Configuration {
	return executorinit.Configuration{
		GardenNetwork:                      *gardenNetwork,
		GardenAddr:                         *gardenAddr,
		ContainerOwnerName:                 *containerOwnerName,
		TempDir:                            *tempDir,
		CachePath:                          *cachePath,
		MaxCacheSizeInBytes:                *maxCacheSizeInBytes,
		SkipCertVerify:                     *skipCertVerify,
		CACertsForDownloads:                caCertsForDownloads,
		ExportNetworkEnvVars:               *exportNetworkEnvVars,
		ContainerMaxCpuShares:              *containerMaxCpuShares,
		ContainerInodeLimit:                *containerInodeLimit,
		HealthyMonitoringInterval:          *healthyMonitoringInterval,
		UnhealthyMonitoringInterval:        *unhealthyMonitoringInterval,
		HealthCheckWorkPoolSize:            *healthCheckWorkPoolSize,
		CreateWorkPoolSize:                 *createWorkPoolSize,
		DeleteWorkPoolSize:                 *deleteWorkPoolSize,
		ReadWorkPoolSize:                   *readWorkPoolSize,
		MetricsWorkPoolSize:                *metricsWorkPoolSize,
		ReservedExpirationTime:             *reservedExpirationTime,
		ContainerReapInterval:              *containerReapInterval,
		MemoryMB:                           *memoryMBFlag,
		DiskMB:                             *diskMBFlag,
		MaxConcurrentDownloads:             *maxConcurrentDownloads,
		GardenHealthcheckInterval:          *gardenHealthcheckInterval,
		GardenHealthcheckEmissionInterval:  *gardenHealthcheckEmissionInterval,
		GardenHealthcheckTimeout:           *gardenHealthcheckTimeout,
		GardenHealthcheckCommandRetryPause: *gardenHealthcheckCommandRetryPause,
		GardenHealthcheckRootFS:            gardenHealthcheckRootFS,
		GardenHealthcheckProcessPath:       *gardenHealthcheckProcessPath,
		GardenHealthcheckProcessUser:       *gardenHealthcheckProcessUser,
		GardenHealthcheckProcessDir:        *gardenHealthcheckProcessDir,
		GardenHealthcheckProcessArgs:       gardenHealthcheckArgs,
		GardenHealthcheckProcessEnv:        gardenHealthcheckEnv,
		PostSetupHook:                      *postSetupHook,
		PostSetupUser:                      *postSetupUser,
		TrustedSystemCertificatesPath:      *trustedSystemCertificatesPath,
		VolmanDriverPaths:                  filepath.SplitList(*volmanDriverPaths),
	}
}
<commit_msg>Add parameter for container metric reporting interval<commit_after>package main

import (
	"flag"

	"path/filepath"

	executorinit "code.cloudfoundry.org/executor/initializer"
)

var gardenNetwork = flag.String(
	"gardenNetwork",
	executorinit.DefaultConfiguration.GardenNetwork,
	"network mode for garden server (tcp, unix)",
)

var gardenAddr = flag.String(
	"gardenAddr",
	executorinit.DefaultConfiguration.GardenAddr,
	"network address for garden server",
)

var memoryMBFlag = flag.String(
	"memoryMB",
	executorinit.DefaultConfiguration.MemoryMB,
	"the amount of memory the executor has available in megabytes",
)

var diskMBFlag = flag.String(
	"diskMB",
	executorinit.DefaultConfiguration.DiskMB,
	"the amount of disk the executor has available in megabytes",
)

var tempDir = flag.String(
	"tempDir",
	executorinit.DefaultConfiguration.TempDir,
	"location to store temporary assets",
)

var reservedExpirationTime = flag.Duration(
	"reservedExpirationTime",
	executorinit.DefaultConfiguration.ReservedExpirationTime,
	"amount of time during which a container can remain in the allocated state",
)

var containerReapInterval = flag.Duration(
	"containerReapInterval",
	executorinit.DefaultConfiguration.ContainerReapInterval,
	"interval at which the executor reaps extra/missing containers",
)

var containerMetricsReportInterval = flag.Duration(
	"containerMetricsReportInterval",
	executorinit.DefaultConfiguration.ContainerMetricsReportInterval,
	"interval at which the executor reports metrics from containers",
)

var containerInodeLimit = flag.Uint64(
	"containerInodeLimit",
	executorinit.DefaultConfiguration.ContainerInodeLimit,
	"max number of inodes per container",
)

var containerMaxCpuShares = flag.Uint64(
	"containerMaxCpuShares",
	executorinit.DefaultConfiguration.ContainerMaxCpuShares,
	"cpu shares allocatable to a container",
)

var cachePath = flag.String(
	"cachePath",
	executorinit.DefaultConfiguration.CachePath,
	"location to cache assets",
)

var maxCacheSizeInBytes = flag.Uint64(
	"maxCacheSizeInBytes",
	executorinit.DefaultConfiguration.MaxCacheSizeInBytes,
	"maximum size of the cache (in bytes) - you should include a healthy amount of overhead",
)

var skipCertVerify = flag.Bool(
	"skipCertVerify",
	executorinit.DefaultConfiguration.SkipCertVerify,
	"skip SSL certificate verification",
)

var pathToCACertsForDownloads = flag.String(
	"caCertsForDownloads",
	"",
	"path to CA certificate bundle to be used when downloading assets",
)

var healthyMonitoringInterval = flag.Duration(
	"healthyMonitoringInterval",
	executorinit.DefaultConfiguration.HealthyMonitoringInterval,
	"interval on which to check healthy containers",
)

var unhealthyMonitoringInterval = flag.Duration(
	"unhealthyMonitoringInterval",
	executorinit.DefaultConfiguration.UnhealthyMonitoringInterval,
	"interval on which to check unhealthy containers",
)

var exportNetworkEnvVars = flag.Bool(
	"exportNetworkEnvVars",
	executorinit.DefaultConfiguration.ExportNetworkEnvVars,
	"export network environment variables into container (e.g. CF_INSTANCE_IP, CF_INSTANCE_PORT)",
)

var containerOwnerName = flag.String(
	"containerOwnerName",
	executorinit.DefaultConfiguration.ContainerOwnerName,
	"owner name with which to tag containers",
)

var createWorkPoolSize = flag.Int(
	"createWorkPoolSize",
	executorinit.DefaultConfiguration.CreateWorkPoolSize,
	"Number of concurrent create operations in garden",
)

var deleteWorkPoolSize = flag.Int(
	"deleteWorkPoolSize",
	executorinit.DefaultConfiguration.DeleteWorkPoolSize,
	"Number of concurrent delete operations in garden",
)

var readWorkPoolSize = flag.Int(
	"readWorkPoolSize",
	executorinit.DefaultConfiguration.ReadWorkPoolSize,
	"Number of concurrent read operations in garden",
)

var metricsWorkPoolSize = flag.Int(
	"metricsWorkPoolSize",
	executorinit.DefaultConfiguration.MetricsWorkPoolSize,
	"Number of concurrent metrics operations in garden",
)

var healthCheckWorkPoolSize = flag.Int(
	"healthCheckWorkPoolSize",
	executorinit.DefaultConfiguration.HealthCheckWorkPoolSize,
	"Number of concurrent ping operations in garden",
)

var maxConcurrentDownloads = flag.Int(
	"maxConcurrentDownloads",
	executorinit.DefaultConfiguration.MaxConcurrentDownloads,
	"Number of concurrent download steps",
)

var gardenHealthcheckInterval = flag.Duration(
	"gardenHealthcheckInterval",
	executorinit.DefaultConfiguration.GardenHealthcheckInterval,
	"Frequency for healthchecking garden",
)

var gardenHealthcheckEmissionInterval = flag.Duration(
	"gardenHealthcheckEmissionInterval",
	executorinit.DefaultConfiguration.GardenHealthcheckEmissionInterval,
	"Frequency for emitting UnhealthyCell metric",
)

var gardenHealthcheckTimeout = flag.Duration(
	"gardenHealthcheckTimeout",
	executorinit.DefaultConfiguration.GardenHealthcheckTimeout,
	"Maximum allowed time for garden healthcheck",
)

var gardenHealthcheckCommandRetryPause = flag.Duration(
	"gardenHealthcheckCommandRetryPause",
	executorinit.DefaultConfiguration.GardenHealthcheckCommandRetryPause,
	"Time to wait between retrying garden commands",
)

var gardenHealthcheckProcessPath = flag.String(
	"gardenHealthcheckProcessPath",
	executorinit.DefaultConfiguration.GardenHealthcheckProcessPath,
	"Path of the command to run to perform a container healthcheck",
)

var gardenHealthcheckProcessUser = flag.String(
	"gardenHealthcheckProcessUser",
	executorinit.DefaultConfiguration.GardenHealthcheckProcessUser,
	"User to use while performing a container healthcheck",
)

var gardenHealthcheckProcessDir = flag.String(
	"gardenHealthcheckProcessDir",
	executorinit.DefaultConfiguration.GardenHealthcheckProcessDir,
	"Directory to run the healthcheck process from",
)

var postSetupHook = flag.String(
	"postSetupHook",
	"",
	"Set of commands to run after the setup action",
)

var postSetupUser = flag.String(
	"postSetupUser",
	"root",
	"User to run the post setup hook",
)

var trustedSystemCertificatesPath = flag.String(
	"trustedSystemCertificatesPath",
	"",
	"path to directory containing trusted system ca certs.",
)

var volmanDriverPaths = flag.String(
	"volmanDriverPaths",
	"",
	"path to directories that the volume manager uses to discover drivers.  May be an OS-specific path-separated set of paths; e.g. /path/to/a:/path/to/b",
)

func executorConfig(caCertsForDownloads []byte, gardenHealthcheckRootFS string, gardenHealthcheckArgs, gardenHealthcheckEnv []string) executorinit.Configuration {
	return executorinit.Configuration{
		GardenNetwork:                      *gardenNetwork,
		GardenAddr:                         *gardenAddr,
		ContainerOwnerName:                 *containerOwnerName,
		TempDir:                            *tempDir,
		CachePath:                          *cachePath,
		MaxCacheSizeInBytes:                *maxCacheSizeInBytes,
		SkipCertVerify:                     *skipCertVerify,
		CACertsForDownloads:                caCertsForDownloads,
		ExportNetworkEnvVars:               *exportNetworkEnvVars,
		ContainerMaxCpuShares:              *containerMaxCpuShares,
		ContainerInodeLimit:                *containerInodeLimit,
		HealthyMonitoringInterval:          *healthyMonitoringInterval,
		UnhealthyMonitoringInterval:        *unhealthyMonitoringInterval,
		HealthCheckWorkPoolSize:            *healthCheckWorkPoolSize,
		CreateWorkPoolSize:                 *createWorkPoolSize,
		DeleteWorkPoolSize:                 *deleteWorkPoolSize,
		ReadWorkPoolSize:                   *readWorkPoolSize,
		MetricsWorkPoolSize:                *metricsWorkPoolSize,
		ReservedExpirationTime:             *reservedExpirationTime,
		ContainerReapInterval:              *containerReapInterval,
		MemoryMB:                           *memoryMBFlag,
		DiskMB:                             *diskMBFlag,
		MaxConcurrentDownloads:             *maxConcurrentDownloads,
		GardenHealthcheckInterval:          *gardenHealthcheckInterval,
		GardenHealthcheckEmissionInterval:  *gardenHealthcheckEmissionInterval,
		GardenHealthcheckTimeout:           *gardenHealthcheckTimeout,
		GardenHealthcheckCommandRetryPause: *gardenHealthcheckCommandRetryPause,
		GardenHealthcheckRootFS:            gardenHealthcheckRootFS,
		GardenHealthcheckProcessPath:       *gardenHealthcheckProcessPath,
		GardenHealthcheckProcessUser:       *gardenHealthcheckProcessUser,
		GardenHealthcheckProcessDir:        *gardenHealthcheckProcessDir,
		GardenHealthcheckProcessArgs:       gardenHealthcheckArgs,
		GardenHealthcheckProcessEnv:        gardenHealthcheckEnv,
		PostSetupHook:                      *postSetupHook,
		PostSetupUser:                      *postSetupUser,
		TrustedSystemCertificatesPath:      *trustedSystemCertificatesPath,
		VolmanDriverPaths:                  filepath.SplitList(*volmanDriverPaths),
		ContainerMetricsReportInterval:     *containerMetricsReportInterval,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"gopkg.in/urfave/cli.v1"
	"log"
	"os"
)

func main() {
	errLog := log.New(os.Stderr, "", 0)
	app := cli.NewApp()
	app.Name = "The SciPipe Tool"
	app.Usage = "A helper tool to ease working with SciPipe workflows"
	app.Version = "0.0.1"
	app.Commands = []cli.Command{
		{
			Name:    "new",
			Aliases: []string{"n"},
			Usage:   "Create a new example workflow file, in a filename provided as first argument after 'new' (default: mynewworkflow.go).",
			Action: func(c *cli.Context) error {
				wfcode := `// Workflow written in SciPipe.
// For more information about SciPipe, see: http://scipipe.org
package main

import sp "github.com/scipipe/scipipe"

func main() {
	// --------------------------------
	// Create a pipeline runner
	// --------------------------------

	run := sp.NewPipelineRunner()

	// --------------------------------
	// Initialize processes and add to runner
	// --------------------------------

	foo := sp.NewProc("fooer",
		"echo foo > {o:foo}")
	foo.SetPathStatic("foo", "foo.txt")
	run.AddProcess(foo)

	f2b := sp.NewProc("foo2bar",
		"sed 's/foo/bar/g' {i:foo} > {o:bar}")
	f2b.SetPathExtend("foo", "bar", ".bar.txt")
	run.AddProcess(f2b)

	snk := sp.NewSink()
	run.AddProcess(snk)

	// --------------------------------
	// Connect workflow dependency network
	// --------------------------------

	f2b.In["foo"].Connect(foo.Out["foo"])
	snk.Connect(f2b.Out["bar"])

	// --------------------------------
	// Run the pipeline!
	// --------------------------------

	run.Run()
}`
				fileName := c.Args().First()
				if fileName == "" {
					fileName = "new_scipipe_workflow.go"
					fmt.Printf("No filename specified, so using the default '%s' ...\n", fileName)
				}

				f, err := os.Create(fileName)
				if err != nil {
					errLog.Println("Could not create file:", fileName)
					os.Exit(1)
				}
				defer f.Close()

				_, err = f.Write([]byte(wfcode))
				if err != nil {
					errLog.Println("Could not write to file:", fileName)
					os.Exit(1)
				}

				fmt.Println("Successfully wrote new workflow file to:", fileName)
				return nil
			},
		},
	}
	app.Run(os.Args)
}
<commit_msg>Use github URL for urfave/cli<commit_after>package main

import (
	"fmt"
	"github.com/urfave/cli"
	"log"
	"os"
)

func main() {
	errLog := log.New(os.Stderr, "", 0)
	app := cli.NewApp()
	app.Name = "The SciPipe Tool"
	app.Usage = "A helper tool to ease working with SciPipe workflows"
	app.Version = "0.0.1"
	app.Commands = []cli.Command{
		{
			Name:    "new",
			Aliases: []string{"n"},
			Usage:   "Create a new example workflow file, in a filename provided as first argument after 'new' (default: mynewworkflow.go).",
			Action: func(c *cli.Context) error {
				wfcode := `// Workflow written in SciPipe.
// For more information about SciPipe, see: http://scipipe.org
package main

import sp "github.com/scipipe/scipipe"

func main() {
	// --------------------------------
	// Create a pipeline runner
	// --------------------------------

	run := sp.NewPipelineRunner()

	// --------------------------------
	// Initialize processes and add to runner
	// --------------------------------

	foo := sp.NewProc("fooer",
		"echo foo > {o:foo}")
	foo.SetPathStatic("foo", "foo.txt")
	run.AddProcess(foo)

	f2b := sp.NewProc("foo2bar",
		"sed 's/foo/bar/g' {i:foo} > {o:bar}")
	f2b.SetPathExtend("foo", "bar", ".bar.txt")
	run.AddProcess(f2b)

	snk := sp.NewSink()
	run.AddProcess(snk)

	// --------------------------------
	// Connect workflow dependency network
	// --------------------------------

	f2b.In["foo"].Connect(foo.Out["foo"])
	snk.Connect(f2b.Out["bar"])

	// --------------------------------
	// Run the pipeline!
	// --------------------------------

	run.Run()
}`
				fileName := c.Args().First()
				if fileName == "" {
					fileName = "new_scipipe_workflow.go"
					fmt.Printf("No filename specified, so using the default '%s' ...\n", fileName)
				}

				f, err := os.Create(fileName)
				if err != nil {
					errLog.Println("Could not create file:", fileName)
					os.Exit(1)
				}
				defer f.Close()

				_, err = f.Write([]byte(wfcode))
				if err != nil {
					errLog.Println("Could not write to file:", fileName)
					os.Exit(1)
				}

				fmt.Println("Successfully wrote new workflow file to:", fileName)
				return nil
			},
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

// The git commit that was compiled. This will be filled in by the compiler.
var GitCommit string

// The main version number that is being run at the moment.
const Version = "0.8.0"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
const VersionPrerelease = ""
<commit_msg>Puts tree back into development mode.<commit_after>package main

// The git commit that was compiled. This will be filled in by the compiler.
var GitCommit string

// The main version number that is being run at the moment.
const Version = "0.8.1"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
const VersionPrerelease = "dev"
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"log"
	"os"

	"github.com/davecgh/go-spew/spew"

	"github.com/donatj/sqlread"
	"github.com/donatj/sqlread/mapcache"
)

var filename string

var (
	nocache = flag.Bool("nocache", false, "disable caching")
)

func init() {
	flag.Parse()

	filename = flag.Arg(0)
}

func main() {
	// return
	log.Println("starting initial pass")

	unbuff, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}

	cache := mapcache.New(unbuff)
	tree, err := cache.Get()
	if err != nil && err != mapcache.ErrCacheMiss {
		log.Fatal(err)
	}

	fmt.Println(*nocache)

	if err == mapcache.ErrCacheMiss || *nocache {
		l, li := sqlread.Lex(unbuff)
		go func() {
			l.Run(sqlread.StartState)
		}()

		sp := sqlread.NewSummaryParser()

		p := sqlread.Parse(li)
		err = p.Run(sp.ParseStart)
		if err != nil {
			log.Fatal(err)
		}

		if !*nocache {
			cache.Store(sp.Tree)
		}

		tree = sp.Tree
	} else {
		log.Println("loaded from cache")
	}

	log.Println("finished initial pass")

	//for tbl, _ := range t {
	//	fmt.Println(tbl)
	//}

	_ = tree

	interactive(tree, unbuff)
}

func interactive(tree sqlread.SummaryTree, buff io.ReaderAt) {
	w := csv.NewWriter(os.Stdout)
	sw := NewStdinWrap(os.Stdin)

	intp := sqlread.Intp{}
	for {
		stdinlex, stdli := sqlread.Lex(sw)
		go func() {
			stdinlex.Run(intp.StartIntpState)
		}()

		qp := sqlread.NewQueryParser()

		p := sqlread.Parse(stdli)
		err := p.Run(qp.ParseStart)
		if err != nil {
			log.Println("query error: ", err)
			sw.Flush()
			continue
		}

		spew.Dump(qp.Tree)

	queryloop:
		for _, qry := range qp.Tree.Queries {
			tbl, tok := tree[qry.Table]
			if !tok {
				log.Printf("table `%s` not found", qry.Table)
				continue
			}

			colind := []int{}

			for _, col := range qry.Columns {
				found := false
				for tci, tcol := range tbl.Cols {
					if col == "*" || col == tcol.Name {
						found = true
						colind = append(colind, tci)
					}
				}

				if !found {
					log.Printf("error: column `%s` not found", col)
					continue queryloop
				}
			}

			spew.Dump(colind)

			for _, loc := range tbl.DataLocs {
				start := loc.Start.Pos
				end := loc.End.Pos

				sl, sli := sqlread.LexSection(buff, start, end-start+1)
				go func() {
					sl.Run(sqlread.StartState)
				}()

				sp := sqlread.NewInsertDetailParser()

				spr := sqlread.Parse(sli)
				go func() {
					err := spr.Run(sp.ParseStart)
					if err != nil {
						log.Fatal(err)
					}
				}()

				for {
					row, ok := <-sp.Out
					if !ok {
						w.Flush()
						break
					}

					out := make([]string, len(colind))
					for i, ci := range colind {
						out[i] = row[ci]
					}

					w.Write(out)
				}
			}
		}

		for i := uint(0); i < qp.Tree.ShowTables; i++ {
			for cv, _ := range tree {
				w.Write([]string{cv})
			}

			w.Flush()
		}

		for _, sctbl := range qp.Tree.ShowColumns {
			tbl, tok := tree[sctbl]
			if !tok {
				log.Printf("table `%s` not found", sctbl)
			}

			for _, col := range tbl.Cols {
				w.Write([]string{col.Name, col.Type})
			}

			w.Flush()
		}

		if qp.Tree.Quit {
			return
		}

		// for {
		// 	x, ok := <-stdli
		// 	if !ok {
		// 		break
		// 	}

		// 	spew.Dump(x)

		// }

		sw.Flush()
		log.Println("restarting lexer")
	}
}
<commit_msg>Small refactor<commit_after>package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"log"
	"os"

	"github.com/davecgh/go-spew/spew"

	"github.com/donatj/sqlread"
	"github.com/donatj/sqlread/mapcache"
)

var filename string

var (
	nocache = flag.Bool("nocache", false, "disable caching")
)

func init() {
	flag.Parse()

	filename = flag.Arg(0)
}

func main() {
	// return
	log.Println("starting initial pass")

	unbuff, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}

	cache := mapcache.New(unbuff)
	tree, err := cache.Get()
	if err != nil && err != mapcache.ErrCacheMiss {
		log.Fatal(err)
	}

	fmt.Println(*nocache)

	if err == mapcache.ErrCacheMiss || *nocache {
		l, li := sqlread.Lex(unbuff)
		go func() {
			l.Run(sqlread.StartState)
		}()

		sp := sqlread.NewSummaryParser()

		p := sqlread.Parse(li)
		err = p.Run(sp.ParseStart)
		if err != nil {
			log.Fatal(err)
		}

		if !*nocache {
			cache.Store(sp.Tree)
		}

		tree = sp.Tree
	} else {
		log.Println("loaded from cache")
	}

	log.Println("finished initial pass")

	//for tbl, _ := range t {
	//	fmt.Println(tbl)
	//}

	_ = tree

	interactive(tree, unbuff)
}

func interactive(tree sqlread.SummaryTree, buff io.ReaderAt) {
	w := csv.NewWriter(os.Stdout)
	sw := NewStdinWrap(os.Stdin)

	intp := sqlread.Intp{}
	for {
		stdinlex, stdli := sqlread.Lex(sw)
		go func() {
			stdinlex.Run(intp.StartIntpState)
		}()

		qp := sqlread.NewQueryParser()

		p := sqlread.Parse(stdli)
		err := p.Run(qp.ParseStart)
		if err != nil {
			log.Println("query error: ", err)
			sw.Flush()
			continue
		}

	queryloop:
		for _, qry := range qp.Tree.Queries {
			tbl, tok := tree[qry.Table]
			if !tok {
				log.Printf("table `%s` not found", qry.Table)
				continue
			}

			colind := []int{}

			for _, col := range qry.Columns {
				found := false
				for tci, tcol := range tbl.Cols {
					if col == "*" || col == tcol.Name {
						found = true
						colind = append(colind, tci)
					}
				}

				if !found {
					log.Printf("error: column `%s` not found", col)
					continue queryloop
				}
			}

			for _, loc := range tbl.DataLocs {
				start := loc.Start.Pos
				end := loc.End.Pos

				sl, sli := sqlread.LexSection(buff, start, end-start+1)
				go func() {
					sl.Run(sqlread.StartState)
				}()

				sp := sqlread.NewInsertDetailParser()

				spr := sqlread.Parse(sli)
				go func() {
					err := spr.Run(sp.ParseStart)
					if err != nil {
						log.Fatal(err)
					}
				}()

				for {
					row, ok := <-sp.Out
					if !ok {
						w.Flush()
						break
					}

					out := make([]string, len(colind))
					for i, ci := range colind {
						out[i] = row[ci]
					}

					w.Write(out)
				}
			}
		}

		for i := uint(0); i < qp.Tree.ShowTables; i++ {
			showTables(tree, w)
		}

		for _, sctbl := range qp.Tree.ShowColumns {
			if err := showColumns(tree, sctbl, w); err != nil {
				log.Println(err)
			}
		}

		if qp.Tree.Quit {
			return
		}

		if intp.EOF {
			break
		}

		sw.Flush()
		log.Println("restarting lexer")
	}
}

func showColumns(tree sqlread.SummaryTree, sctbl string, w *csv.Writer) error {
	tbl, tok := tree[sctbl]
	if !tok {
		return fmt.Errorf("table `%s` not found", sctbl)
	}
	for _, col := range tbl.Cols {
		w.Write([]string{col.Name, col.Type})
	}
	w.Flush()

	return nil
}

func showTables(tree sqlread.SummaryTree, w *csv.Writer) {
	for cv, _ := range tree {
		w.Write([]string{cv})
	}
	w.Flush()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"os/user"
	"strings"

	"github.com/alexcesaro/log"
	"github.com/alexcesaro/log/golog"
	flags "github.com/jessevdk/go-flags"
	"golang.org/x/crypto/ssh"

	sshchat "github.com/shazow/ssh-chat"
	"github.com/shazow/ssh-chat/chat"
	"github.com/shazow/ssh-chat/chat/message"
	"github.com/shazow/ssh-chat/sshd"
)
import _ "net/http/pprof"

// Version of the binary, assigned during build.
var Version string = "dev"

// Options contains the flag options
type Options struct {
	Verbose   []bool `short:"v" long:"verbose" description:"Show verbose logging."`
	Version   bool   `long:"version" description:"Print version and exit."`
	Identity  string `short:"i" long:"identity" description:"Private key to identify server with." default:"~/.ssh/id_rsa"`
	Bind      string `long:"bind" description:"Host and port to listen on." default:"0.0.0.0:2022"`
	Admin     string `long:"admin" description:"File of public keys who are admins."`
	Whitelist string `long:"whitelist" description:"Optional file of public keys who are allowed to connect."`
	Motd      string `long:"motd" description:"Optional Message of the Day file."`
	Log       string `long:"log" description:"Write chat log to this file."`
	Pprof     int    `long:"pprof" description:"Enable pprof http server for profiling."`
}

var logLevels = []log.Level{
	log.Warning,
	log.Info,
	log.Debug,
}

func fail(code int, format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format, args...)
	os.Exit(code)
}

func main() {
	options := Options{}
	parser := flags.NewParser(&options, flags.Default)
	p, err := parser.Parse()
	if err != nil {
		if p == nil {
			fmt.Print(err)
		}
		return
	}

	if options.Pprof != 0 {
		go func() {
			fmt.Println(http.ListenAndServe(fmt.Sprintf("localhost:%d", options.Pprof), nil))
		}()
	}

	if options.Version {
		fmt.Println(Version)
		return
	}

	// Figure out the log level
	numVerbose := len(options.Verbose)
	if numVerbose > len(logLevels) {
		numVerbose = len(logLevels) - 1
	}

	logLevel := logLevels[numVerbose]
	sshchat.SetLogger(golog.New(os.Stderr, logLevel))

	if logLevel == log.Debug {
		// Enable logging from submodules
		chat.SetLogger(os.Stderr)
		sshd.SetLogger(os.Stderr)
		message.SetLogger(os.Stderr)
	}

	privateKeyPath := options.Identity
	if strings.HasPrefix(privateKeyPath, "~/") {
		user, err := user.Current()
		if err == nil {
			privateKeyPath = strings.Replace(privateKeyPath, "~", user.HomeDir, 1)
		}
	}

	privateKey, err := ReadPrivateKey(privateKeyPath)
	if err != nil {
		fail(2, "Couldn't read private key: %v\n", err)
	}

	signer, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		fail(3, "Failed to parse key: %v\n", err)
	}

	auth := sshchat.NewAuth()
	config := sshd.MakeAuth(auth)
	config.AddHostKey(signer)
	config.ServerVersion = "SSH-2.0-Go ssh-chat"

	s, err := sshd.ListenSSH(options.Bind, config)
	if err != nil {
		fail(4, "Failed to listen on socket: %v\n", err)
	}
	defer s.Close()
	s.RateLimit = sshd.NewInputLimiter

	fmt.Printf("Listening for connections on %v\n", s.Addr().String())

	host := sshchat.NewHost(s, auth)
	host.SetTheme(message.Themes[0])
	host.Version = Version

	err = fromFile(options.Admin, func(line []byte) error {
		key, _, _, _, err := ssh.ParseAuthorizedKey(line)
		if err != nil {
			return err
		}
		auth.Op(key, 0)
		return nil
	})
	if err != nil {
		fail(5, "Failed to load admins: %v\n", err)
	}

	err = fromFile(options.Whitelist, func(line []byte) error {
		key, _, _, _, err := ssh.ParseAuthorizedKey(line)
		if err != nil {
			return err
		}
		auth.Whitelist(key, 0)
		return nil
	})
	if err != nil {
		fail(6, "Failed to load whitelist: %v\n", err)
	}

	if options.Motd != "" {
		motd, err := ioutil.ReadFile(options.Motd)
		if err != nil {
			fail(7, "Failed to load MOTD file: %v\n", err)
		}
		motdString := string(motd)
		// hack to normalize line endings into \r\n
		motdString = strings.Replace(motdString, "\r\n", "\n", -1)
		motdString = strings.Replace(motdString, "\n", "\r\n", -1)
		host.SetMotd(motdString)
	}

	if options.Log == "-" {
		host.SetLogging(os.Stdout)
	} else if options.Log != "" {
		fp, err := os.OpenFile(options.Log, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
		if err != nil {
			fail(8, "Failed to open log file for writing: %v", err)
		}
		host.SetLogging(fp)
	}

	go host.Serve()

	// Construct interrupt handler
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt)

	<-sig // Wait for ^C signal
	fmt.Fprintln(os.Stderr, "Interrupt signal detected, shutting down.")
}

func fromFile(path string, handler func(line []byte) error) error {
	if path == "" {
		// Skip
		return nil
	}

	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		err := handler(scanner.Bytes())
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>main: --admin and --whitelist: Skip non-key lines<commit_after>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"os/user"
	"strings"

	"github.com/alexcesaro/log"
	"github.com/alexcesaro/log/golog"
	flags "github.com/jessevdk/go-flags"
	"golang.org/x/crypto/ssh"

	sshchat "github.com/shazow/ssh-chat"
	"github.com/shazow/ssh-chat/chat"
	"github.com/shazow/ssh-chat/chat/message"
	"github.com/shazow/ssh-chat/sshd"
)
import _ "net/http/pprof"

// Version of the binary, assigned during build.
var Version string = "dev"

// Options contains the flag options
type Options struct {
	Verbose   []bool `short:"v" long:"verbose" description:"Show verbose logging."`
	Version   bool   `long:"version" description:"Print version and exit."`
	Identity  string `short:"i" long:"identity" description:"Private key to identify server with." default:"~/.ssh/id_rsa"`
	Bind      string `long:"bind" description:"Host and port to listen on." default:"0.0.0.0:2022"`
	Admin     string `long:"admin" description:"File of public keys who are admins."`
	Whitelist string `long:"whitelist" description:"Optional file of public keys who are allowed to connect."`
	Motd      string `long:"motd" description:"Optional Message of the Day file."`
	Log       string `long:"log" description:"Write chat log to this file."`
	Pprof     int    `long:"pprof" description:"Enable pprof http server for profiling."`
}

var logLevels = []log.Level{
	log.Warning,
	log.Info,
	log.Debug,
}

func fail(code int, format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format, args...)
	os.Exit(code)
}

func main() {
	options := Options{}
	parser := flags.NewParser(&options, flags.Default)
	p, err := parser.Parse()
	if err != nil {
		if p == nil {
			fmt.Print(err)
		}
		return
	}

	if options.Pprof != 0 {
		go func() {
			fmt.Println(http.ListenAndServe(fmt.Sprintf("localhost:%d", options.Pprof), nil))
		}()
	}

	if options.Version {
		fmt.Println(Version)
		return
	}

	// Figure out the log level
	numVerbose := len(options.Verbose)
	if numVerbose > len(logLevels) {
		numVerbose = len(logLevels) - 1
	}

	logLevel := logLevels[numVerbose]
	sshchat.SetLogger(golog.New(os.Stderr, logLevel))

	if logLevel == log.Debug {
		// Enable logging from submodules
		chat.SetLogger(os.Stderr)
		sshd.SetLogger(os.Stderr)
		message.SetLogger(os.Stderr)
	}

	privateKeyPath := options.Identity
	if strings.HasPrefix(privateKeyPath, "~/") {
		user, err := user.Current()
		if err == nil {
			privateKeyPath = strings.Replace(privateKeyPath, "~", user.HomeDir, 1)
		}
	}

	privateKey, err := ReadPrivateKey(privateKeyPath)
	if err != nil {
		fail(2, "Couldn't read private key: %v\n", err)
	}

	signer, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		fail(3, "Failed to parse key: %v\n", err)
	}

	auth := sshchat.NewAuth()
	config := sshd.MakeAuth(auth)
	config.AddHostKey(signer)
	config.ServerVersion = "SSH-2.0-Go ssh-chat"

	s, err := sshd.ListenSSH(options.Bind, config)
	if err != nil {
		fail(4, "Failed to listen on socket: %v\n", err)
	}
	defer s.Close()
	s.RateLimit = sshd.NewInputLimiter

	fmt.Printf("Listening for connections on %v\n", s.Addr().String())

	host := sshchat.NewHost(s, auth)
	host.SetTheme(message.Themes[0])
	host.Version = Version

	err = fromFile(options.Admin, func(line []byte) error {
		key, _, _, _, err := ssh.ParseAuthorizedKey(line)
		if err != nil {
			if err.Error() == "ssh: no key found" {
				return nil // Skip line
			}
			return err
		}
		auth.Op(key, 0)
		return nil
	})
	if err != nil {
		fail(5, "Failed to load admins: %v\n", err)
	}

	err = fromFile(options.Whitelist, func(line []byte) error {
		key, _, _, _, err := ssh.ParseAuthorizedKey(line)
		if err != nil {
			if err.Error() == "ssh: no key found" {
				return nil // Skip line
			}
			return err
		}
		auth.Whitelist(key, 0)
		return nil
	})
	if err != nil {
		fail(6, "Failed to load whitelist: %v\n", err)
	}

	if options.Motd != "" {
		motd, err := ioutil.ReadFile(options.Motd)
		if err != nil {
			fail(7, "Failed to load MOTD file: %v\n", err)
		}
		motdString := string(motd)
		// hack to normalize line endings into \r\n
		motdString = strings.Replace(motdString, "\r\n", "\n", -1)
		motdString = strings.Replace(motdString, "\n", "\r\n", -1)
		host.SetMotd(motdString)
	}

	if options.Log == "-" {
		host.SetLogging(os.Stdout)
	} else if options.Log != "" {
		fp, err := os.OpenFile(options.Log, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
		if err != nil {
			fail(8, "Failed to open log file for writing: %v", err)
		}
		host.SetLogging(fp)
	}

	go host.Serve()

	// Construct interrupt handler
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt)

	<-sig // Wait for ^C signal
	fmt.Fprintln(os.Stderr, "Interrupt signal detected, shutting down.")
}

func fromFile(path string, handler func(line []byte) error) error {
	if path == "" {
		// Skip
		return nil
	}

	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		err := handler(scanner.Bytes())
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 * MinIO Cloud Storage, (C) 2018, 2019 MinIO, Inc.
 *
 * 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 cmd

import (
	"context"
	"fmt"
	"net/http"

	"github.com/gorilla/mux"
	"github.com/minio/minio/cmd/logger"
	"github.com/minio/minio/pkg/auth"
	"github.com/minio/minio/pkg/iam/validator"
)

const (
	// STS API version.
	stsAPIVersion = "2011-06-15"

	// STS API action constants
	clientGrants = "AssumeRoleWithClientGrants"
	webIdentity  = "AssumeRoleWithWebIdentity"
	assumeRole   = "AssumeRole"
)

// stsAPIHandlers implements and provides http handlers for AWS STS API.
type stsAPIHandlers struct{}

// registerSTSRouter - registers AWS STS compatible APIs.
func registerSTSRouter(router *mux.Router) {
	// Initialize STS.
	sts := &stsAPIHandlers{}

	// STS Router
	stsRouter := router.NewRoute().PathPrefix("/").Subrouter()

	// Assume roles with no JWT, handles AssumeRole.
	stsRouter.Methods("POST").HeadersRegexp("Content-Type", "application/x-www-form-urlencoded*").
		HeadersRegexp("Authorization", "AWS4-HMAC-SHA256*").
		HandlerFunc(httpTraceAll(sts.AssumeRole))

	// Assume roles with JWT handler, handles both ClientGrants and WebIdentity.
	stsRouter.Methods("POST").HeadersRegexp("Content-Type", "application/x-www-form-urlencoded*").
		HandlerFunc(httpTraceAll(sts.AssumeRoleWithJWT))

	// AssumeRoleWithClientGrants
	stsRouter.Methods("POST").HandlerFunc(httpTraceAll(sts.AssumeRoleWithClientGrants)).
		Queries("Action", clientGrants).
		Queries("Version", stsAPIVersion).
		Queries("Token", "{Token:.*}")

	// AssumeRoleWithWebIdentity
	stsRouter.Methods("POST").HandlerFunc(httpTraceAll(sts.AssumeRoleWithWebIdentity)).
		Queries("Action", webIdentity).
		Queries("Version", stsAPIVersion).
		Queries("WebIdentityToken", "{Token:.*}")

}

func checkAssumeRoleAuth(ctx context.Context, r *http.Request) (user auth.Credentials, stsErr STSErrorCode) {
	switch getRequestAuthType(r) {
	default:
		return user, ErrSTSAccessDenied
	case authTypeSigned:
		s3Err := isReqAuthenticated(ctx, r, globalServerConfig.GetRegion(), serviceSTS)
		if STSErrorCode(s3Err) != ErrSTSNone {
			return user, STSErrorCode(s3Err)
		}
		var owner bool
		user, owner, s3Err = getReqAccessKeyV4(r, globalServerConfig.GetRegion(), serviceSTS)
		if STSErrorCode(s3Err) != ErrSTSNone {
			return user, STSErrorCode(s3Err)
		}
		// Root credentials are not allowed to use STS API
		if owner {
			return user, ErrSTSAccessDenied
		}
	}

	// Session tokens are not allowed in STS AssumeRole requests.
	if getSessionToken(r) != "" {
		return user, ErrSTSAccessDenied
	}

	return user, ErrSTSNone
}

// AssumeRole - implementation of AWS STS API AssumeRole to get temporary
// credentials for regular users on Minio.
// https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html
func (sts *stsAPIHandlers) AssumeRole(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "AssumeRole")

	user, stsErr := checkAssumeRoleAuth(ctx, r)
	if stsErr != ErrSTSNone {
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(stsErr))
		return
	}

	if err := r.ParseForm(); err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	if r.Form.Get("Policy") != "" {
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	if r.Form.Get("Version") != stsAPIVersion {
		logger.LogIf(ctx, fmt.Errorf("Invalid STS API version %s, expecting %s", r.Form.Get("Version"), stsAPIVersion))
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSMissingParameter))
		return
	}

	action := r.Form.Get("Action")
	switch action {
	case assumeRole:
	default:
		logger.LogIf(ctx, fmt.Errorf("Unsupported action %s", action))
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	ctx = newContext(r, w, action)
	defer logger.AuditLog(w, r, action, nil)

	var err error
	m := make(map[string]interface{})
	m["exp"], err = validator.GetDefaultExpiration(r.Form.Get("DurationSeconds"))
	if err != nil {
		switch err {
		case validator.ErrInvalidDuration:
			writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		default:
			logger.LogIf(ctx, err)
			writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		}
		return
	}

	policyName, err := globalIAMSys.GetUserPolicy(user.AccessKey)
	if err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	// This policy is the policy associated with the user
	// requesting for temporary credentials. The temporary
	// credentials will inherit the same policy requirements.
	m["policy"] = policyName

	secret := globalServerConfig.GetCredential().SecretKey
	cred, err := auth.GetNewCredentialsWithMetadata(m, secret)
	if err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInternalError))
		return
	}

	// Set the newly generated credentials.
	if err = globalIAMSys.SetTempUser(cred.AccessKey, cred, policyName); err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInternalError))
		return
	}

	// Notify all other MinIO peers to reload temp users
	for _, nerr := range globalNotificationSys.LoadUsers() {
		if nerr.Err != nil {
			logger.GetReqInfo(ctx).SetTags("peerAddress", nerr.Host.String())
			logger.LogIf(ctx, nerr.Err)
		}
	}

	assumeRoleResponse := &AssumeRoleResponse{
		Result: AssumeRoleResult{
			Credentials: cred,
		},
	}

	assumeRoleResponse.ResponseMetadata.RequestID = w.Header().Get(responseRequestIDKey)
	writeSuccessResponseXML(w, encodeResponse(assumeRoleResponse))
}

func (sts *stsAPIHandlers) AssumeRoleWithJWT(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "AssumeRoleJWTCommon")

	// Parse the incoming form data.
	if err := r.ParseForm(); err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	if r.Form.Get("Policy") != "" {
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	if r.Form.Get("Version") != stsAPIVersion {
		logger.LogIf(ctx, fmt.Errorf("Invalid STS API version %s, expecting %s", r.Form.Get("Version"), stsAPIVersion))
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSMissingParameter))
		return
	}

	action := r.Form.Get("Action")
	switch action {
	case clientGrants, webIdentity:
	default:
		logger.LogIf(ctx, fmt.Errorf("Unsupported action %s", action))
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	ctx = newContext(r, w, action)
	defer logger.AuditLog(w, r, action, nil)

	if globalIAMValidators == nil {
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSNotInitialized))
		return
	}

	v, err := globalIAMValidators.Get("jwt")
	if err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	token := r.Form.Get("Token")
	if token == "" {
		token = r.Form.Get("WebIdentityToken")
	}

	m, err := v.Validate(token, r.Form.Get("DurationSeconds"))
	if err != nil {
		switch err {
		case validator.ErrTokenExpired:
			switch action {
			case clientGrants:
				writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSClientGrantsExpiredToken))
			case webIdentity:
				writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSWebIdentityExpiredToken))
			}
			return
		case validator.ErrInvalidDuration:
			writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
			return
		}
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	secret := globalServerConfig.GetCredential().SecretKey
	cred, err := auth.GetNewCredentialsWithMetadata(m, secret)
	if err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInternalError))
		return
	}

	// JWT has requested a custom claim with policy value set.
	// This is a MinIO STS API specific value, this value should
	// be set and configured on your identity provider as part of
	// JWT custom claims.
	var policyName string
	if v, ok := m["policy"]; ok {
		policyName, _ = v.(string)
	}

	var subFromToken string
	if v, ok := m["sub"]; ok {
		subFromToken, _ = v.(string)
	}

	// Set the newly generated credentials.
	if err = globalIAMSys.SetTempUser(cred.AccessKey, cred, policyName); err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInternalError))
		return
	}

	// Notify all other MinIO peers to reload temp users
	for _, nerr := range globalNotificationSys.LoadUsers() {
		if nerr.Err != nil {
			logger.GetReqInfo(ctx).SetTags("peerAddress", nerr.Host.String())
			logger.LogIf(ctx, nerr.Err)
		}
	}

	var encodedSuccessResponse []byte
	switch action {
	case clientGrants:
		clientGrantsResponse := &AssumeRoleWithClientGrantsResponse{
			Result: ClientGrantsResult{
				Credentials:      cred,
				SubjectFromToken: subFromToken,
			},
		}
		clientGrantsResponse.ResponseMetadata.RequestID = w.Header().Get(responseRequestIDKey)
		encodedSuccessResponse = encodeResponse(clientGrantsResponse)
	case webIdentity:
		webIdentityResponse := &AssumeRoleWithWebIdentityResponse{
			Result: WebIdentityResult{
				Credentials:                 cred,
				SubjectFromWebIdentityToken: subFromToken,
			},
		}
		webIdentityResponse.ResponseMetadata.RequestID = w.Header().Get(responseRequestIDKey)
		encodedSuccessResponse = encodeResponse(webIdentityResponse)
	}

	writeSuccessResponseXML(w, encodedSuccessResponse)
}

// AssumeRoleWithWebIdentity - implementation of AWS STS API supporting OAuth2.0
// users from web identity provider such as Facebook, Google, or any OpenID
// Connect-compatible identity provider.
//
// Eg:-
//    $ curl https://minio:9000/?Action=AssumeRoleWithWebIdentity&WebIdentityToken=<jwt>
func (sts *stsAPIHandlers) AssumeRoleWithWebIdentity(w http.ResponseWriter, r *http.Request) {
	sts.AssumeRoleWithJWT(w, r)
}

// AssumeRoleWithClientGrants - implementation of AWS STS extension API supporting
// OAuth2.0 client credential grants.
//
// Eg:-
//    $ curl https://minio:9000/?Action=AssumeRoleWithClientGrants&Token=<jwt>
func (sts *stsAPIHandlers) AssumeRoleWithClientGrants(w http.ResponseWriter, r *http.Request) {
	sts.AssumeRoleWithJWT(w, r)
}
<commit_msg>Fix STS AssumeRole route conflict with MultipartUpload (#7574)<commit_after>/*
 * MinIO Cloud Storage, (C) 2018, 2019 MinIO, Inc.
 *
 * 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 cmd

import (
	"context"
	"fmt"
	"net/http"

	"github.com/gorilla/mux"
	"github.com/minio/minio/cmd/logger"
	"github.com/minio/minio/pkg/auth"
	"github.com/minio/minio/pkg/iam/validator"
	"github.com/minio/minio/pkg/wildcard"
)

const (
	// STS API version.
	stsAPIVersion = "2011-06-15"

	// STS API action constants
	clientGrants = "AssumeRoleWithClientGrants"
	webIdentity  = "AssumeRoleWithWebIdentity"
	assumeRole   = "AssumeRole"
)

// stsAPIHandlers implements and provides http handlers for AWS STS API.
type stsAPIHandlers struct{}

// registerSTSRouter - registers AWS STS compatible APIs.
func registerSTSRouter(router *mux.Router) {
	// Initialize STS.
	sts := &stsAPIHandlers{}

	// STS Router
	stsRouter := router.NewRoute().PathPrefix("/").Subrouter()

	// Assume roles with no JWT, handles AssumeRole.
	stsRouter.Methods("POST").MatcherFunc(func(r *http.Request, rm *mux.RouteMatch) bool {
		ctypeOk := wildcard.MatchSimple("application/x-www-form-urlencoded*", r.Header.Get("Content-Type"))
		authOk := wildcard.MatchSimple("AWS4-HMAC-SHA256*", r.Header.Get("Authorization"))
		noQueries := len(r.URL.Query()) == 0
		return ctypeOk && authOk && noQueries
	}).HandlerFunc(httpTraceAll(sts.AssumeRole))

	// Assume roles with JWT handler, handles both ClientGrants and WebIdentity.
	stsRouter.Methods("POST").MatcherFunc(func(r *http.Request, rm *mux.RouteMatch) bool {
		ctypeOk := wildcard.MatchSimple("application/x-www-form-urlencoded*", r.Header.Get("Content-Type"))
		noQueries := len(r.URL.Query()) == 0
		return ctypeOk && noQueries
	}).HandlerFunc(httpTraceAll(sts.AssumeRoleWithJWT))

	// AssumeRoleWithClientGrants
	stsRouter.Methods("POST").HandlerFunc(httpTraceAll(sts.AssumeRoleWithClientGrants)).
		Queries("Action", clientGrants).
		Queries("Version", stsAPIVersion).
		Queries("Token", "{Token:.*}")

	// AssumeRoleWithWebIdentity
	stsRouter.Methods("POST").HandlerFunc(httpTraceAll(sts.AssumeRoleWithWebIdentity)).
		Queries("Action", webIdentity).
		Queries("Version", stsAPIVersion).
		Queries("WebIdentityToken", "{Token:.*}")

}

func checkAssumeRoleAuth(ctx context.Context, r *http.Request) (user auth.Credentials, stsErr STSErrorCode) {
	switch getRequestAuthType(r) {
	default:
		return user, ErrSTSAccessDenied
	case authTypeSigned:
		s3Err := isReqAuthenticated(ctx, r, globalServerConfig.GetRegion(), serviceSTS)
		if STSErrorCode(s3Err) != ErrSTSNone {
			return user, STSErrorCode(s3Err)
		}
		var owner bool
		user, owner, s3Err = getReqAccessKeyV4(r, globalServerConfig.GetRegion(), serviceSTS)
		if STSErrorCode(s3Err) != ErrSTSNone {
			return user, STSErrorCode(s3Err)
		}
		// Root credentials are not allowed to use STS API
		if owner {
			return user, ErrSTSAccessDenied
		}
	}

	// Session tokens are not allowed in STS AssumeRole requests.
	if getSessionToken(r) != "" {
		return user, ErrSTSAccessDenied
	}

	return user, ErrSTSNone
}

// AssumeRole - implementation of AWS STS API AssumeRole to get temporary
// credentials for regular users on Minio.
// https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html
func (sts *stsAPIHandlers) AssumeRole(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "AssumeRole")

	user, stsErr := checkAssumeRoleAuth(ctx, r)
	if stsErr != ErrSTSNone {
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(stsErr))
		return
	}

	if err := r.ParseForm(); err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	if r.Form.Get("Policy") != "" {
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	if r.Form.Get("Version") != stsAPIVersion {
		logger.LogIf(ctx, fmt.Errorf("Invalid STS API version %s, expecting %s", r.Form.Get("Version"), stsAPIVersion))
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSMissingParameter))
		return
	}

	action := r.Form.Get("Action")
	switch action {
	case assumeRole:
	default:
		logger.LogIf(ctx, fmt.Errorf("Unsupported action %s", action))
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	ctx = newContext(r, w, action)
	defer logger.AuditLog(w, r, action, nil)

	var err error
	m := make(map[string]interface{})
	m["exp"], err = validator.GetDefaultExpiration(r.Form.Get("DurationSeconds"))
	if err != nil {
		switch err {
		case validator.ErrInvalidDuration:
			writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		default:
			logger.LogIf(ctx, err)
			writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		}
		return
	}

	policyName, err := globalIAMSys.GetUserPolicy(user.AccessKey)
	if err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	// This policy is the policy associated with the user
	// requesting for temporary credentials. The temporary
	// credentials will inherit the same policy requirements.
	m["policy"] = policyName

	secret := globalServerConfig.GetCredential().SecretKey
	cred, err := auth.GetNewCredentialsWithMetadata(m, secret)
	if err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInternalError))
		return
	}

	// Set the newly generated credentials.
	if err = globalIAMSys.SetTempUser(cred.AccessKey, cred, policyName); err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInternalError))
		return
	}

	// Notify all other MinIO peers to reload temp users
	for _, nerr := range globalNotificationSys.LoadUsers() {
		if nerr.Err != nil {
			logger.GetReqInfo(ctx).SetTags("peerAddress", nerr.Host.String())
			logger.LogIf(ctx, nerr.Err)
		}
	}

	assumeRoleResponse := &AssumeRoleResponse{
		Result: AssumeRoleResult{
			Credentials: cred,
		},
	}

	assumeRoleResponse.ResponseMetadata.RequestID = w.Header().Get(responseRequestIDKey)
	writeSuccessResponseXML(w, encodeResponse(assumeRoleResponse))
}

func (sts *stsAPIHandlers) AssumeRoleWithJWT(w http.ResponseWriter, r *http.Request) {
	ctx := newContext(r, w, "AssumeRoleJWTCommon")

	// Parse the incoming form data.
	if err := r.ParseForm(); err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	if r.Form.Get("Policy") != "" {
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	if r.Form.Get("Version") != stsAPIVersion {
		logger.LogIf(ctx, fmt.Errorf("Invalid STS API version %s, expecting %s", r.Form.Get("Version"), stsAPIVersion))
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSMissingParameter))
		return
	}

	action := r.Form.Get("Action")
	switch action {
	case clientGrants, webIdentity:
	default:
		logger.LogIf(ctx, fmt.Errorf("Unsupported action %s", action))
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	ctx = newContext(r, w, action)
	defer logger.AuditLog(w, r, action, nil)

	if globalIAMValidators == nil {
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSNotInitialized))
		return
	}

	v, err := globalIAMValidators.Get("jwt")
	if err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	token := r.Form.Get("Token")
	if token == "" {
		token = r.Form.Get("WebIdentityToken")
	}

	m, err := v.Validate(token, r.Form.Get("DurationSeconds"))
	if err != nil {
		switch err {
		case validator.ErrTokenExpired:
			switch action {
			case clientGrants:
				writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSClientGrantsExpiredToken))
			case webIdentity:
				writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSWebIdentityExpiredToken))
			}
			return
		case validator.ErrInvalidDuration:
			writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
			return
		}
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInvalidParameterValue))
		return
	}

	secret := globalServerConfig.GetCredential().SecretKey
	cred, err := auth.GetNewCredentialsWithMetadata(m, secret)
	if err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInternalError))
		return
	}

	// JWT has requested a custom claim with policy value set.
	// This is a MinIO STS API specific value, this value should
	// be set and configured on your identity provider as part of
	// JWT custom claims.
	var policyName string
	if v, ok := m["policy"]; ok {
		policyName, _ = v.(string)
	}

	var subFromToken string
	if v, ok := m["sub"]; ok {
		subFromToken, _ = v.(string)
	}

	// Set the newly generated credentials.
	if err = globalIAMSys.SetTempUser(cred.AccessKey, cred, policyName); err != nil {
		logger.LogIf(ctx, err)
		writeSTSErrorResponse(w, stsErrCodes.ToSTSErr(ErrSTSInternalError))
		return
	}

	// Notify all other MinIO peers to reload temp users
	for _, nerr := range globalNotificationSys.LoadUsers() {
		if nerr.Err != nil {
			logger.GetReqInfo(ctx).SetTags("peerAddress", nerr.Host.String())
			logger.LogIf(ctx, nerr.Err)
		}
	}

	var encodedSuccessResponse []byte
	switch action {
	case clientGrants:
		clientGrantsResponse := &AssumeRoleWithClientGrantsResponse{
			Result: ClientGrantsResult{
				Credentials:      cred,
				SubjectFromToken: subFromToken,
			},
		}
		clientGrantsResponse.ResponseMetadata.RequestID = w.Header().Get(responseRequestIDKey)
		encodedSuccessResponse = encodeResponse(clientGrantsResponse)
	case webIdentity:
		webIdentityResponse := &AssumeRoleWithWebIdentityResponse{
			Result: WebIdentityResult{
				Credentials:                 cred,
				SubjectFromWebIdentityToken: subFromToken,
			},
		}
		webIdentityResponse.ResponseMetadata.RequestID = w.Header().Get(responseRequestIDKey)
		encodedSuccessResponse = encodeResponse(webIdentityResponse)
	}

	writeSuccessResponseXML(w, encodedSuccessResponse)
}

// AssumeRoleWithWebIdentity - implementation of AWS STS API supporting OAuth2.0
// users from web identity provider such as Facebook, Google, or any OpenID
// Connect-compatible identity provider.
//
// Eg:-
//    $ curl https://minio:9000/?Action=AssumeRoleWithWebIdentity&WebIdentityToken=<jwt>
func (sts *stsAPIHandlers) AssumeRoleWithWebIdentity(w http.ResponseWriter, r *http.Request) {
	sts.AssumeRoleWithJWT(w, r)
}

// AssumeRoleWithClientGrants - implementation of AWS STS extension API supporting
// OAuth2.0 client credential grants.
//
// Eg:-
//    $ curl https://minio:9000/?Action=AssumeRoleWithClientGrants&Token=<jwt>
func (sts *stsAPIHandlers) AssumeRoleWithClientGrants(w http.ResponseWriter, r *http.Request) {
	sts.AssumeRoleWithJWT(w, r)
}
<|endoftext|>