<commit_before>package template

import "fmt"
import "bytes"
import "gnd.la/util/generic"

type instructions []inst

func (i instructions) replace(idx int, count int, repl []inst) []inst {
	// look for jumps before the block which need to be adjusted
	for ii, v := range i[:idx] {
		switch v.op {
		case opJMP, opJMPT, opJMPF, opNEXT:
			val := int(int32(v.val))
			if ii+val > idx {
				i[ii] = inst{v.op, valType(int32(val + len(repl) - count))}
			}
		}
	}
	// look for jumps after the block which need to be adjusted
	start := idx + count
	for ii, v := range i[start:] {
		switch v.op {
		case opJMP, opJMPT, opJMPF, opNEXT:
			val := int(int32(v.val))
			if ii+val < 0 {
				i[ii+start] = inst{v.op, valType(int32(val - len(repl) + count))}
			}
		}
	}
	var ret []inst
	ret = append(ret, i[:idx]...)
	ret = append(ret, repl...)
	ret = append(ret, i[idx+count:]...)
	return ret
}

func (i instructions) replaceTemplateInvocation(idx int, repl []inst) []inst {
	ti := i[idx]
	if ti.op != opTEMPLATE {
		panic(fmt.Errorf("OP at index %d is not opTEMPLATE", idx))
	}
	// Check if we can also remove the pipeline which provides arguments
	// to the template invocation.
	count := 1
	if idx < len(i)-1 {
		pi := i[idx+1]
		if pi.op == opPOP {
			// Remove the POP
			count++
			// This opPOP removes the arguments which were set up for
			// the template call
			if pi.val == 0 {
				// POP until the mark
				for {
					count++
					idx--
					if i[idx].op == opMARK {
						break
					}
				}
			} else {
				// Remove until we find as many pushes as this
				// pop removes
				stack := int(pi.val)
				for stack > 0 {
					count++
					idx--
					ii := i[idx]
					stack -= ii.pushes()
					stack += ii.pops()
				}
			}
		}
	}
	return i.replace(idx, count, repl)
}

// Returns true iff the program references the byte slice determined
// by val
func (p *program) referencesBytes(val valType) bool {
	for _, code := range p.code {
		for _, pi := range code {
			if pi.op == opWB && pi.val == val {
				return true
			}
		}
	}
	return false
}

// removeInternedBytes removes the internet byte slice determined
// by val, adjusting all required instructions.
func (p *program) removeInternedBytes(val valType) {
	idx := int(val)
	p.bs = append(p.bs[:idx], p.bs[idx+1:]...)
	for _, code := range p.code {
		for ii, pi := range code {
			if pi.op == opWB && pi.val >= val {
				ni := pi
				ni.val = pi.val - 1
				code[ii] = ni
			}
		}
	}
}

func (p *program) optimize() {
	p.removeEmptyTemplateInvocations()
	p.stitch()
	p.mergeWriteBytes()
}

// mergeWriteBytes merges adjacent opWB operations into a single WB
func (p *program) mergeWriteBytes() {
	removedReferences := make(map[valType]struct{})
	for name, code := range p.code {
		wb := -1
		ii := 0
		checkWBMerge := func() {
			if wb >= 0 {
				// We had a WB sequence, check its length
				// and merge it if appropriate
				count := ii - wb
				if count > 1 {
					var buf bytes.Buffer
					var refs []valType
					for c := wb; c < ii; c++ {
						wbInst := code[c]
						removedReferences[wbInst.val] = struct{}{}
						refs = append(refs, wbInst.val)
						buf.Write(p.bytesValue(wbInst.val))
					}
					repl := []inst{
						inst{op: opWB, val: p.internBytes(buf.Bytes())},
					}
					code = instructions(code).replace(wb, count, repl)
					compilerDebugf("merged %d WB instructions with byte slice refs %v at PC %d (new ref %d)\n",
						count, refs, wb, repl[0].val)
				}
			}
		}
		for ; ii < len(code); ii++ {
			v := code[ii]
			if v.op == opWB {
				if wb < 0 {
					wb = ii
				}
			} else {
				checkWBMerge()
				wb = -1
			}
		}
		checkWBMerge()
		p.code[name] = code
	}
	// Sort references from higher to lower, so we
	// don't invalidate them while they're being removed
	sortedReferences := make([]valType, 0, len(removedReferences))
	for r := range removedReferences {
		sortedReferences = append(sortedReferences, r)
	}
	generic.SortFunc(sortedReferences, func(a, b valType) bool {
		return a > b
	})
	compilerDebugln("candidates for byte slice removal", sortedReferences)
	for _, r := range sortedReferences {
		if !p.referencesBytes(r) {
			compilerDebugln("will remove byte slice reference", r)
			p.removeInternedBytes(r)
		}
	}
}

// Remove {{ template }} invocations which call into an
// empty template.
func (p *program) removeEmptyTemplateInvocations() {
	for name, code := range p.code {
		for ii := 0; ii < len(code); ii++ {
			v := code[ii]
			if v.op == opTEMPLATE {
				_, t := decodeVal(v.val)
				tmplName := p.strings[t]
				tmplCode := p.code[tmplName]
				if len(tmplCode) == 0 {
					// Empty template invocation, remove it
					compilerDebugf("remove empty template invocation %q from %q\n", tmplName, name)
					code = instructions(code).replaceTemplateInvocation(ii, nil)
				}
			}
		}
		p.code[name] = code
	}
}

func (p *program) stitchTree(name string) {
	// TODO: Save the name of the original template somewhere
	// so we can recover it for error messages. Until we fix
	// that problem we're only stitching trees which are just
	// a WB. In most cases, this will inline the top and bottom
	// hooks, giving already a nice performance boost.
	code := p.code[name]
	for ii := 0; ii < len(code); ii++ {
		v := code[ii]
		if v.op == opTEMPLATE {
			_, t := decodeVal(v.val)
			tmpl := p.strings[t]
			repl := p.code[tmpl]
			if len(repl) == 1 && repl[0].op == opWB {
				// replace the tree
				code = instructions(code).replaceTemplateInvocation(ii, repl)
				compilerDebugf("replaced template %q invocation from %q with WB at PC %d\n", tmpl, name, ii)
				ii--
			}
		}
	}
	p.code[name] = code
}

func (p *program) stitch() {
	p.stitchTree(p.tmpl.root)
}
<commit_msg>Make sure we fail on invalid template invocations while compiling<commit_after>package template

import "fmt"
import "bytes"
import "gnd.la/util/generic"

type instructions []inst

func (i instructions) replace(idx int, count int, repl []inst) []inst {
	// look for jumps before the block which need to be adjusted
	for ii, v := range i[:idx] {
		switch v.op {
		case opJMP, opJMPT, opJMPF, opNEXT:
			val := int(int32(v.val))
			if ii+val > idx {
				i[ii] = inst{v.op, valType(int32(val + len(repl) - count))}
			}
		}
	}
	// look for jumps after the block which need to be adjusted
	start := idx + count
	for ii, v := range i[start:] {
		switch v.op {
		case opJMP, opJMPT, opJMPF, opNEXT:
			val := int(int32(v.val))
			if ii+val < 0 {
				i[ii+start] = inst{v.op, valType(int32(val - len(repl) + count))}
			}
		}
	}
	var ret []inst
	ret = append(ret, i[:idx]...)
	ret = append(ret, repl...)
	ret = append(ret, i[idx+count:]...)
	return ret
}

func (i instructions) replaceTemplateInvocation(idx int, repl []inst) []inst {
	ti := i[idx]
	if ti.op != opTEMPLATE {
		panic(fmt.Errorf("OP at index %d is not opTEMPLATE", idx))
	}
	// Check if we can also remove the pipeline which provides arguments
	// to the template invocation.
	count := 1
	if idx < len(i)-1 {
		pi := i[idx+1]
		if pi.op == opPOP {
			// Remove the POP
			count++
			// This opPOP removes the arguments which were set up for
			// the template call
			if pi.val == 0 {
				// POP until the mark
				for {
					count++
					idx--
					if i[idx].op == opMARK {
						break
					}
				}
			} else {
				// Remove until we find as many pushes as this
				// pop removes
				stack := int(pi.val)
				for stack > 0 {
					count++
					idx--
					ii := i[idx]
					stack -= ii.pushes()
					stack += ii.pops()
				}
			}
		}
	}
	return i.replace(idx, count, repl)
}

// Returns true iff the program references the byte slice determined
// by val
func (p *program) referencesBytes(val valType) bool {
	for _, code := range p.code {
		for _, pi := range code {
			if pi.op == opWB && pi.val == val {
				return true
			}
		}
	}
	return false
}

// removeInternedBytes removes the internet byte slice determined
// by val, adjusting all required instructions.
func (p *program) removeInternedBytes(val valType) {
	idx := int(val)
	p.bs = append(p.bs[:idx], p.bs[idx+1:]...)
	for _, code := range p.code {
		for ii, pi := range code {
			if pi.op == opWB && pi.val >= val {
				ni := pi
				ni.val = pi.val - 1
				code[ii] = ni
			}
		}
	}
}

func (p *program) optimize() {
	p.removeEmptyTemplateInvocations()
	p.stitch()
	p.mergeWriteBytes()
}

// mergeWriteBytes merges adjacent opWB operations into a single WB
func (p *program) mergeWriteBytes() {
	removedReferences := make(map[valType]struct{})
	for name, code := range p.code {
		wb := -1
		ii := 0
		checkWBMerge := func() {
			if wb >= 0 {
				// We had a WB sequence, check its length
				// and merge it if appropriate
				count := ii - wb
				if count > 1 {
					var buf bytes.Buffer
					var refs []valType
					for c := wb; c < ii; c++ {
						wbInst := code[c]
						removedReferences[wbInst.val] = struct{}{}
						refs = append(refs, wbInst.val)
						buf.Write(p.bytesValue(wbInst.val))
					}
					repl := []inst{
						inst{op: opWB, val: p.internBytes(buf.Bytes())},
					}
					code = instructions(code).replace(wb, count, repl)
					compilerDebugf("merged %d WB instructions with byte slice refs %v at PC %d (new ref %d)\n",
						count, refs, wb, repl[0].val)
				}
			}
		}
		for ; ii < len(code); ii++ {
			v := code[ii]
			if v.op == opWB {
				if wb < 0 {
					wb = ii
				}
			} else {
				checkWBMerge()
				wb = -1
			}
		}
		checkWBMerge()
		p.code[name] = code
	}
	// Sort references from higher to lower, so we
	// don't invalidate them while they're being removed
	sortedReferences := make([]valType, 0, len(removedReferences))
	for r := range removedReferences {
		sortedReferences = append(sortedReferences, r)
	}
	generic.SortFunc(sortedReferences, func(a, b valType) bool {
		return a > b
	})
	compilerDebugln("candidates for byte slice removal", sortedReferences)
	for _, r := range sortedReferences {
		if !p.referencesBytes(r) {
			compilerDebugln("will remove byte slice reference", r)
			p.removeInternedBytes(r)
		}
	}
}

// Remove {{ template }} invocations which call into an
// empty template.
func (p *program) removeEmptyTemplateInvocations() {
	for name, code := range p.code {
		for ii := 0; ii < len(code); ii++ {
			v := code[ii]
			if v.op == opTEMPLATE {
				_, t := decodeVal(v.val)
				tmplName := p.strings[t]
				tmplCode, ok := p.code[tmplName]
				if !ok {
					panic(fmt.Errorf("missing template %q", tmplName))
				}
				if len(tmplCode) == 0 {
					// Empty template invocation, remove it
					compilerDebugf("remove empty template invocation %q from %q\n", tmplName, name)
					code = instructions(code).replaceTemplateInvocation(ii, nil)
				}
			}
		}
		p.code[name] = code
	}
}

func (p *program) stitchTree(name string) {
	// TODO: Save the name of the original template somewhere
	// so we can recover it for error messages. Until we fix
	// that problem we're only stitching trees which are just
	// a WB. In most cases, this will inline the top and bottom
	// hooks, giving already a nice performance boost.
	code := p.code[name]
	for ii := 0; ii < len(code); ii++ {
		v := code[ii]
		if v.op == opTEMPLATE {
			_, t := decodeVal(v.val)
			tmpl := p.strings[t]
			repl := p.code[tmpl]
			if len(repl) == 1 && repl[0].op == opWB {
				// replace the tree
				code = instructions(code).replaceTemplateInvocation(ii, repl)
				compilerDebugf("replaced template %q invocation from %q with WB at PC %d\n", tmpl, name, ii)
				ii--
			}
		}
	}
	p.code[name] = code
}

func (p *program) stitch() {
	p.stitchTree(p.tmpl.root)
}
<|endoftext|>


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

//go:build !android && !js && !ppc64le
// +build !android,!js,!ppc64le

// Note: we don't run on Android or ppc64 because if there is any non-race test
// file in this package, the OS tries to link the .syso file into the
// test (even when we're not in race mode), which fails. I'm not sure
// why, but easiest to just punt - as long as a single builder runs
// this test, we're good.

package race

import (
	"bytes"
	"os/exec"
	"path/filepath"
	"runtime"
	"testing"
)

func TestIssue37485(t *testing.T) {
	files, err := filepath.Glob("./*.syso")
	if err != nil {
		t.Fatalf("can't find syso files: %s", err)
	}
	for _, f := range files {
		cmd := exec.Command(filepath.Join(runtime.GOROOT(), "bin", "go"), "tool", "nm", f)
		res, err := cmd.CombinedOutput()
		if err != nil {
			t.Errorf("nm of %s failed: %s", f, err)
			continue
		}
		if bytes.Contains(res, []byte("getauxval")) {
			t.Errorf("%s contains getauxval", f)
		}
	}
}
<commit_msg>runtime/race: use race build tag on syso_test.go<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build race
// +build race

package race

import (
	"bytes"
	"os/exec"
	"path/filepath"
	"runtime"
	"testing"
)

func TestIssue37485(t *testing.T) {
	files, err := filepath.Glob("./*.syso")
	if err != nil {
		t.Fatalf("can't find syso files: %s", err)
	}
	for _, f := range files {
		cmd := exec.Command(filepath.Join(runtime.GOROOT(), "bin", "go"), "tool", "nm", f)
		res, err := cmd.CombinedOutput()
		if err != nil {
			t.Errorf("nm of %s failed: %s", f, err)
			continue
		}
		if bytes.Contains(res, []byte("getauxval")) {
			t.Errorf("%s contains getauxval", f)
		}
	}
}
<|endoftext|>


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

// A minimal client for adding Host records to Datastore for testing. This
// command is ONLY for testing. Host record management by direct access to
// Datastore will not be supported by ePoxy.
package main

import (
	"flag"
	"log"

	"cloud.google.com/go/datastore"
	"github.com/kylelemons/godebug/pretty"
	"github.com/m-lab/epoxy/storage"
	"golang.org/x/net/context"
)

var (
	project  = flag.String("project", "mlab-sandbox", "GCP project ID.")
	hostname = flag.String("hostname", "mlab3.iad1t.measurement-lab.org", "Hostname of new record.")
	address  = flag.String("address", "165.117.240.35", "IP address of hostname.")
	stage1   = flag.String("stage1",
		"https://storage.googleapis.com/epoxy-sandbox/stage1/stage1.ipxe",
		"Absolute URL to a stage1.ipxe script.")
	stage2 = flag.String("stage2",
		"https://storage.googleapis.com/epoxy-sandbox/stage2/stage2.json",
		"Absolute URL to a stage2.json config.")
	stage3 = flag.String("stage3",
		"https://storage.googleapis.com/epoxy-sandbox/stage3/stage3.json",
		"Absolute URL to a stage2.json config.")
)

const usage = `USAGE:
**Only use for testing.**
EXAMPLE:
    create_sample_data --project mlab-sandbox \
        --hostname mlab3.iad1t.measurement-lab.org \
        --address 165.117.240.35 \
        --stage1 https://storage.googleapis.com/epoxy-sandbox/stage1/stage1.ipxe
`

func main() {
	flag.Parse()

	// Print usage unconditionally.
	log.Println(usage)

	// Setup Datastore client.
	ctx := context.Background()
	client, err := datastore.NewClient(ctx, *project)
	if err != nil {
		log.Fatalf("Failed to create new datastore client: %s", err)
	}

	// Save the host record to Datstore.
	ds := storage.NewDatastoreConfig(client)
	h := &storage.Host{
		Name:     *hostname,
		IPv4Addr: *address,
		Boot: storage.Sequence{
			Stage1ChainURL: *stage1,
			Stage2ChainURL: *stage2,
			Stage3ChainURL: *stage3,
		},
	}
	if err = ds.Save(h); err != nil {
		log.Fatalf("%s", err)
	}

	// Retrieve the host record from Datastore to exercise the full save & load path.
	h2, err := ds.Load(h.Name)
	if err != nil {
		log.Fatalf("%s", err)
	}
	pretty.Print(h2)
}
<commit_msg>Fixup flags and usage text<commit_after>// Copyright 2016 ePoxy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//////////////////////////////////////////////////////////////////////////////

// A minimal client for adding Host records to Datastore for testing. This
// command is ONLY for testing. Host record management by direct access to
// Datastore will not be supported by ePoxy.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"cloud.google.com/go/datastore"
	"github.com/kr/pretty"
	"github.com/m-lab/epoxy/storage"
	"golang.org/x/net/context"
)

const usage = `USAGE:
**Only use for testing.**

EXAMPLE:
    epoxy_admin --project mlab-sandbox \
        --hostname mlab3.iad1t.measurement-lab.org \
        --address 165.117.240.35 \
        --stage1 https://storage.googleapis.com/epoxy-sandbox/stage1/stage1.ipxe
        --stage2 https://storage.googleapis.com/epoxy-sandbox/stage2/stage2.ipxe
        --stage3 https://storage.googleapis.com/epoxy-sandbox/stage3/stage3.ipxe
`

var (
	fProject  string
	fHostname string
	fAddress  string
	fStage1   string
	fStage2   string
	fStage3   string
)

func init() {
	// Add an alternate usage message.
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, usage)
		flag.PrintDefaults()
	}
	flag.StringVar(&fProject, "project", "mlab-sandbox", "GCP project ID.")
	flag.StringVar(&fHostname, "hostname", "mlab3.iad1t.measurement-lab.org", "Hostname of new record.")
	flag.StringVar(&fAddress, "address", "165.117.240.35", "IP address of hostname.")
	flag.StringVar(&fStage1, "stage1",
		"https://storage.googleapis.com/epoxy-sandbox/stage1/stage1.ipxe",
		"Absolute URL to a stage1.ipxe script.")
	flag.StringVar(&fStage2, "stage2",
		"https://storage.googleapis.com/epoxy-sandbox/stage2/stage2.json",
		"Absolute URL to a stage2.json config.")
	flag.StringVar(&fStage3, "stage3",
		"https://storage.googleapis.com/epoxy-sandbox/stage3/stage3.json",
		"Absolute URL to a stage2.json config.")
}

func main() {
	flag.Parse()

	// Setup Datastore client.
	ctx := context.Background()
	client, err := datastore.NewClient(ctx, fProject)
	if err != nil {
		log.Fatalf("Failed to create new datastore client: %s", err)
	}

	// Save the host record to Datstore.
	ds := storage.NewDatastoreConfig(client)
	h := &storage.Host{
		Name:     fHostname,
		IPv4Addr: fAddress,
		Boot: storage.Sequence{
			Stage1ChainURL: fStage1,
			Stage2ChainURL: fStage2,
			Stage3ChainURL: fStage3,
		},
	}
	if err = ds.Save(h); err != nil {
		log.Fatalf("%s", err)
	}

	// Retrieve the host record from Datastore to exercise the full save & load path.
	h2, err := ds.Load(h.Name)
	if err != nil {
		log.Fatalf("%s", err)
	}
	pretty.Print(h2)
}
<|endoftext|>


<commit_before>// chris 073015

// extractebnf extracts and prints EBNF grammars.
//
// In particular, it extracts grammars from HTML documents, as specified
// by golang.org/x/exp/ebnflint, which includes the Go Programming
// Language Specification HTML page.
package main

import (
	"bytes"
	"flag"
	"io"
	"log"
	"os"
	"path"

	"io/ioutil"

	"chrispennello.com/go/rebnf"
)

var args struct {
	progName string
	strip    bool
}

func init() {
	log.SetFlags(0)
	strip := flag.Bool("strip", false, "strip superfluous newlines")
	flag.Parse()
	args.progName = path.Base(os.Args[0])
	args.strip = *strip
}

func usage() {
	log.Printf("usage: %s [grammarfile]\n", args.progName)
	os.Exit(2)
}

// stripNewlines consolidates contiguous newlines and returns a new
// slice of bytes with duplicated newlines omitted.
func stripNewlines(x []byte) []byte {
	buf := new(bytes.Buffer)
	last := byte('\n')
	for _, b := range x {
		if last != '\n' || b != '\n' {
			buf.WriteByte(b)
		}
		last = b
	}
	return buf.Bytes()
}

func main() {
	var (
		filename string
		r        io.Reader
	)

	if flag.NArg() == 0 {
		filename, r = "-", os.Stdin
	} else if flag.NArg() == 1 {
		filename = flag.Args()[0]
		file, err := os.Open(filename)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()
		r = file
	} else {
		usage()
	}

	src, err := ioutil.ReadAll(r)
	if err != nil {
		log.Fatal(err)
	}

	src = rebnf.CheckExtract(filename, src)

	if args.strip {
		src = stripNewlines(src)
	}

	_, err = io.Copy(os.Stdout, bytes.NewBuffer(src))
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Adds option documentation to extractebnf.<commit_after>// chris 073015

// extractebnf extracts and prints EBNF grammars.
//
// In particular, it extracts grammars from HTML documents, as specified
// by golang.org/x/exp/ebnflint, which includes the Go Programming
// Language Specification HTML page.
//
// Options are:
//
//	-strip=false: strip superfluous newlines
//
package main

import (
	"bytes"
	"flag"
	"io"
	"log"
	"os"
	"path"

	"io/ioutil"

	"chrispennello.com/go/rebnf"
)

var args struct {
	progName string
	strip    bool
}

func init() {
	log.SetFlags(0)
	strip := flag.Bool("strip", false, "strip superfluous newlines")
	flag.Parse()
	args.progName = path.Base(os.Args[0])
	args.strip = *strip
}

func usage() {
	log.Printf("usage: %s [grammarfile]\n", args.progName)
	os.Exit(2)
}

// stripNewlines consolidates contiguous newlines and returns a new
// slice of bytes with duplicated newlines omitted.
func stripNewlines(x []byte) []byte {
	buf := new(bytes.Buffer)
	last := byte('\n')
	for _, b := range x {
		if last != '\n' || b != '\n' {
			buf.WriteByte(b)
		}
		last = b
	}
	return buf.Bytes()
}

func main() {
	var (
		filename string
		r        io.Reader
	)

	if flag.NArg() == 0 {
		filename, r = "-", os.Stdin
	} else if flag.NArg() == 1 {
		filename = flag.Args()[0]
		file, err := os.Open(filename)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()
		r = file
	} else {
		usage()
	}

	src, err := ioutil.ReadAll(r)
	if err != nil {
		log.Fatal(err)
	}

	src = rebnf.CheckExtract(filename, src)

	if args.strip {
		src = stripNewlines(src)
	}

	_, err = io.Copy(os.Stdout, bytes.NewBuffer(src))
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"go/build"
	"go/format"
	"gondola/admin"
	"gondola/mux"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

func TemplatesMap(ctx *mux.Context) {
	var dir string
	var name string
	extensions := map[string]struct{}{
		".html": struct{}{},
	}
	ctx.ParseParamValue("dir", &dir)
	if dir == "" {
		fmt.Fprintf(os.Stderr, "dir can't be empty")
		return
	}
	ctx.ParseParamValue("name", &name)
	if name == "" {
		fmt.Fprintf(os.Stderr, "name can't be empty")
		return
	}
	var exts string
	ctx.ParseParamValue("extensions", &exts)
	if exts != "" {
		for _, v := range strings.Split(exts, ",") {
			e := strings.ToLower(strings.TrimSpace(v))
			if e != "" {
				if e[0] != '.' {
					e = "." + e
				}
				extensions[e] = struct{}{}
			}
		}
	}
	var out string
	ctx.ParseParamValue("out", &out)
	var buf bytes.Buffer
	if out != "" {
		// Try to guess package name. Do it before writing the file, otherwise the package becomes invalid.
		odir := filepath.Dir(out)
		p, err := build.ImportDir(odir, 0)
		if err == nil {
			buf.WriteString(fmt.Sprintf("package %s\n", p.Name))
		}
	}
	buf.WriteString("import \"gondola/loaders\"\n")
	buf.WriteString(fmt.Sprintf("// AUTOMATICALLY GENERATED WITH %s. DO NOT EDIT!\n", strings.Join(os.Args, " ")))
	buf.WriteString(fmt.Sprintf("var %s = loaders.MapLoader(map[string][]byte{\n", name))
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.Mode().IsRegular() {
			if _, ok := extensions[strings.ToLower(filepath.Ext(path))]; ok {
				contents, err := ioutil.ReadFile(path)
				if err != nil {
					return fmt.Errorf("error reading %s: %s", path, err)
				}
				rel := path[len(dir):]
				if rel[0] == '/' {
					rel = rel[1:]
				}
				buf.WriteString(fmt.Sprintf("%q", rel))
				buf.WriteByte(':')
				buf.WriteString(" []byte{")
				for ii, v := range contents {
					buf.WriteString(fmt.Sprintf("0x%02X", v))
					buf.WriteByte(',')
					if ii%8 == 0 {
						buf.WriteByte('\n')
					}
				}
				buf.Truncate(buf.Len() - 1)
				buf.WriteString("},\n")
			}
		}
		return nil
	})
	if err != nil {
		panic(err)
	}
	buf.WriteString("})\n")
	b, err := format.Source(buf.Bytes())
	if err != nil {
		panic(err)
	}
	var w io.Writer
	if out != "" {
		flags := os.O_CREATE | os.O_WRONLY
		var force bool
		ctx.ParseParamValue("f", &force)
		if !force {
			flags |= os.O_EXCL
		}
		f, err := os.OpenFile(out, flags, 0644)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error creating output file %q: %s\n", out, err)
			return
		}
		defer f.Close()
		w = f
	} else {
		w = os.Stdout
	}
	_, err = w.Write(b)
	if err != nil {
		panic(err)
	}
}

func init() {
	admin.Register(TemplatesMap, &admin.Options{
		Help: "Converts all templates in <dir> into Go code and generates a MapLoader named with <name>",
		Flags: admin.Flags(
			admin.StringFlag("dir", "", "Directory with the html templates"),
			admin.StringFlag("name", "", "Name of the generated MapLoader"),
			admin.StringFlag("out", "", "Output filename. If empty, output is printed to standard output"),
			admin.BoolFlag("f", false, "When creating output file, overwrite any existing file with the same name"),
			admin.StringFlag("extensions", "", "Additional extensions (besides .html) to include, separated by commas"),
		),
	})
}
<commit_msg>Truncate the file when overwriting it<commit_after>package main

import (
	"bytes"
	"fmt"
	"go/build"
	"go/format"
	"gondola/admin"
	"gondola/mux"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

func TemplatesMap(ctx *mux.Context) {
	var dir string
	var name string
	extensions := map[string]struct{}{
		".html": struct{}{},
	}
	ctx.ParseParamValue("dir", &dir)
	if dir == "" {
		fmt.Fprintf(os.Stderr, "dir can't be empty")
		return
	}
	ctx.ParseParamValue("name", &name)
	if name == "" {
		fmt.Fprintf(os.Stderr, "name can't be empty")
		return
	}
	var exts string
	ctx.ParseParamValue("extensions", &exts)
	if exts != "" {
		for _, v := range strings.Split(exts, ",") {
			e := strings.ToLower(strings.TrimSpace(v))
			if e != "" {
				if e[0] != '.' {
					e = "." + e
				}
				extensions[e] = struct{}{}
			}
		}
	}
	var out string
	ctx.ParseParamValue("out", &out)
	var buf bytes.Buffer
	if out != "" {
		// Try to guess package name. Do it before writing the file, otherwise the package becomes invalid.
		odir := filepath.Dir(out)
		p, err := build.ImportDir(odir, 0)
		if err == nil {
			buf.WriteString(fmt.Sprintf("package %s\n", p.Name))
		}
	}
	buf.WriteString("import \"gondola/loaders\"\n")
	buf.WriteString(fmt.Sprintf("// AUTOMATICALLY GENERATED WITH %s. DO NOT EDIT!\n", strings.Join(os.Args, " ")))
	buf.WriteString(fmt.Sprintf("var %s = loaders.MapLoader(map[string][]byte{\n", name))
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.Mode().IsRegular() {
			if _, ok := extensions[strings.ToLower(filepath.Ext(path))]; ok {
				contents, err := ioutil.ReadFile(path)
				if err != nil {
					return fmt.Errorf("error reading %s: %s", path, err)
				}
				rel := path[len(dir):]
				if rel[0] == '/' {
					rel = rel[1:]
				}
				buf.WriteString(fmt.Sprintf("%q", rel))
				buf.WriteByte(':')
				buf.WriteString(" []byte{")
				for ii, v := range contents {
					buf.WriteString(fmt.Sprintf("0x%02X", v))
					buf.WriteByte(',')
					if ii%8 == 0 {
						buf.WriteByte('\n')
					}
				}
				buf.Truncate(buf.Len() - 1)
				buf.WriteString("},\n")
			}
		}
		return nil
	})
	if err != nil {
		panic(err)
	}
	buf.WriteString("})\n")
	b, err := format.Source(buf.Bytes())
	if err != nil {
		panic(err)
	}
	var w io.Writer
	if out != "" {
		flags := os.O_CREATE | os.O_WRONLY | os.O_TRUNC
		var force bool
		ctx.ParseParamValue("f", &force)
		if !force {
			flags |= os.O_EXCL
		}
		f, err := os.OpenFile(out, flags, 0644)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error creating output file %q: %s\n", out, err)
			return
		}
		defer f.Close()
		w = f
	} else {
		w = os.Stdout
	}
	_, err = w.Write(b)
	if err != nil {
		panic(err)
	}
}

func init() {
	admin.Register(TemplatesMap, &admin.Options{
		Help: "Converts all templates in <dir> into Go code and generates a MapLoader named with <name>",
		Flags: admin.Flags(
			admin.StringFlag("dir", "", "Directory with the html templates"),
			admin.StringFlag("name", "", "Name of the generated MapLoader"),
			admin.StringFlag("out", "", "Output filename. If empty, output is printed to standard output"),
			admin.BoolFlag("f", false, "When creating output file, overwrite any existing file with the same name"),
			admin.StringFlag("extensions", "", "Additional extensions (besides .html) to include, separated by commas"),
		),
	})
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"

	"launchpad.net/gnuflag"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/names"
)

const getConstraintsDoc = `
get-constraints returns a list of constraints that have been set on
the environment using juju set-constraints.  You can also view constraints set
for a specific service by using juju get-constraints <service>.

See Also:
   juju help constraints
   juju help set-constraints
`

const setConstraintsDoc = `
set-constraints sets machine constraints on the system, which are used as the
default constraints for all new machines provisioned in the environment (unless
overridden).  You can also set constraints on a specific service by using juju
set-constraints <service>.

Constraints set on a service are combined with environment constraints for
commands (such as juju deploy) that provision machines for services.  Where
environment and service constraints overlap, the service constraints take
precedence.

Examples:

   set-constraints mem=8G                         (all new machines in the environment must have at least 8GB of RAM)
   set-constraints --service wordpress mem=4G     (all new wordpress machines can ignore the 8G constraint above, and require only 4G)

See Also:
   juju help constraints
   juju help get-constraints
   juju help deploy
   juju help add-machine
   juju help add-unit
`

// GetConstraintsCommand shows the constraints for a service or environment.
type GetConstraintsCommand struct {
	cmd.EnvCommandBase
	ServiceName string
	out         cmd.Output
}

func (c *GetConstraintsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "get-constraints",
		Args:    "[<service>]",
		Purpose: "view constraints on the environment or a service",
		Doc:     getConstraintsDoc,
	}
}

func formatConstraints(value interface{}) ([]byte, error) {
	return []byte(value.(constraints.Value).String()), nil
}

func (c *GetConstraintsCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	c.out.AddFlags(f, "constraints", map[string]cmd.Formatter{
		"constraints": formatConstraints,
		"yaml":        cmd.FormatYaml,
		"json":        cmd.FormatJson,
	})
}

func (c *GetConstraintsCommand) Init(args []string) error {
	if len(args) > 0 {
		if !names.IsService(args[0]) {
			return fmt.Errorf("invalid service name %q", args[0])
		}
		c.ServiceName, args = args[0], args[1:]
	}
	return cmd.CheckEmpty(args)
}

func (c *GetConstraintsCommand) Run(ctx *cmd.Context) error {
	apiclient, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer apiclient.Close()

	var cons constraints.Value
	if c.ServiceName == "" {
		cons, err = apiclient.GetEnvironmentConstraints()
	} else {
		cons, err = apiclient.GetServiceConstraints(c.ServiceName)
	}
	if err != nil {
		return err
	}
	return c.out.Write(ctx, cons)
}

// SetConstraintsCommand shows the constraints for a service or environment.
type SetConstraintsCommand struct {
	cmd.EnvCommandBase
	ServiceName string
	Constraints constraints.Value
}

func (c *SetConstraintsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "set-constraints",
		Args:    "[key=[value] ...]",
		Purpose: "set constraints on the environment or a service",
		Doc:     setConstraintsDoc,
	}
}

func (c *SetConstraintsCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.StringVar(&c.ServiceName, "s", "", "set service constraints")
	f.StringVar(&c.ServiceName, "service", "", "")
}

func (c *SetConstraintsCommand) Init(args []string) (err error) {
	if c.ServiceName != "" && !names.IsService(c.ServiceName) {
		return fmt.Errorf("invalid service name %q", c.ServiceName)
	}
	c.Constraints, err = constraints.Parse(args...)
	return err
}

func (c *SetConstraintsCommand) Run(_ *cmd.Context) (err error) {
	apiclient, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer apiclient.Close()
	if c.ServiceName == "" {
		return apiclient.SetEnvironmentConstraints(c.Constraints)
	}
	return apiclient.SetServiceConstraints(c.ServiceName, c.Constraints)
}
<commit_msg>[r=jameinel],[bug=1253633] cmd/juju/constraints.go: restore 1.16 compat<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"

	"launchpad.net/gnuflag"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/names"
	"launchpad.net/juju-core/rpc"
)

const getConstraintsDoc = `
get-constraints returns a list of constraints that have been set on
the environment using juju set-constraints.  You can also view constraints set
for a specific service by using juju get-constraints <service>.

See Also:
   juju help constraints
   juju help set-constraints
`

const setConstraintsDoc = `
set-constraints sets machine constraints on the system, which are used as the
default constraints for all new machines provisioned in the environment (unless
overridden).  You can also set constraints on a specific service by using juju
set-constraints <service>.

Constraints set on a service are combined with environment constraints for
commands (such as juju deploy) that provision machines for services.  Where
environment and service constraints overlap, the service constraints take
precedence.

Examples:

   set-constraints mem=8G                         (all new machines in the environment must have at least 8GB of RAM)
   set-constraints --service wordpress mem=4G     (all new wordpress machines can ignore the 8G constraint above, and require only 4G)

See Also:
   juju help constraints
   juju help get-constraints
   juju help deploy
   juju help add-machine
   juju help add-unit
`

// GetConstraintsCommand shows the constraints for a service or environment.
type GetConstraintsCommand struct {
	cmd.EnvCommandBase
	ServiceName string
	out         cmd.Output
}

func (c *GetConstraintsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "get-constraints",
		Args:    "[<service>]",
		Purpose: "view constraints on the environment or a service",
		Doc:     getConstraintsDoc,
	}
}

func formatConstraints(value interface{}) ([]byte, error) {
	return []byte(value.(constraints.Value).String()), nil
}

func (c *GetConstraintsCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	c.out.AddFlags(f, "constraints", map[string]cmd.Formatter{
		"constraints": formatConstraints,
		"yaml":        cmd.FormatYaml,
		"json":        cmd.FormatJson,
	})
}

func (c *GetConstraintsCommand) Init(args []string) error {
	if len(args) > 0 {
		if !names.IsService(args[0]) {
			return fmt.Errorf("invalid service name %q", args[0])
		}
		c.ServiceName, args = args[0], args[1:]
	}
	return cmd.CheckEmpty(args)
}

// getEnvironConstraints1dot16 uses direct DB access to get the Environment
// constraints against an API server running 1.16 or older (when GetEnvironmentConstraints
// was not available). This fallback can be removed when we no longer maintain
// 1.16 compatibility.
// This only does the GetEnvironmentConstraints portion of Run, since
// GetServiceConstraints was already implemented.
func (c *GetConstraintsCommand) getEnvironConstraints1dot16() (constraints.Value, error) {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return constraints.Value{}, err
	}
	defer conn.Close()
	return conn.State.EnvironConstraints()
}

func (c *GetConstraintsCommand) Run(ctx *cmd.Context) error {
	apiclient, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer apiclient.Close()

	var cons constraints.Value
	if c.ServiceName == "" {
		cons, err = apiclient.GetEnvironmentConstraints()
		if rpc.IsNoSuchRequest(err) {
			logger.Infof("GetEnvironmentConstraints not supported by the API server, " +
				"falling back to 1.16 compatibility mode (direct DB access)")
			cons, err = c.getEnvironConstraints1dot16()
		}
	} else {
		cons, err = apiclient.GetServiceConstraints(c.ServiceName)
	}
	if err != nil {
		return err
	}
	return c.out.Write(ctx, cons)
}

// SetConstraintsCommand shows the constraints for a service or environment.
type SetConstraintsCommand struct {
	cmd.EnvCommandBase
	ServiceName string
	Constraints constraints.Value
}

func (c *SetConstraintsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "set-constraints",
		Args:    "[key=[value] ...]",
		Purpose: "set constraints on the environment or a service",
		Doc:     setConstraintsDoc,
	}
}

func (c *SetConstraintsCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.StringVar(&c.ServiceName, "s", "", "set service constraints")
	f.StringVar(&c.ServiceName, "service", "", "")
}

func (c *SetConstraintsCommand) Init(args []string) (err error) {
	if c.ServiceName != "" && !names.IsService(c.ServiceName) {
		return fmt.Errorf("invalid service name %q", c.ServiceName)
	}
	c.Constraints, err = constraints.Parse(args...)
	return err
}

// setEnvironConstraints1dot16 uses direct DB access to get the Environment
// constraints against an API server running 1.16 or older (when SetEnvironmentConstraints
// was not available). This fallback can be removed when we no longer maintain
// 1.16 compatibility.
// This only does the SetEnvironmentConstraints portion of Run, since
// SetServiceConstraints was already implemented.
func (c *SetConstraintsCommand) setEnvironConstraints1dot16() error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.State.SetEnvironConstraints(c.Constraints)
}

func (c *SetConstraintsCommand) Run(_ *cmd.Context) (err error) {
	apiclient, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer apiclient.Close()
	if c.ServiceName == "" {
		err = apiclient.SetEnvironmentConstraints(c.Constraints)
		if rpc.IsNoSuchRequest(err) {
			logger.Infof("SetEnvironmentConstraints not supported by the API server, " +
				"falling back to 1.16 compatibility mode (direct DB access)")
			err = c.setEnvironConstraints1dot16()
		}
		return err
	}
	return apiclient.SetServiceConstraints(c.ServiceName, c.Constraints)
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	stderrors "errors"
	"fmt"

	"launchpad.net/gnuflag"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/storage"
	"launchpad.net/juju-core/environs/sync"
	envtools "launchpad.net/juju-core/environs/tools"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state/api/params"
	coretools "launchpad.net/juju-core/tools"
	"launchpad.net/juju-core/version"
)

// UpgradeJujuCommand upgrades the agents in a juju installation.
type UpgradeJujuCommand struct {
	cmd.EnvCommandBase
	vers        string
	Version     version.Number
	UploadTools bool
	Series      []string
}

var upgradeJujuDoc = `
The upgrade-juju command upgrades a running environment by setting a version
number for all juju agents to run. By default, it chooses the most recent
supported version compatible with the command-line tools version.

A development version is defined to be any version with an odd minor
version or a nonzero build component (for example version 2.1.1, 3.3.0
and 2.0.0.1 are development versions; 2.0.3 and 3.4.1 are not). A
development version may be chosen in two cases:

 - when the current agent version is a development one and there is
   a more recent version available with the same major.minor numbers;
 - when an explicit --version major.minor is given (e.g. --version 1.17,
   or 1.17.2, but not just 1)

For development use, the --upload-tools flag specifies that the juju tools will
packaged (or compiled locally, if no jujud binaries exists, for which you will
need the golang packages installed) and uploaded before the version is set.
Currently the tools will be uploaded as if they had the version of the current
juju tool, unless specified otherwise by the --version flag.

When run without arguments. upgrade-juju will try to upgrade to the
following versions, in order of preference, depending on the current
value of the environment's agent-version setting:

 - The highest patch.build version of the *next* stable major.minor version.
 - The highest patch.build version of the *current* major.minor version.

Both of these depend on tools availability, which some situations (no
outgoing internet access) and provider types (such as maas) require that
you manage yourself; see the documentation for "sync-tools".
`

func (c *UpgradeJujuCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "upgrade-juju",
		Purpose: "upgrade the tools in a juju environment",
		Doc:     upgradeJujuDoc,
	}
}

func (c *UpgradeJujuCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.StringVar(&c.vers, "version", "", "upgrade to specific version")
	f.BoolVar(&c.UploadTools, "upload-tools", false, "upload local version of tools")
	f.Var(seriesVar{&c.Series}, "series", "upload tools for supplied comma-separated series list")
}

func (c *UpgradeJujuCommand) Init(args []string) error {
	if c.vers != "" {
		vers, err := version.Parse(c.vers)
		if err != nil {
			return err
		}
		if vers.Major != version.Current.Major {
			return fmt.Errorf("cannot upgrade to version incompatible with CLI")
		}
		if c.UploadTools && vers.Build != 0 {
			// TODO(fwereade): when we start taking versions from actual built
			// code, we should disable --version when used with --upload-tools.
			// For now, it's the only way to experiment with version upgrade
			// behaviour live, so the only restriction is that Build cannot
			// be used (because its value needs to be chosen internally so as
			// not to collide with existing tools).
			return fmt.Errorf("cannot specify build number when uploading tools")
		}
		c.Version = vers
	}
	if len(c.Series) > 0 && !c.UploadTools {
		return fmt.Errorf("--series requires --upload-tools")
	}
	return cmd.CheckEmpty(args)
}

var errUpToDate = stderrors.New("no upgrades available")

// Run changes the version proposed for the juju envtools.
func (c *UpgradeJujuCommand) Run(_ *cmd.Context) (err error) {
	client, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer client.Close()
	defer func() {
		if err == errUpToDate {
			log.Noticef(err.Error())
			err = nil
		}
	}()

	// Determine the version to upgrade to, uploading tools if necessary.
	attrs, err := client.EnvironmentGet()
	if params.IsCodeNotImplemented(err) {
		return c.run1dot16()
	}
	if err != nil {
		return err
	}
	cfg, err := config.New(config.NoDefaults, attrs)
	if err != nil {
		return err
	}
	env, err := environs.New(cfg)
	if err != nil {
		return err
	}
	v, err := c.initVersions(cfg, env)
	if err != nil {
		return err
	}
	if c.UploadTools {
		series := getUploadSeries(cfg, c.Series)
		if err := v.uploadTools(env.Storage(), series); err != nil {
			return err
		}
	}
	if err := v.validate(); err != nil {
		return err
	}
	log.Infof("upgrade version chosen: %s", v.chosen)
	// TODO(fwereade): this list may be incomplete, pending envtools.Upload change.
	log.Infof("available tools: %s", v.tools)

	if err := client.SetEnvironAgentVersion(v.chosen); err != nil {
		return err
	}
	log.Noticef("started upgrade to %s", v.chosen)
	return nil
}

// run1dot16 implements the command without access to the API. This is
// needed for compatibility, so 1.16 can be upgraded to newer
// releases. It should be removed in 1.18.
func (c *UpgradeJujuCommand) run1dot16() error {
	log.Warningf("running in 1.16 compatibility mode")
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	defer func() {
		if err == errUpToDate {
			log.Noticef(err.Error())
			err = nil
		}
	}()

	// Determine the version to upgrade to, uploading tools if necessary.
	env := conn.Environ
	cfg, err := conn.State.EnvironConfig()
	if err != nil {
		return err
	}
	v, err := c.initVersions(cfg, env)
	if err != nil {
		return err
	}
	if c.UploadTools {
		series := getUploadSeries(cfg, c.Series)
		if err := v.uploadTools(env.Storage(), series); err != nil {
			return err
		}
	}
	if err := v.validate(); err != nil {
		return err
	}
	log.Infof("upgrade version chosen: %s", v.chosen)
	// TODO(fwereade): this list may be incomplete, pending envtools.Upload change.
	log.Infof("available tools: %s", v.tools)

	if err := conn.State.SetEnvironAgentVersion(v.chosen); err != nil {
		return err
	}
	log.Noticef("started upgrade to %s", v.chosen)
	return nil
}

// initVersions collects state relevant to an upgrade decision. The returned
// agent and client versions, and the list of currently available tools, will
// always be accurate; the chosen version, and the flag indicating development
// mode, may remain blank until uploadTools or validate is called.
func (c *UpgradeJujuCommand) initVersions(cfg *config.Config, env environs.Environ) (*upgradeVersions, error) {
	agent, ok := cfg.AgentVersion()
	if !ok {
		// Can't happen. In theory.
		return nil, fmt.Errorf("incomplete environment configuration")
	}
	if c.Version == agent {
		return nil, errUpToDate
	}
	client := version.Current.Number
	// TODO use an API call rather than requiring the environment,
	// so that we can restrict access to the provider secrets
	// while still allowing users to upgrade.
	available, err := envtools.FindTools(env, client.Major, -1, coretools.Filter{}, envtools.DoNotAllowRetry)
	if err != nil {
		if !errors.IsNotFoundError(err) {
			return nil, err
		}
		if !c.UploadTools {
			// No tools found and we shouldn't upload any, so if we are not asking for a
			// major upgrade, pretend there is no more recent version available.
			if c.Version == version.Zero && agent.Major == client.Major {
				return nil, errUpToDate
			}
			return nil, err
		}
	}
	return &upgradeVersions{
		agent:  agent,
		client: client,
		chosen: c.Version,
		tools:  available,
	}, nil
}

// upgradeVersions holds the version information for making upgrade decisions.
type upgradeVersions struct {
	agent  version.Number
	client version.Number
	chosen version.Number
	tools  coretools.List
}

// uploadTools compiles jujud from $GOPATH and uploads it into the supplied
// storage. If no version has been explicitly chosen, the version number
// reported by the built tools will be based on the client version number.
// In any case, the version number reported will have a build component higher
// than that of any otherwise-matching available envtools.
// uploadTools resets the chosen version and replaces the available tools
// with the ones just uploaded.
func (v *upgradeVersions) uploadTools(storage storage.Storage, series []string) error {
	// TODO(fwereade): this is kinda crack: we should not assume that
	// version.Current matches whatever source happens to be built. The
	// ideal would be:
	//  1) compile jujud from $GOPATH into some build dir
	//  2) get actual version with `jujud version`
	//  3) check actual version for compatibility with CLI tools
	//  4) generate unique build version with reference to available tools
	//  5) force-version that unique version into the dir directly
	//  6) archive and upload the build dir
	// ...but there's no way we have time for that now. In the meantime,
	// considering the use cases, this should work well enough; but it
	// won't detect an incompatible major-version change, which is a shame.
	if v.chosen == version.Zero {
		v.chosen = v.client
	}
	v.chosen = uploadVersion(v.chosen, v.tools)

	// TODO(fwereade): envtools.Upload should return envtools.List, and should
	// include all the extra series we build, so we can set *that* onto
	// v.available and maybe one day be able to check that a given upgrade
	// won't leave out-of-date machines lying around, starved of envtools.
	uploaded, err := sync.Upload(storage, &v.chosen, series...)
	if err != nil {
		return err
	}
	v.tools = coretools.List{uploaded}
	return nil
}

// validate chooses an upgrade version, if one has not already been chosen,
// and ensures the tools list contains no entries that do not have that version.
// If validate returns no error, the environment agent-version can be set to
// the value of the chosen field.
func (v *upgradeVersions) validate() (err error) {
	if v.chosen == version.Zero {
		// No explicitly specified version, so find the version to which we
		// need to upgrade. If the CLI and agent major versions match, we find
		// next available stable release to upgrade to by incrementing the
		// minor version, starting from the current agent version and doing
		// major.minor+1 or +2 as needed. If the CLI has a greater major version,
		// we just use the CLI version as is.
		nextVersion := v.agent
		if nextVersion.Major == v.client.Major {
			if v.agent.IsDev() {
				nextVersion.Minor += 1
			} else {
				nextVersion.Minor += 2
			}
		} else {
			nextVersion = v.client
		}

		newestNextStable, found := v.tools.NewestCompatible(nextVersion)
		if found {
			log.Debugf("found a more recent stable version %s", newestNextStable)
			v.chosen = newestNextStable
		} else {
			newestCurrent, found := v.tools.NewestCompatible(v.agent)
			if found {
				log.Debugf("found more recent current version %s", newestCurrent)
				v.chosen = newestCurrent
			} else {
				if v.agent.Major != v.client.Major {
					return fmt.Errorf("no compatible tools available")
				} else {
					return fmt.Errorf("no more recent supported versions available")
				}
			}
		}
	} else {
		// If not completely specified already, pick a single tools version.
		filter := coretools.Filter{Number: v.chosen, Released: !v.chosen.IsDev()}
		if v.tools, err = v.tools.Match(filter); err != nil {
			return err
		}
		v.chosen, v.tools = v.tools.Newest()
	}
	if v.chosen == v.agent {
		return errUpToDate
	}

	// Disallow major.minor version downgrades.
	if v.chosen.Major < v.agent.Major || v.chosen.Major == v.agent.Major && v.chosen.Minor < v.agent.Minor {
		// TODO(fwereade): I'm a bit concerned about old agent/CLI tools even
		// *connecting* to environments with higher agent-versions; but ofc they
		// have to connect in order to discover they shouldn't. However, once
		// any of our tools detect an incompatible version, they should act to
		// minimize damage: the CLI should abort politely, and the agents should
		// run an Upgrader but no other tasks.
		return fmt.Errorf("cannot change version from %d.%d to %d.%d",
			v.agent.Major, v.agent.Minor, v.chosen.Major, v.chosen.Minor)
	}

	return nil
}

// uploadVersion returns a copy of the supplied version with a build number
// higher than any of the supplied tools that share its major, minor and patch.
func uploadVersion(vers version.Number, existing coretools.List) version.Number {
	vers.Build++
	for _, t := range existing {
		if t.Version.Major != vers.Major || t.Version.Minor != vers.Minor || t.Version.Patch != vers.Patch {
			continue
		}
		if t.Version.Build >= vers.Build {
			vers.Build = t.Version.Build + 1
		}
	}
	return vers
}
<commit_msg>Tweak error message<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	stderrors "errors"
	"fmt"

	"launchpad.net/gnuflag"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/storage"
	"launchpad.net/juju-core/environs/sync"
	envtools "launchpad.net/juju-core/environs/tools"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state/api/params"
	coretools "launchpad.net/juju-core/tools"
	"launchpad.net/juju-core/version"
)

// UpgradeJujuCommand upgrades the agents in a juju installation.
type UpgradeJujuCommand struct {
	cmd.EnvCommandBase
	vers        string
	Version     version.Number
	UploadTools bool
	Series      []string
}

var upgradeJujuDoc = `
The upgrade-juju command upgrades a running environment by setting a version
number for all juju agents to run. By default, it chooses the most recent
supported version compatible with the command-line tools version.

A development version is defined to be any version with an odd minor
version or a nonzero build component (for example version 2.1.1, 3.3.0
and 2.0.0.1 are development versions; 2.0.3 and 3.4.1 are not). A
development version may be chosen in two cases:

 - when the current agent version is a development one and there is
   a more recent version available with the same major.minor numbers;
 - when an explicit --version major.minor is given (e.g. --version 1.17,
   or 1.17.2, but not just 1)

For development use, the --upload-tools flag specifies that the juju tools will
packaged (or compiled locally, if no jujud binaries exists, for which you will
need the golang packages installed) and uploaded before the version is set.
Currently the tools will be uploaded as if they had the version of the current
juju tool, unless specified otherwise by the --version flag.

When run without arguments. upgrade-juju will try to upgrade to the
following versions, in order of preference, depending on the current
value of the environment's agent-version setting:

 - The highest patch.build version of the *next* stable major.minor version.
 - The highest patch.build version of the *current* major.minor version.

Both of these depend on tools availability, which some situations (no
outgoing internet access) and provider types (such as maas) require that
you manage yourself; see the documentation for "sync-tools".
`

func (c *UpgradeJujuCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "upgrade-juju",
		Purpose: "upgrade the tools in a juju environment",
		Doc:     upgradeJujuDoc,
	}
}

func (c *UpgradeJujuCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.StringVar(&c.vers, "version", "", "upgrade to specific version")
	f.BoolVar(&c.UploadTools, "upload-tools", false, "upload local version of tools")
	f.Var(seriesVar{&c.Series}, "series", "upload tools for supplied comma-separated series list")
}

func (c *UpgradeJujuCommand) Init(args []string) error {
	if c.vers != "" {
		vers, err := version.Parse(c.vers)
		if err != nil {
			return err
		}
		if vers.Major != version.Current.Major {
			return fmt.Errorf("cannot upgrade to version incompatible with CLI")
		}
		if c.UploadTools && vers.Build != 0 {
			// TODO(fwereade): when we start taking versions from actual built
			// code, we should disable --version when used with --upload-tools.
			// For now, it's the only way to experiment with version upgrade
			// behaviour live, so the only restriction is that Build cannot
			// be used (because its value needs to be chosen internally so as
			// not to collide with existing tools).
			return fmt.Errorf("cannot specify build number when uploading tools")
		}
		c.Version = vers
	}
	if len(c.Series) > 0 && !c.UploadTools {
		return fmt.Errorf("--series requires --upload-tools")
	}
	return cmd.CheckEmpty(args)
}

var errUpToDate = stderrors.New("no upgrades available")

// Run changes the version proposed for the juju envtools.
func (c *UpgradeJujuCommand) Run(_ *cmd.Context) (err error) {
	client, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer client.Close()
	defer func() {
		if err == errUpToDate {
			log.Noticef(err.Error())
			err = nil
		}
	}()

	// Determine the version to upgrade to, uploading tools if necessary.
	attrs, err := client.EnvironmentGet()
	if params.IsCodeNotImplemented(err) {
		return c.run1dot16()
	}
	if err != nil {
		return err
	}
	cfg, err := config.New(config.NoDefaults, attrs)
	if err != nil {
		return err
	}
	env, err := environs.New(cfg)
	if err != nil {
		return err
	}
	v, err := c.initVersions(cfg, env)
	if err != nil {
		return err
	}
	if c.UploadTools {
		series := getUploadSeries(cfg, c.Series)
		if err := v.uploadTools(env.Storage(), series); err != nil {
			return err
		}
	}
	if err := v.validate(); err != nil {
		return err
	}
	log.Infof("upgrade version chosen: %s", v.chosen)
	// TODO(fwereade): this list may be incomplete, pending envtools.Upload change.
	log.Infof("available tools: %s", v.tools)

	if err := client.SetEnvironAgentVersion(v.chosen); err != nil {
		return err
	}
	log.Noticef("started upgrade to %s", v.chosen)
	return nil
}

// run1dot16 implements the command without access to the API. This is
// needed for compatibility, so 1.16 can be upgraded to newer
// releases. It should be removed in 1.18.
func (c *UpgradeJujuCommand) run1dot16() error {
	log.Warningf("running in 1.16 compatibility mode")
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	defer func() {
		if err == errUpToDate {
			log.Noticef(err.Error())
			err = nil
		}
	}()

	// Determine the version to upgrade to, uploading tools if necessary.
	env := conn.Environ
	cfg, err := conn.State.EnvironConfig()
	if err != nil {
		return err
	}
	v, err := c.initVersions(cfg, env)
	if err != nil {
		return err
	}
	if c.UploadTools {
		series := getUploadSeries(cfg, c.Series)
		if err := v.uploadTools(env.Storage(), series); err != nil {
			return err
		}
	}
	if err := v.validate(); err != nil {
		return err
	}
	log.Infof("upgrade version chosen: %s", v.chosen)
	// TODO(fwereade): this list may be incomplete, pending envtools.Upload change.
	log.Infof("available tools: %s", v.tools)

	if err := conn.State.SetEnvironAgentVersion(v.chosen); err != nil {
		return err
	}
	log.Noticef("started upgrade to %s", v.chosen)
	return nil
}

// initVersions collects state relevant to an upgrade decision. The returned
// agent and client versions, and the list of currently available tools, will
// always be accurate; the chosen version, and the flag indicating development
// mode, may remain blank until uploadTools or validate is called.
func (c *UpgradeJujuCommand) initVersions(cfg *config.Config, env environs.Environ) (*upgradeVersions, error) {
	agent, ok := cfg.AgentVersion()
	if !ok {
		// Can't happen. In theory.
		return nil, fmt.Errorf("incomplete environment configuration")
	}
	if c.Version == agent {
		return nil, errUpToDate
	}
	client := version.Current.Number
	// TODO use an API call rather than requiring the environment,
	// so that we can restrict access to the provider secrets
	// while still allowing users to upgrade.
	available, err := envtools.FindTools(env, client.Major, -1, coretools.Filter{}, envtools.DoNotAllowRetry)
	if err != nil {
		if !errors.IsNotFoundError(err) {
			return nil, err
		}
		if !c.UploadTools {
			// No tools found and we shouldn't upload any, so if we are not asking for a
			// major upgrade, pretend there is no more recent version available.
			if c.Version == version.Zero && agent.Major == client.Major {
				return nil, errUpToDate
			}
			return nil, err
		}
	}
	return &upgradeVersions{
		agent:  agent,
		client: client,
		chosen: c.Version,
		tools:  available,
	}, nil
}

// upgradeVersions holds the version information for making upgrade decisions.
type upgradeVersions struct {
	agent  version.Number
	client version.Number
	chosen version.Number
	tools  coretools.List
}

// uploadTools compiles jujud from $GOPATH and uploads it into the supplied
// storage. If no version has been explicitly chosen, the version number
// reported by the built tools will be based on the client version number.
// In any case, the version number reported will have a build component higher
// than that of any otherwise-matching available envtools.
// uploadTools resets the chosen version and replaces the available tools
// with the ones just uploaded.
func (v *upgradeVersions) uploadTools(storage storage.Storage, series []string) error {
	// TODO(fwereade): this is kinda crack: we should not assume that
	// version.Current matches whatever source happens to be built. The
	// ideal would be:
	//  1) compile jujud from $GOPATH into some build dir
	//  2) get actual version with `jujud version`
	//  3) check actual version for compatibility with CLI tools
	//  4) generate unique build version with reference to available tools
	//  5) force-version that unique version into the dir directly
	//  6) archive and upload the build dir
	// ...but there's no way we have time for that now. In the meantime,
	// considering the use cases, this should work well enough; but it
	// won't detect an incompatible major-version change, which is a shame.
	if v.chosen == version.Zero {
		v.chosen = v.client
	}
	v.chosen = uploadVersion(v.chosen, v.tools)

	// TODO(fwereade): envtools.Upload should return envtools.List, and should
	// include all the extra series we build, so we can set *that* onto
	// v.available and maybe one day be able to check that a given upgrade
	// won't leave out-of-date machines lying around, starved of envtools.
	uploaded, err := sync.Upload(storage, &v.chosen, series...)
	if err != nil {
		return err
	}
	v.tools = coretools.List{uploaded}
	return nil
}

// validate chooses an upgrade version, if one has not already been chosen,
// and ensures the tools list contains no entries that do not have that version.
// If validate returns no error, the environment agent-version can be set to
// the value of the chosen field.
func (v *upgradeVersions) validate() (err error) {
	if v.chosen == version.Zero {
		// No explicitly specified version, so find the version to which we
		// need to upgrade. If the CLI and agent major versions match, we find
		// next available stable release to upgrade to by incrementing the
		// minor version, starting from the current agent version and doing
		// major.minor+1 or +2 as needed. If the CLI has a greater major version,
		// we just use the CLI version as is.
		nextVersion := v.agent
		if nextVersion.Major == v.client.Major {
			if v.agent.IsDev() {
				nextVersion.Minor += 1
			} else {
				nextVersion.Minor += 2
			}
		} else {
			nextVersion = v.client
		}

		newestNextStable, found := v.tools.NewestCompatible(nextVersion)
		if found {
			log.Debugf("found a more recent stable version %s", newestNextStable)
			v.chosen = newestNextStable
		} else {
			newestCurrent, found := v.tools.NewestCompatible(v.agent)
			if found {
				log.Debugf("found more recent current version %s", newestCurrent)
				v.chosen = newestCurrent
			} else {
				if v.agent.Major != v.client.Major {
					return fmt.Errorf("no compatible tools available")
				} else {
					return fmt.Errorf("no more recent supported versions available")
				}
			}
		}
	} else {
		// If not completely specified already, pick a single tools version.
		filter := coretools.Filter{Number: v.chosen, Released: !v.chosen.IsDev()}
		if v.tools, err = v.tools.Match(filter); err != nil {
			return err
		}
		v.chosen, v.tools = v.tools.Newest()
	}
	if v.chosen == v.agent {
		return errUpToDate
	}

	// Disallow major.minor version downgrades.
	if v.chosen.Major < v.agent.Major || v.chosen.Major == v.agent.Major && v.chosen.Minor < v.agent.Minor {
		// TODO(fwereade): I'm a bit concerned about old agent/CLI tools even
		// *connecting* to environments with higher agent-versions; but ofc they
		// have to connect in order to discover they shouldn't. However, once
		// any of our tools detect an incompatible version, they should act to
		// minimize damage: the CLI should abort politely, and the agents should
		// run an Upgrader but no other tasks.
		return fmt.Errorf("cannot change version from %s to %s", v.agent, v.chosen)
	}

	return nil
}

// uploadVersion returns a copy of the supplied version with a build number
// higher than any of the supplied tools that share its major, minor and patch.
func uploadVersion(vers version.Number, existing coretools.List) version.Number {
	vers.Build++
	for _, t := range existing {
		if t.Version.Major != vers.Major || t.Version.Minor != vers.Minor || t.Version.Patch != vers.Patch {
			continue
		}
		if t.Version.Build >= vers.Build {
			vers.Build = t.Version.Build + 1
		}
	}
	return vers
}
<|endoftext|>


<commit_before>package kfs

import (
	"log"
	"sync"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/Shopify/sarama"
	"golang.org/x/net/context"
)

type Dir struct {
	sync.RWMutex
	attr fuse.Attr

	path string
	fs   *KafkaFS
}

func (d *Dir) Attr(ctx context.Context, o *fuse.Attr) error {
	d.RLock()
	defer d.RUnlock()

	*o = d.attr

	return nil
}

func (d *Dir) Lookup(ctx context.Context, name string) (fs.Node, error) {
	d.RLock()
	defer d.RUnlock()

	log.Printf("[Dir] Lookup, name=%s", name)

	return nil, fuse.ENOENT
}

func (d *Dir) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	d.RLock()
	defer d.RUnlock()

	var out []fuse.Dirent

	kfk, err := sarama.NewClient(d.fs.zkcluster.BrokerList(), sarama.NewConfig())
	if err != nil {
		log.Println(err)

		return nil, err
	}
	defer kfk.Close()

	topics, err := kfk.Topics()
	if err != nil {
		log.Println(err)

		return nil, err
	}

	for _, topic := range topics {
		de := fuse.Dirent{
			Name: topic,
			Type: fuse.DT_File,
		}

		out = append(out, de)
	}

	return out, nil
}

func (d *Dir) Mkdir(ctx context.Context, req *fuse.MkdirRequest) (fs.Node, error) {
	return nil, fuse.EPERM
}

func (d *Dir) Create(ctx context.Context, req *fuse.CreateRequest, resp *fuse.CreateResponse) (fs.Node, fs.Handle, error) {
	return nil, nil, fuse.EPERM
}

func (d *Dir) Rename(ctx context.Context, req *fuse.RenameRequest, newDir fs.Node) error {
	return fuse.EPERM
}

func (d *Dir) Remove(ctx context.Context, req *fuse.RemoveRequest) error {
	return fuse.EPERM
}
<commit_msg>list topic name suffixed with partition id<commit_after>package kfs

import (
	"fmt"
	"log"
	"sync"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/Shopify/sarama"
	"golang.org/x/net/context"
)

type Dir struct {
	sync.RWMutex
	attr fuse.Attr

	path string
	fs   *KafkaFS
}

func (d *Dir) Attr(ctx context.Context, o *fuse.Attr) error {
	d.RLock()
	defer d.RUnlock()

	*o = d.attr

	return nil
}

func (d *Dir) Lookup(ctx context.Context, name string) (fs.Node, error) {
	d.RLock()
	defer d.RUnlock()

	log.Printf("Dir Lookup, name=%s", name)

	return nil, fuse.ENOENT
}

func (d *Dir) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	d.RLock()
	defer d.RUnlock()

	var out []fuse.Dirent

	kfk, err := sarama.NewClient(d.fs.zkcluster.BrokerList(), sarama.NewConfig())
	if err != nil {
		log.Println(err)

		return nil, err
	}
	defer kfk.Close()

	topics, err := kfk.Topics()
	if err != nil {
		log.Println(err)

		return nil, err
	}

	for _, topic := range topics {
		partitions, err := kfk.Partitions(topic)
		if err != nil {
			log.Println(err)

			return nil, err
		}

		for _, p := range partitions {
			de := fuse.Dirent{
				Name: fmt.Sprintf("%s.%d", topic, p),
				Type: fuse.DT_File,
			}

			out = append(out, de)
		}

	}

	return out, nil
}

func (d *Dir) Mkdir(ctx context.Context, req *fuse.MkdirRequest) (fs.Node, error) {
	return nil, fuse.EPERM
}

func (d *Dir) Create(ctx context.Context, req *fuse.CreateRequest, resp *fuse.CreateResponse) (fs.Node, fs.Handle, error) {
	return nil, nil, fuse.EPERM
}

func (d *Dir) Rename(ctx context.Context, req *fuse.RenameRequest, newDir fs.Node) error {
	return fuse.EPERM
}

func (d *Dir) Remove(ctx context.Context, req *fuse.RemoveRequest) error {
	return fuse.EPERM
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package main

import (
	"context"
	"fmt"
	"io"
	"strings"

	"github.com/spf13/cobra"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kops/cmd/kops/util"
	kopsapi "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/registry"
	"k8s.io/kops/pkg/commands/commandutils"
	"k8s.io/kops/pkg/kopscodecs"
	"k8s.io/kops/util/pkg/tables"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

var (
	getClusterLong = templates.LongDesc(i18n.T(`
	Display one or many cluster resources.`))

	getClusterExample = templates.Examples(i18n.T(`
	# Get all clusters in a state store
	kops get clusters

	# Get a cluster
	kops get cluster k8s-cluster.example.com

	# Get a cluster YAML desired configuration
	kops get cluster k8s-cluster.example.com -o yaml

	# Save a cluster desired configuration to YAML file
	kops get cluster k8s-cluster.example.com -o yaml > cluster-desired-config.yaml
	`))

	getClusterShort = i18n.T(`Get one or many clusters.`)

	// Warning for --full.  Since we are not using the template from kubectl
	// we have to have zero white space before the comment characters otherwise
	// output to stdout is going to be off.
	get_cluster_full_warning = i18n.T(`
//
//   WARNING: Do not use a '--full' cluster specification to define a Kubernetes installation.
//   You may experience unexpected behavior and other bugs.  Use only the required elements
//   and any modifications that you require.
//
//   Use the following command to retrieve only the required elements:
//   $ kops get cluster -o yaml
//

`)
)

type GetClusterOptions struct {
	*GetOptions

	// FullSpec determines if we should output the completed (fully populated) spec
	FullSpec bool

	// ClusterNames is a list of cluster names to show; if not specified all clusters will be shown
	ClusterNames []string
}

func NewCmdGetCluster(f *util.Factory, out io.Writer, getOptions *GetOptions) *cobra.Command {
	options := GetClusterOptions{
		GetOptions: getOptions,
	}

	cmd := &cobra.Command{
		Use:     "clusters [CLUSTER]...",
		Aliases: []string{"cluster"},
		Short:   getClusterShort,
		Long:    getClusterLong,
		Example: getClusterExample,
		Args: func(cmd *cobra.Command, args []string) error {
			if len(args) != 0 {
				if rootCommand.clusterName != "" {
					return fmt.Errorf("cannot mix --name for cluster with positional arguments")
				}
				options.ClusterNames = append(options.ClusterNames, args...)
			} else {
				options.ClusterNames = append(options.ClusterNames, rootCommand.ClusterName(true))
			}

			return nil
		},
		ValidArgsFunction: commandutils.CompleteClusterName(f, false, true),
		RunE: func(cmd *cobra.Command, args []string) error {
			return RunGetClusters(context.TODO(), f, out, &options)
		},
	}

	cmd.Flags().BoolVar(&options.FullSpec, "full", options.FullSpec, "Show fully populated configuration")

	return cmd
}

func RunGetClusters(ctx context.Context, f commandutils.Factory, out io.Writer, options *GetClusterOptions) error {
	client, err := f.Clientset()
	if err != nil {
		return err
	}

	var clusterList []*kopsapi.Cluster
	if len(options.ClusterNames) != 1 {
		list, err := client.ListClusters(ctx, metav1.ListOptions{})
		if err != nil {
			return err
		}
		for i := range list.Items {
			clusterList = append(clusterList, &list.Items[i])
		}
	} else {
		// Optimization - avoid fetching all clusters if we're only querying one
		cluster, err := client.GetCluster(ctx, options.ClusterNames[0])
		if err != nil {
			if !apierrors.IsNotFound(err) {
				return err
			}
		} else {
			clusterList = append(clusterList, cluster)
		}
	}

	clusters, err := filterClustersByName(options.ClusterNames, clusterList)
	if err != nil {
		return err
	}

	if len(clusters) == 0 {
		return fmt.Errorf("no clusters found")
	}

	if options.FullSpec {
		var err error
		clusters, err = fullClusterSpecs(clusters)
		if err != nil {
			return err
		}

		fmt.Fprint(out, get_cluster_full_warning)
	}

	var obj []runtime.Object
	if options.Output != OutputTable {
		for _, c := range clusters {
			obj = append(obj, c)
		}
	}

	switch options.Output {
	case OutputTable:
		return clusterOutputTable(clusters, out)
	case OutputYaml:
		return fullOutputYAML(out, obj...)
	case OutputJSON:
		return fullOutputJSON(out, obj...)
	default:
		return fmt.Errorf("Unknown output format: %q", options.Output)
	}
}

// filterClustersByName returns the clusters matching the specified names.
// If names are specified and no cluster is found with a name, we return an error.
func filterClustersByName(clusterNames []string, clusters []*kopsapi.Cluster) ([]*kopsapi.Cluster, error) {
	if len(clusterNames) != 0 {
		// Build a map as we want to return them in the same order as args
		m := make(map[string]*kopsapi.Cluster)
		for _, c := range clusters {
			m[c.ObjectMeta.Name] = c
		}
		var filtered []*kopsapi.Cluster
		for _, clusterName := range clusterNames {
			c := m[clusterName]
			if c == nil {
				return nil, fmt.Errorf("cluster not found %q", clusterName)
			}

			filtered = append(filtered, c)
		}
		return filtered, nil
	}

	return clusters, nil
}

func clusterOutputTable(clusters []*kopsapi.Cluster, out io.Writer) error {
	t := &tables.Table{}
	t.AddColumn("NAME", func(c *kopsapi.Cluster) string {
		return c.ObjectMeta.Name
	})
	t.AddColumn("CLOUD", func(c *kopsapi.Cluster) string {
		return c.Spec.CloudProvider
	})
	t.AddColumn("ZONES", func(c *kopsapi.Cluster) string {
		zones := sets.NewString()
		for _, s := range c.Spec.Subnets {
			if s.Zone != "" {
				zones.Insert(s.Zone)
			}
		}
		return strings.Join(zones.List(), ",")
	})

	return t.Render(clusters, out, "NAME", "CLOUD", "ZONES")
}

// fullOutputJson outputs the marshalled JSON of a list of clusters and instance groups.  It will handle
// nils for clusters and instanceGroups slices.
func fullOutputJSON(out io.Writer, args ...runtime.Object) error {
	argsLen := len(args)

	if argsLen > 1 {
		if _, err := fmt.Fprint(out, "["); err != nil {
			return err
		}
	}

	for i, arg := range args {
		if i != 0 {
			if _, err := fmt.Fprint(out, ","); err != nil {
				return err
			}
		}
		if err := marshalToWriter(arg, marshalJSON, out); err != nil {
			return err
		}
	}

	if argsLen > 1 {
		if _, err := fmt.Fprint(out, "]"); err != nil {
			return err
		}
	}

	return nil
}

// fullOutputJson outputs the marshalled JSON of a list of clusters and instance groups.  It will handle
// nils for clusters and instanceGroups slices.
func fullOutputYAML(out io.Writer, args ...runtime.Object) error {
	for i, obj := range args {
		if i != 0 {
			if err := writeYAMLSep(out); err != nil {
				return fmt.Errorf("error writing to stdout: %v", err)
			}
		}
		if err := marshalToWriter(obj, marshalYaml, out); err != nil {
			return err
		}
	}
	return nil
}

func fullClusterSpecs(clusters []*kopsapi.Cluster) ([]*kopsapi.Cluster, error) {
	var fullSpecs []*kopsapi.Cluster
	for _, cluster := range clusters {
		configBase, err := registry.ConfigBase(cluster)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		configPath := configBase.Join(registry.PathClusterCompleted)
		b, err := configPath.ReadFile()
		if err != nil {
			return nil, fmt.Errorf("error loading Cluster %q: %v", configPath, err)
		}

		o, _, err := kopscodecs.Decode(b, nil)
		if err != nil {
			return nil, fmt.Errorf("error parsing Cluster %q: %v", configPath, err)
		}
		if fullSpec, ok := o.(*kopsapi.Cluster); ok {
			fullSpecs = append(fullSpecs, fullSpec)
		} else {
			return nil, fmt.Errorf("unexpected object type for Cluster %q: %T", configPath, o)
		}
	}
	return fullSpecs, nil
}
<commit_msg>Fix cluster list action<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package main

import (
	"context"
	"fmt"
	"io"
	"strings"

	"github.com/spf13/cobra"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kops/cmd/kops/util"
	kopsapi "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/registry"
	"k8s.io/kops/pkg/commands/commandutils"
	"k8s.io/kops/pkg/kopscodecs"
	"k8s.io/kops/util/pkg/tables"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

var (
	getClusterLong = templates.LongDesc(i18n.T(`
	Display one or many cluster resources.`))

	getClusterExample = templates.Examples(i18n.T(`
	# Get all clusters in a state store
	kops get clusters

	# Get a cluster
	kops get cluster k8s-cluster.example.com

	# Get a cluster YAML desired configuration
	kops get cluster k8s-cluster.example.com -o yaml

	# Save a cluster desired configuration to YAML file
	kops get cluster k8s-cluster.example.com -o yaml > cluster-desired-config.yaml
	`))

	getClusterShort = i18n.T(`Get one or many clusters.`)

	// Warning for --full.  Since we are not using the template from kubectl
	// we have to have zero white space before the comment characters otherwise
	// output to stdout is going to be off.
	get_cluster_full_warning = i18n.T(`
//
//   WARNING: Do not use a '--full' cluster specification to define a Kubernetes installation.
//   You may experience unexpected behavior and other bugs.  Use only the required elements
//   and any modifications that you require.
//
//   Use the following command to retrieve only the required elements:
//   $ kops get cluster -o yaml
//

`)
)

type GetClusterOptions struct {
	*GetOptions

	// FullSpec determines if we should output the completed (fully populated) spec
	FullSpec bool

	// ClusterNames is a list of cluster names to show; if not specified all clusters will be shown
	ClusterNames []string
}

func NewCmdGetCluster(f *util.Factory, out io.Writer, getOptions *GetOptions) *cobra.Command {
	options := GetClusterOptions{
		GetOptions: getOptions,
	}

	cmd := &cobra.Command{
		Use:     "clusters [CLUSTER]...",
		Aliases: []string{"cluster"},
		Short:   getClusterShort,
		Long:    getClusterLong,
		Example: getClusterExample,
		Args: func(cmd *cobra.Command, args []string) error {
			if len(args) != 0 {
				if rootCommand.clusterName != "" {
					return fmt.Errorf("cannot mix --name for cluster with positional arguments")
				}
				options.ClusterNames = append(options.ClusterNames, args...)
			} else if rootCommand.clusterName != "" {
				options.ClusterNames = append(options.ClusterNames, rootCommand.clusterName)
			}

			return nil
		},
		ValidArgsFunction: commandutils.CompleteClusterName(f, false, true),
		RunE: func(cmd *cobra.Command, args []string) error {
			return RunGetClusters(context.TODO(), f, out, &options)
		},
	}

	cmd.Flags().BoolVar(&options.FullSpec, "full", options.FullSpec, "Show fully populated configuration")

	return cmd
}

func RunGetClusters(ctx context.Context, f commandutils.Factory, out io.Writer, options *GetClusterOptions) error {
	client, err := f.Clientset()
	if err != nil {
		return err
	}

	var clusterList []*kopsapi.Cluster
	if len(options.ClusterNames) != 1 {
		list, err := client.ListClusters(ctx, metav1.ListOptions{})
		if err != nil {
			return err
		}
		for i := range list.Items {
			clusterList = append(clusterList, &list.Items[i])
		}
	} else {
		// Optimization - avoid fetching all clusters if we're only querying one
		cluster, err := client.GetCluster(ctx, options.ClusterNames[0])
		if err != nil {
			if !apierrors.IsNotFound(err) {
				return err
			}
		} else {
			clusterList = append(clusterList, cluster)
		}
	}

	clusters, err := filterClustersByName(options.ClusterNames, clusterList)
	if err != nil {
		return err
	}

	if len(clusters) == 0 {
		return fmt.Errorf("no clusters found")
	}

	if options.FullSpec {
		var err error
		clusters, err = fullClusterSpecs(clusters)
		if err != nil {
			return err
		}

		fmt.Fprint(out, get_cluster_full_warning)
	}

	var obj []runtime.Object
	if options.Output != OutputTable {
		for _, c := range clusters {
			obj = append(obj, c)
		}
	}

	switch options.Output {
	case OutputTable:
		return clusterOutputTable(clusters, out)
	case OutputYaml:
		return fullOutputYAML(out, obj...)
	case OutputJSON:
		return fullOutputJSON(out, obj...)
	default:
		return fmt.Errorf("Unknown output format: %q", options.Output)
	}
}

// filterClustersByName returns the clusters matching the specified names.
// If names are specified and no cluster is found with a name, we return an error.
func filterClustersByName(clusterNames []string, clusters []*kopsapi.Cluster) ([]*kopsapi.Cluster, error) {
	if len(clusterNames) != 0 {
		// Build a map as we want to return them in the same order as args
		m := make(map[string]*kopsapi.Cluster)
		for _, c := range clusters {
			m[c.ObjectMeta.Name] = c
		}
		var filtered []*kopsapi.Cluster
		for _, clusterName := range clusterNames {
			c := m[clusterName]
			if c == nil {
				return nil, fmt.Errorf("cluster not found %q", clusterName)
			}

			filtered = append(filtered, c)
		}
		return filtered, nil
	}

	return clusters, nil
}

func clusterOutputTable(clusters []*kopsapi.Cluster, out io.Writer) error {
	t := &tables.Table{}
	t.AddColumn("NAME", func(c *kopsapi.Cluster) string {
		return c.ObjectMeta.Name
	})
	t.AddColumn("CLOUD", func(c *kopsapi.Cluster) string {
		return c.Spec.CloudProvider
	})
	t.AddColumn("ZONES", func(c *kopsapi.Cluster) string {
		zones := sets.NewString()
		for _, s := range c.Spec.Subnets {
			if s.Zone != "" {
				zones.Insert(s.Zone)
			}
		}
		return strings.Join(zones.List(), ",")
	})

	return t.Render(clusters, out, "NAME", "CLOUD", "ZONES")
}

// fullOutputJson outputs the marshalled JSON of a list of clusters and instance groups.  It will handle
// nils for clusters and instanceGroups slices.
func fullOutputJSON(out io.Writer, args ...runtime.Object) error {
	argsLen := len(args)

	if argsLen > 1 {
		if _, err := fmt.Fprint(out, "["); err != nil {
			return err
		}
	}

	for i, arg := range args {
		if i != 0 {
			if _, err := fmt.Fprint(out, ","); err != nil {
				return err
			}
		}
		if err := marshalToWriter(arg, marshalJSON, out); err != nil {
			return err
		}
	}

	if argsLen > 1 {
		if _, err := fmt.Fprint(out, "]"); err != nil {
			return err
		}
	}

	return nil
}

// fullOutputJson outputs the marshalled JSON of a list of clusters and instance groups.  It will handle
// nils for clusters and instanceGroups slices.
func fullOutputYAML(out io.Writer, args ...runtime.Object) error {
	for i, obj := range args {
		if i != 0 {
			if err := writeYAMLSep(out); err != nil {
				return fmt.Errorf("error writing to stdout: %v", err)
			}
		}
		if err := marshalToWriter(obj, marshalYaml, out); err != nil {
			return err
		}
	}
	return nil
}

func fullClusterSpecs(clusters []*kopsapi.Cluster) ([]*kopsapi.Cluster, error) {
	var fullSpecs []*kopsapi.Cluster
	for _, cluster := range clusters {
		configBase, err := registry.ConfigBase(cluster)
		if err != nil {
			return nil, fmt.Errorf("error reading full cluster spec for %q: %v", cluster.ObjectMeta.Name, err)
		}
		configPath := configBase.Join(registry.PathClusterCompleted)
		b, err := configPath.ReadFile()
		if err != nil {
			return nil, fmt.Errorf("error loading Cluster %q: %v", configPath, err)
		}

		o, _, err := kopscodecs.Decode(b, nil)
		if err != nil {
			return nil, fmt.Errorf("error parsing Cluster %q: %v", configPath, err)
		}
		if fullSpec, ok := o.(*kopsapi.Cluster); ok {
			fullSpecs = append(fullSpecs, fullSpec)
		} else {
			return nil, fmt.Errorf("unexpected object type for Cluster %q: %T", configPath, o)
		}
	}
	return fullSpecs, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strings"
	"unicode"

	"github.com/bmatsuo/lark/lib/lark/core"
	"github.com/codegangsta/cli"
	"github.com/yuin/gopher-lua"
)

// CommandRun implements the "run" action (the default action)
var CommandRun = Command(func(lark *Context, cmd *cli.Command) {
	cmd.Name = "run"
	cmd.Aliases = []string{"make"}
	cmd.Usage = "Run lark project task(s)"
	cmd.ArgsUsage = `task ...

    The arguments are the names of tasks from lark.lua.`
	cmd.Flags = []cli.Flag{
		cli.IntFlag{
			Name:   "j",
			Usage:  "Number of parallel processes.",
			EnvVar: "LARK_RUN_PARALLEL",
		},
		cli.BoolFlag{
			Name:        "v",
			Usage:       "Enable verbose reporting of errors.",
			EnvVar:      "LARK_VERBOSE",
			Destination: Verbose,
		},
	}
	cmd.Action = lark.Action(Run)
})

// Run loads a lua vm and runs tasks specified in the command line.
func Run(c *Context) {
	args := c.Args()
	var tasks []*Task
	for {
		t, n, err := ParseTask(args)
		if err != nil {
			log.Fatalf("task %d: %v", len(tasks), err)
		}
		if n == 0 {
			break
		}
		tasks = append(tasks, t)
		args = args[n:]
	}
	if len(tasks) == 0 {
		tasks = []*Task{{}}
	}

	luaFiles, err := FindTaskFiles("")
	if err != nil {
		log.Fatal(err)
	}

	luaConfig := &LuaConfig{}
	c.Lua, err = LoadVM(luaConfig)
	if err != nil {
		log.Fatal(err)
	}
	defer c.Lua.Close()

	err = InitLark(c, luaFiles)
	if err != nil {
		log.Fatal(err)
	}

	for _, task := range tasks {
		err := RunTask(c, task)
		if err != nil {
			os.Exit(1)
		}
	}
}

func normTasks(args []string) ([]string, error) {
	if len(args) == 0 {
		return []string{""}, nil
	}
	for _, task := range args {
		if task == "" {
			return nil, fmt.Errorf("invalid task name")
		}
	}
	return args, nil
}

// RunTask calls lark.run in state to execute task.
func RunTask(c *Context, task *Task) error {
	lark := c.Lua.GetGlobal("lark")
	run := c.Lua.GetField(lark, "run")
	trace := c.Lua.NewFunction(errTraceback)

	narg := 1
	c.Lua.Push(run)
	if task.Name == "" {
		c.Lua.Push(lua.LNil)
	} else {
		c.Lua.Push(lua.LString(task.Name))
	}
	if len(task.Params) > 0 {
		params := c.Lua.NewTable()
		for k, v := range task.Params {
			c.Lua.SetField(params, k, lua.LString(v))
		}

		c.Lua.Push(params)
		narg++
	}
	err := c.Lua.PCall(narg, 0, trace)
	if err != nil {
		handleErr(c, err)
	}

	wait := c.Lua.GetField(lark, "wait")
	for {
		c.Lua.Push(wait)
		errwait := c.Lua.PCall(0, 0, trace)
		if errwait == nil {
			break
		}
		if err == nil {
			handleErr(c, errwait)

			// prevent handleErr from being called multiple times.
			err = errwait
		}
	}

	return err
}

func handleErr(c *Context, err error) {
	core.Log(fmt.Sprint(err), &core.LogOpt{
		Color: "red",
	})
}

// Task is a task invocation from the command line.
type Task struct {
	Name   string
	Params map[string]string
}

// ToLua returns a string representing the task in lua table syntax.
func (t *Task) ToLua() string {
	var name string
	if t.Name == "" {
		name = "nil"
	} else {
		name = fmt.Sprintf("%q", t.Name)
	}
	return fmt.Sprintf("{name=%s,params=%s}", name, luamap(t.Params))
}

func luamap(m map[string]string) string {
	buf := bytes.NewBuffer(nil)
	io.WriteString(buf, "{")
	for k, v := range m {
		io.WriteString(buf, k)
		io.WriteString(buf, "=")
		fmt.Fprintf(buf, "%q", v)
	}
	io.WriteString(buf, "}")
	return buf.String()
}

// ParseTask parses a task from command line arguments and returns it along
// with the number of args consumed.
func ParseTask(args []string) (*Task, int, error) {
	t := &Task{}
	if len(args) == 0 {
		return t, 0, nil
	}
	if args[0] == "--" {
		return t, 1, nil
	}
	if !strings.Contains(args[0], "=") {
		t.Name = args[0]
		args = args[1:]
		if len(t.Name) == 0 {
			return nil, 0, fmt.Errorf("missing name")
		}
	}
	t.Params = make(map[string]string)
	for _, p := range args {
		pieces := strings.SplitN(p, "=", 2)
		if len(pieces) == 1 {
			break
		}
		err := sanitizeParam(pieces[0])
		if err != nil {
			return nil, 0, err
		}
		t.Params[pieces[0]] = pieces[1]
	}
	return t, 1 + len(t.Params), nil
}

func sanitizeParam(name string) error {
	if len(name) == 0 {
		return fmt.Errorf("missing param name")
	}
	badchars := strings.TrimFunc(name, func(c rune) bool {
		return unicode.IsLetter(c) || unicode.IsNumber(c) || c == '_'
	})
	if len(badchars) > 0 {
		return fmt.Errorf("invalid character in param: %q", badchars[0])
	}
	return nil
}

var reLoc = regexp.MustCompile(`^[^:]+:\d+:\s*`)

func trimLoc(msg string) string {
	return reLoc.ReplaceAllString(msg, "")
}

func errTraceback(L *lua.LState) int {
	msg := L.Get(1)
	L.SetTop(0)
	L.Push(L.GetField(L.GetGlobal("debug"), "traceback"))
	L.Push(msg)
	L.Push(lua.LNumber(2))
	L.Call(2, 1)
	return 1
}
<commit_msg>run: the -C flag changes the working directory before loading files<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strings"
	"unicode"

	"github.com/bmatsuo/lark/lib/lark/core"
	"github.com/codegangsta/cli"
	"github.com/yuin/gopher-lua"
)

// CommandRun implements the "run" action (the default action)
var CommandRun = Command(func(lark *Context, cmd *cli.Command) {
	cmd.Name = "run"
	cmd.Aliases = []string{"make"}
	cmd.Usage = "Run lark project task(s)"
	cmd.ArgsUsage = `task ...

    The arguments are the names of tasks from lark.lua.`
	cmd.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "C",
			Usage:  "Change the working directory before loading files and running tasks",
			EnvVar: "LARK_RUN_DIRECTORY",
		},
		cli.IntFlag{
			Name:   "j",
			Usage:  "Number of parallel processes.",
			EnvVar: "LARK_RUN_PARALLEL",
		},
		cli.BoolFlag{
			Name:        "v",
			Usage:       "Enable verbose reporting of errors.",
			EnvVar:      "LARK_VERBOSE",
			Destination: Verbose,
		},
	}
	cmd.Action = lark.Action(Run)
})

// Run loads a lua vm and runs tasks specified in the command line.
func Run(c *Context) {
	chdir := c.String("C")
	if chdir != "" {
		err := os.Chdir(chdir)
		if err != nil {
			log.Fatal(err)
		}
	}

	args := c.Args()
	var tasks []*Task
	for {
		t, n, err := ParseTask(args)
		if err != nil {
			log.Fatalf("task %d: %v", len(tasks), err)
		}
		if n == 0 {
			break
		}
		tasks = append(tasks, t)
		args = args[n:]
	}
	if len(tasks) == 0 {
		tasks = []*Task{{}}
	}

	luaFiles, err := FindTaskFiles("")
	if err != nil {
		log.Fatal(err)
	}

	luaConfig := &LuaConfig{}
	c.Lua, err = LoadVM(luaConfig)
	if err != nil {
		log.Fatal(err)
	}
	defer c.Lua.Close()

	err = InitLark(c, luaFiles)
	if err != nil {
		log.Fatal(err)
	}

	for _, task := range tasks {
		err := RunTask(c, task)
		if err != nil {
			os.Exit(1)
		}
	}
}

func normTasks(args []string) ([]string, error) {
	if len(args) == 0 {
		return []string{""}, nil
	}
	for _, task := range args {
		if task == "" {
			return nil, fmt.Errorf("invalid task name")
		}
	}
	return args, nil
}

// RunTask calls lark.run in state to execute task.
func RunTask(c *Context, task *Task) error {
	lark := c.Lua.GetGlobal("lark")
	run := c.Lua.GetField(lark, "run")
	trace := c.Lua.NewFunction(errTraceback)

	narg := 1
	c.Lua.Push(run)
	if task.Name == "" {
		c.Lua.Push(lua.LNil)
	} else {
		c.Lua.Push(lua.LString(task.Name))
	}
	if len(task.Params) > 0 {
		params := c.Lua.NewTable()
		for k, v := range task.Params {
			c.Lua.SetField(params, k, lua.LString(v))
		}

		c.Lua.Push(params)
		narg++
	}
	err := c.Lua.PCall(narg, 0, trace)
	if err != nil {
		handleErr(c, err)
	}

	wait := c.Lua.GetField(lark, "wait")
	for {
		c.Lua.Push(wait)
		errwait := c.Lua.PCall(0, 0, trace)
		if errwait == nil {
			break
		}
		if err == nil {
			handleErr(c, errwait)

			// prevent handleErr from being called multiple times.
			err = errwait
		}
	}

	return err
}

func handleErr(c *Context, err error) {
	core.Log(fmt.Sprint(err), &core.LogOpt{
		Color: "red",
	})
}

// Task is a task invocation from the command line.
type Task struct {
	Name   string
	Params map[string]string
}

// ToLua returns a string representing the task in lua table syntax.
func (t *Task) ToLua() string {
	var name string
	if t.Name == "" {
		name = "nil"
	} else {
		name = fmt.Sprintf("%q", t.Name)
	}
	return fmt.Sprintf("{name=%s,params=%s}", name, luamap(t.Params))
}

func luamap(m map[string]string) string {
	buf := bytes.NewBuffer(nil)
	io.WriteString(buf, "{")
	for k, v := range m {
		io.WriteString(buf, k)
		io.WriteString(buf, "=")
		fmt.Fprintf(buf, "%q", v)
	}
	io.WriteString(buf, "}")
	return buf.String()
}

// ParseTask parses a task from command line arguments and returns it along
// with the number of args consumed.
func ParseTask(args []string) (*Task, int, error) {
	t := &Task{}
	if len(args) == 0 {
		return t, 0, nil
	}
	if args[0] == "--" {
		return t, 1, nil
	}
	if !strings.Contains(args[0], "=") {
		t.Name = args[0]
		args = args[1:]
		if len(t.Name) == 0 {
			return nil, 0, fmt.Errorf("missing name")
		}
	}
	t.Params = make(map[string]string)
	for _, p := range args {
		pieces := strings.SplitN(p, "=", 2)
		if len(pieces) == 1 {
			break
		}
		err := sanitizeParam(pieces[0])
		if err != nil {
			return nil, 0, err
		}
		t.Params[pieces[0]] = pieces[1]
	}
	return t, 1 + len(t.Params), nil
}

func sanitizeParam(name string) error {
	if len(name) == 0 {
		return fmt.Errorf("missing param name")
	}
	badchars := strings.TrimFunc(name, func(c rune) bool {
		return unicode.IsLetter(c) || unicode.IsNumber(c) || c == '_'
	})
	if len(badchars) > 0 {
		return fmt.Errorf("invalid character in param: %q", badchars[0])
	}
	return nil
}

var reLoc = regexp.MustCompile(`^[^:]+:\d+:\s*`)

func trimLoc(msg string) string {
	return reLoc.ReplaceAllString(msg, "")
}

func errTraceback(L *lua.LState) int {
	msg := L.Get(1)
	L.SetTop(0)
	L.Push(L.GetField(L.GetGlobal("debug"), "traceback"))
	L.Push(msg)
	L.Push(lua.LNumber(2))
	L.Call(2, 1)
	return 1
}
<|endoftext|>