<commit_before>package errcheck

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"regexp"
	"runtime"
	"strings"
	"testing"

	"golang.org/x/tools/go/packages"
)

const testPackage = "github.com/kisielk/errcheck/testdata"

var (
	uncheckedMarkers map[marker]bool
	blankMarkers     map[marker]bool
	assertMarkers    map[marker]bool
)

type marker struct {
	file string
	line int
}

func newMarker(e UncheckedError) marker {
	return marker{e.Pos.Filename, e.Pos.Line}
}

func (m marker) String() string {
	return fmt.Sprintf("%s:%d", m.file, m.line)
}

func init() {
	uncheckedMarkers = make(map[marker]bool)
	blankMarkers = make(map[marker]bool)
	assertMarkers = make(map[marker]bool)

	cfg := &packages.Config{
		Mode:  packages.LoadSyntax,
		Tests: true,
	}
	pkgs, err := packages.Load(cfg, testPackage)
	if err != nil {
		panic("failed to import test package")
	}
	for _, pkg := range pkgs {
		for _, file := range pkg.Syntax {
			for _, comment := range file.Comments {
				text := comment.Text()
				pos := pkg.Fset.Position(comment.Pos())
				m := marker{pos.Filename, pos.Line}
				switch text {
				case "UNCHECKED\n":
					uncheckedMarkers[m] = true
				case "BLANK\n":
					blankMarkers[m] = true
				case "ASSERT\n":
					assertMarkers[m] = true
				}
			}
		}
	}
}

type flags uint

const (
	CheckAsserts flags = 1 << iota
	CheckBlank
)

// TestUnchecked runs a test against the example files and ensures all unchecked errors are caught.
func TestUnchecked(t *testing.T) {
	test(t, 0)
}

// TestBlank is like TestUnchecked but also ensures assignments to the blank identifier are caught.
func TestBlank(t *testing.T) {
	test(t, CheckBlank)
}

func TestAll(t *testing.T) {
	// TODO: CheckAsserts should work independently of CheckBlank
	test(t, CheckAsserts|CheckBlank)
}

func TestBuildTags(t *testing.T) {
	const (
		// uses "custom1" build tag and contains 1 unchecked error
		testBuildCustom1Tag = `
` + `// +build custom1

package custom

import "fmt"

func Print1() {
	// returns an error that is not checked
	fmt.Fprintln(nil)
}`
		// uses "custom2" build tag and contains 1 unchecked error
		testBuildCustom2Tag = `
` + `// +build custom2

package custom

import "fmt"

func Print2() {
	// returns an error that is not checked
	fmt.Fprintln(nil)
}`
		// included so that package is not empty when built without specifying tags
		testDoc = `
// Package custom contains code for testing build tags.
package custom
`
	)

	testBuildTagsDir, err := ioutil.TempDir(".", "testbuildtags")
	if err != nil {
		t.Fatalf("unable to create testbuildtags directory: %v", err)
	}
	defer os.RemoveAll(testBuildTagsDir)

	if err := ioutil.WriteFile(path.Join(testBuildTagsDir, "custom1.go"), []byte(testBuildCustom1Tag), 0644); err != nil {
		t.Fatalf("Failed to write testbuildtags custom1: %v", err)
	}
	if err := ioutil.WriteFile(path.Join(testBuildTagsDir, "custom2.go"), []byte(testBuildCustom2Tag), 0644); err != nil {
		t.Fatalf("Failed to write testbuildtags custom2: %v", err)
	}
	if err := ioutil.WriteFile(path.Join(testBuildTagsDir, "doc.go"), []byte(testDoc), 0644); err != nil {
		t.Fatalf("Failed to write testbuildtags doc: %v", err)
	}

	cases := []struct {
		tags            []string
		numExpectedErrs int
	}{
		// with no tags specified, main is ignored and there are no errors
		{
			tags:            nil,
			numExpectedErrs: 0,
		},
		// specifying "custom1" tag includes file with 1 error
		{
			tags:            []string{"custom1"},
			numExpectedErrs: 1,
		},
		// specifying "custom1" and "custom2" tags includes 2 files with 1 error each
		{
			tags:            []string{"custom1", "custom2"},
			numExpectedErrs: 2,
		},
	}

	for i, currCase := range cases {
		checker := NewChecker()
		checker.Tags = currCase.tags
		err := checker.CheckPackages(path.Join("github.com/kisielk/errcheck/internal/errcheck", testBuildTagsDir))

		if currCase.numExpectedErrs == 0 {
			if err != nil {
				t.Errorf("Case %d: expected no errors, but got: %v", i, err)
			}
			continue
		}

		uerr, ok := err.(*UncheckedErrors)
		if !ok {
			t.Errorf("Case %d: wrong error type returned: %v", i, err)
			continue
		}

		if currCase.numExpectedErrs != len(uerr.Errors) {
			t.Errorf("Case %d:\nExpected: %d errors\nActual:   %d errors", i, currCase.numExpectedErrs, len(uerr.Errors))
		}
	}
}

func TestWhitelist(t *testing.T) {

}

func TestIgnore(t *testing.T) {
	const testVendorMain = `
	package main

	import "github.com/testlog"

	func main() {
		// returns an error that is not checked
		testlog.Info()
	}`
	const testLog = `
	package testlog

	func Info() error {
		return nil
	}`

	if strings.HasPrefix(runtime.Version(), "go1.5") && os.Getenv("GO15VENDOREXPERIMENT") != "1" {
		// skip tests if running in go1.5 and vendoring is not enabled
		t.SkipNow()
	}

	// copy testvendor directory into directory for test
	tmpGopath, err := ioutil.TempDir("", "testvendor")
	if err != nil {
		t.Fatalf("unable to create testvendor directory: %v", err)
	}
	testVendorDir := path.Join(tmpGopath, "src", "github.com/testvendor")
	if err := os.MkdirAll(testVendorDir, 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	defer func() {
		os.RemoveAll(tmpGopath)
	}()

	if err := ioutil.WriteFile(path.Join(testVendorDir, "main.go"), []byte(testVendorMain), 0755); err != nil {
		t.Fatalf("Failed to write testvendor main: %v", err)
	}
	if err := os.MkdirAll(path.Join(testVendorDir, "vendor/github.com/testlog"), 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	if err := ioutil.WriteFile(path.Join(testVendorDir, "vendor/github.com/testlog/testlog.go"), []byte(testLog), 0755); err != nil {
		t.Fatalf("Failed to write testlog: %v", err)
	}

	cases := []struct {
		ignore          map[string]*regexp.Regexp
		numExpectedErrs int
	}{
		// basic case has one error
		{
			ignore:          nil,
			numExpectedErrs: 1,
		},
		// ignoring vendored import works
		{
			ignore: map[string]*regexp.Regexp{
				path.Join("github.com/testvendor/vendor/github.com/testlog"): regexp.MustCompile("Info"),
			},
		},
		// non-vendored path ignores vendored import
		{
			ignore: map[string]*regexp.Regexp{
				"github.com/testlog": regexp.MustCompile("Info"),
			},
		},
	}

	for i, currCase := range cases {
		checker := NewChecker()
		checker.Ignore = currCase.ignore
		loadPackages = func(cfg *packages.Config, paths ...string) ([]*packages.Package, error) {
			cfg.Env = append(os.Environ(), "GOPATH="+tmpGopath)
			cfg.Dir = testVendorDir
			pkgs, err := packages.Load(cfg, paths...)
			return pkgs, err
		}
		err := checker.CheckPackages("github.com/testvendor")

		if currCase.numExpectedErrs == 0 {
			if err != nil {
				t.Errorf("Case %d: expected no errors, but got: %v", i, err)
			}
			continue
		}

		uerr, ok := err.(*UncheckedErrors)
		if !ok {
			t.Errorf("Case %d: wrong error type returned: %v", i, err)
			continue
		}

		if currCase.numExpectedErrs != len(uerr.Errors) {
			t.Errorf("Case %d:\nExpected: %d errors\nActual:   %d errors", i, currCase.numExpectedErrs, len(uerr.Errors))
		}
	}
}

func TestWithoutGeneratedCode(t *testing.T) {
	const testVendorMain = `
	// Code generated by protoc-gen-go. DO NOT EDIT.
	package main

	import "github.com/testlog"

	func main() {
		// returns an error that is not checked
		testlog.Info()
	}`
	const testLog = `
	package testlog

	func Info() error {
		return nil
	}`

	if strings.HasPrefix(runtime.Version(), "go1.5") && os.Getenv("GO15VENDOREXPERIMENT") != "1" {
		// skip tests if running in go1.5 and vendoring is not enabled
		t.SkipNow()
	}

	// copy testvendor directory into directory for test
	tmpGopath, err := ioutil.TempDir("", "testvendor")
	if err != nil {
		t.Fatalf("unable to create testvendor directory: %v", err)
	}
	testVendorDir := path.Join(tmpGopath, "src", "github.com/testvendor")
	if err := os.MkdirAll(testVendorDir, 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	defer func() {
		os.RemoveAll(tmpGopath)
	}()

	if err := ioutil.WriteFile(path.Join(testVendorDir, "main.go"), []byte(testVendorMain), 0755); err != nil {
		t.Fatalf("Failed to write testvendor main: %v", err)
	}
	if err := os.MkdirAll(path.Join(testVendorDir, "vendor/github.com/testlog"), 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	if err := ioutil.WriteFile(path.Join(testVendorDir, "vendor/github.com/testlog/testlog.go"), []byte(testLog), 0755); err != nil {
		t.Fatalf("Failed to write testlog: %v", err)
	}

	cases := []struct {
		withoutGeneratedCode bool
		numExpectedErrs      int
	}{
		// basic case has one error
		{
			withoutGeneratedCode: false,
			numExpectedErrs:      1,
		},
		// ignoring generated code works
		{
			withoutGeneratedCode: true,
			numExpectedErrs:      0,
		},
	}

	for i, currCase := range cases {
		checker := NewChecker()
		checker.WithoutGeneratedCode = currCase.withoutGeneratedCode
		loadPackages = func(cfg *packages.Config, paths ...string) ([]*packages.Package, error) {
			cfg.Env = append(os.Environ(), "GOPATH="+tmpGopath)
			cfg.Dir = testVendorDir
			pkgs, err := packages.Load(cfg, paths...)
			return pkgs, err
		}
		err := checker.CheckPackages(path.Join("github.com/testvendor"))

		if currCase.numExpectedErrs == 0 {
			if err != nil {
				t.Errorf("Case %d: expected no errors, but got: %v", i, err)
			}
			continue
		}

		uerr, ok := err.(*UncheckedErrors)
		if !ok {
			t.Errorf("Case %d: wrong error type returned: %v", i, err)
			continue
		}

		if currCase.numExpectedErrs != len(uerr.Errors) {
			t.Errorf("Case %d:\nExpected: %d errors\nActual:   %d errors", i, currCase.numExpectedErrs, len(uerr.Errors))
		}
	}
}

func test(t *testing.T, f flags) {
	var (
		asserts bool = f&CheckAsserts != 0
		blank   bool = f&CheckBlank != 0
	)
	checker := NewChecker()
	checker.Asserts = asserts
	checker.Blank = blank
	checker.SetExclude(map[string]bool{
		fmt.Sprintf("(%s.ErrorMakerInterface).MakeNilError", testPackage): true,
	})
	err := checker.CheckPackages(testPackage)
	uerr, ok := err.(*UncheckedErrors)
	if !ok {
		t.Fatalf("wrong error type returned: %v", err)
	}

	numErrors := len(uncheckedMarkers)
	if blank {
		numErrors += len(blankMarkers)
	}
	if asserts {
		numErrors += len(assertMarkers)
	}

	if len(uerr.Errors) != numErrors {
		t.Errorf("got %d errors, want %d", len(uerr.Errors), numErrors)
	unchecked_loop:
		for k := range uncheckedMarkers {
			for _, e := range uerr.Errors {
				if newMarker(e) == k {
					continue unchecked_loop
				}
			}
			t.Errorf("Expected unchecked at %s", k)
		}
		if blank {
		blank_loop:
			for k := range blankMarkers {
				for _, e := range uerr.Errors {
					if newMarker(e) == k {
						continue blank_loop
					}
				}
				t.Errorf("Expected blank at %s", k)
			}
		}
		if asserts {
		assert_loop:
			for k := range assertMarkers {
				for _, e := range uerr.Errors {
					if newMarker(e) == k {
						continue assert_loop
					}
				}
				t.Errorf("Expected assert at %s", k)
			}
		}
	}

	for i, err := range uerr.Errors {
		m := marker{err.Pos.Filename, err.Pos.Line}
		if !uncheckedMarkers[m] && !blankMarkers[m] && !assertMarkers[m] {
			t.Errorf("%d: unexpected error: %v", i, err)
		}
	}
}
<commit_msg>Add test for build tags<commit_after>package errcheck

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"regexp"
	"runtime"
	"strings"
	"testing"

	"golang.org/x/tools/go/packages"
)

const testPackage = "github.com/kisielk/errcheck/testdata"

var (
	uncheckedMarkers map[marker]bool
	blankMarkers     map[marker]bool
	assertMarkers    map[marker]bool
)

type marker struct {
	file string
	line int
}

func newMarker(e UncheckedError) marker {
	return marker{e.Pos.Filename, e.Pos.Line}
}

func (m marker) String() string {
	return fmt.Sprintf("%s:%d", m.file, m.line)
}

func init() {
	uncheckedMarkers = make(map[marker]bool)
	blankMarkers = make(map[marker]bool)
	assertMarkers = make(map[marker]bool)

	cfg := &packages.Config{
		Mode:  packages.LoadSyntax,
		Tests: true,
	}
	pkgs, err := packages.Load(cfg, testPackage)
	if err != nil {
		panic("failed to import test package")
	}
	for _, pkg := range pkgs {
		for _, file := range pkg.Syntax {
			for _, comment := range file.Comments {
				text := comment.Text()
				pos := pkg.Fset.Position(comment.Pos())
				m := marker{pos.Filename, pos.Line}
				switch text {
				case "UNCHECKED\n":
					uncheckedMarkers[m] = true
				case "BLANK\n":
					blankMarkers[m] = true
				case "ASSERT\n":
					assertMarkers[m] = true
				}
			}
		}
	}
}

type flags uint

const (
	CheckAsserts flags = 1 << iota
	CheckBlank
)

// TestUnchecked runs a test against the example files and ensures all unchecked errors are caught.
func TestUnchecked(t *testing.T) {
	test(t, 0)
}

// TestBlank is like TestUnchecked but also ensures assignments to the blank identifier are caught.
func TestBlank(t *testing.T) {
	test(t, CheckBlank)
}

func TestAll(t *testing.T) {
	// TODO: CheckAsserts should work independently of CheckBlank
	test(t, CheckAsserts|CheckBlank)
}

func TestBuildTags(t *testing.T) {
	const (
		// uses "custom1" build tag and contains 1 unchecked error
		testBuildCustom1Tag = `
` + `// +build custom1

package custom

import "fmt"

func Print1() {
	// returns an error that is not checked
	fmt.Fprintln(nil)
}`
		// uses "custom2" build tag and contains 1 unchecked error
		testBuildCustom2Tag = `
` + `// +build custom2

package custom

import "fmt"

func Print2() {
	// returns an error that is not checked
	fmt.Fprintln(nil)
}`
		// included so that package is not empty when built without specifying tags
		testDoc = `
// Package custom contains code for testing build tags.
package custom
`
	)

	tmpGopath, err := ioutil.TempDir("", "testbuildtags")
	if err != nil {
		t.Fatalf("unable to create testbuildtags directory: %v", err)
	}
	testBuildTagsDir := path.Join(tmpGopath, "src", "github.com/testbuildtags")
	if err := os.MkdirAll(testBuildTagsDir, 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	defer func() {
		os.RemoveAll(tmpGopath)
	}()

	if err := ioutil.WriteFile(path.Join(testBuildTagsDir, "custom1.go"), []byte(testBuildCustom1Tag), 0644); err != nil {
		t.Fatalf("Failed to write testbuildtags custom1: %v", err)
	}
	if err := ioutil.WriteFile(path.Join(testBuildTagsDir, "custom2.go"), []byte(testBuildCustom2Tag), 0644); err != nil {
		t.Fatalf("Failed to write testbuildtags custom2: %v", err)
	}
	if err := ioutil.WriteFile(path.Join(testBuildTagsDir, "doc.go"), []byte(testDoc), 0644); err != nil {
		t.Fatalf("Failed to write testbuildtags doc: %v", err)
	}

	cases := []struct {
		tags            []string
		numExpectedErrs int
	}{
		// with no tags specified, main is ignored and there are no errors
		{
			tags:            nil,
			numExpectedErrs: 0,
		},
		// specifying "custom1" tag includes file with 1 error
		{
			tags:            []string{"custom1"},
			numExpectedErrs: 1,
		},
		// specifying "custom1" and "custom2" tags includes 2 files with 1 error each
		{
			tags:            []string{"custom1", "custom2"},
			numExpectedErrs: 2,
		},
	}

	for i, currCase := range cases {
		checker := NewChecker()
		checker.Tags = currCase.tags

		loadPackages = func(cfg *packages.Config, paths ...string) ([]*packages.Package, error) {
			cfg.Env = append(os.Environ(), "GOPATH="+tmpGopath)
			cfg.Dir = testBuildTagsDir
			pkgs, err := packages.Load(cfg, paths...)
			return pkgs, err
		}
		err := checker.CheckPackages("github.com/testbuildtags")

		if currCase.numExpectedErrs == 0 {
			if err != nil {
				t.Errorf("Case %d: expected no errors, but got: %v", i, err)
			}
			continue
		}

		uerr, ok := err.(*UncheckedErrors)
		if !ok {
			t.Errorf("Case %d: wrong error type returned: %v", i, err)
			continue
		}

		if currCase.numExpectedErrs != len(uerr.Errors) {
			t.Errorf("Case %d:\nExpected: %d errors\nActual:   %d errors", i, currCase.numExpectedErrs, len(uerr.Errors))
		}
	}
}

func TestWhitelist(t *testing.T) {

}

func TestIgnore(t *testing.T) {
	const testVendorMain = `
	package main

	import "github.com/testlog"

	func main() {
		// returns an error that is not checked
		testlog.Info()
	}`
	const testLog = `
	package testlog

	func Info() error {
		return nil
	}`

	if strings.HasPrefix(runtime.Version(), "go1.5") && os.Getenv("GO15VENDOREXPERIMENT") != "1" {
		// skip tests if running in go1.5 and vendoring is not enabled
		t.SkipNow()
	}

	// copy testvendor directory into directory for test
	tmpGopath, err := ioutil.TempDir("", "testvendor")
	if err != nil {
		t.Fatalf("unable to create testvendor directory: %v", err)
	}
	testVendorDir := path.Join(tmpGopath, "src", "github.com/testvendor")
	if err := os.MkdirAll(testVendorDir, 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	defer func() {
		os.RemoveAll(tmpGopath)
	}()

	if err := ioutil.WriteFile(path.Join(testVendorDir, "main.go"), []byte(testVendorMain), 0755); err != nil {
		t.Fatalf("Failed to write testvendor main: %v", err)
	}
	if err := os.MkdirAll(path.Join(testVendorDir, "vendor/github.com/testlog"), 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	if err := ioutil.WriteFile(path.Join(testVendorDir, "vendor/github.com/testlog/testlog.go"), []byte(testLog), 0755); err != nil {
		t.Fatalf("Failed to write testlog: %v", err)
	}

	cases := []struct {
		ignore          map[string]*regexp.Regexp
		numExpectedErrs int
	}{
		// basic case has one error
		{
			ignore:          nil,
			numExpectedErrs: 1,
		},
		// ignoring vendored import works
		{
			ignore: map[string]*regexp.Regexp{
				path.Join("github.com/testvendor/vendor/github.com/testlog"): regexp.MustCompile("Info"),
			},
		},
		// non-vendored path ignores vendored import
		{
			ignore: map[string]*regexp.Regexp{
				"github.com/testlog": regexp.MustCompile("Info"),
			},
		},
	}

	for i, currCase := range cases {
		checker := NewChecker()
		checker.Ignore = currCase.ignore
		loadPackages = func(cfg *packages.Config, paths ...string) ([]*packages.Package, error) {
			cfg.Env = append(os.Environ(), "GOPATH="+tmpGopath)
			cfg.Dir = testVendorDir
			pkgs, err := packages.Load(cfg, paths...)
			return pkgs, err
		}
		err := checker.CheckPackages("github.com/testvendor")

		if currCase.numExpectedErrs == 0 {
			if err != nil {
				t.Errorf("Case %d: expected no errors, but got: %v", i, err)
			}
			continue
		}

		uerr, ok := err.(*UncheckedErrors)
		if !ok {
			t.Errorf("Case %d: wrong error type returned: %v", i, err)
			continue
		}

		if currCase.numExpectedErrs != len(uerr.Errors) {
			t.Errorf("Case %d:\nExpected: %d errors\nActual:   %d errors", i, currCase.numExpectedErrs, len(uerr.Errors))
		}
	}
}

func TestWithoutGeneratedCode(t *testing.T) {
	const testVendorMain = `
	// Code generated by protoc-gen-go. DO NOT EDIT.
	package main

	import "github.com/testlog"

	func main() {
		// returns an error that is not checked
		testlog.Info()
	}`
	const testLog = `
	package testlog

	func Info() error {
		return nil
	}`

	if strings.HasPrefix(runtime.Version(), "go1.5") && os.Getenv("GO15VENDOREXPERIMENT") != "1" {
		// skip tests if running in go1.5 and vendoring is not enabled
		t.SkipNow()
	}

	// copy testvendor directory into directory for test
	tmpGopath, err := ioutil.TempDir("", "testvendor")
	if err != nil {
		t.Fatalf("unable to create testvendor directory: %v", err)
	}
	testVendorDir := path.Join(tmpGopath, "src", "github.com/testvendor")
	if err := os.MkdirAll(testVendorDir, 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	defer func() {
		os.RemoveAll(tmpGopath)
	}()

	if err := ioutil.WriteFile(path.Join(testVendorDir, "main.go"), []byte(testVendorMain), 0755); err != nil {
		t.Fatalf("Failed to write testvendor main: %v", err)
	}
	if err := os.MkdirAll(path.Join(testVendorDir, "vendor/github.com/testlog"), 0755); err != nil {
		t.Fatalf("MkdirAll failed: %v", err)
	}
	if err := ioutil.WriteFile(path.Join(testVendorDir, "vendor/github.com/testlog/testlog.go"), []byte(testLog), 0755); err != nil {
		t.Fatalf("Failed to write testlog: %v", err)
	}

	cases := []struct {
		withoutGeneratedCode bool
		numExpectedErrs      int
	}{
		// basic case has one error
		{
			withoutGeneratedCode: false,
			numExpectedErrs:      1,
		},
		// ignoring generated code works
		{
			withoutGeneratedCode: true,
			numExpectedErrs:      0,
		},
	}

	for i, currCase := range cases {
		checker := NewChecker()
		checker.WithoutGeneratedCode = currCase.withoutGeneratedCode
		loadPackages = func(cfg *packages.Config, paths ...string) ([]*packages.Package, error) {
			cfg.Env = append(os.Environ(), "GOPATH="+tmpGopath)
			cfg.Dir = testVendorDir
			pkgs, err := packages.Load(cfg, paths...)
			return pkgs, err
		}
		err := checker.CheckPackages(path.Join("github.com/testvendor"))

		if currCase.numExpectedErrs == 0 {
			if err != nil {
				t.Errorf("Case %d: expected no errors, but got: %v", i, err)
			}
			continue
		}

		uerr, ok := err.(*UncheckedErrors)
		if !ok {
			t.Errorf("Case %d: wrong error type returned: %v", i, err)
			continue
		}

		if currCase.numExpectedErrs != len(uerr.Errors) {
			t.Errorf("Case %d:\nExpected: %d errors\nActual:   %d errors", i, currCase.numExpectedErrs, len(uerr.Errors))
		}
	}
}

func test(t *testing.T, f flags) {
	var (
		asserts bool = f&CheckAsserts != 0
		blank   bool = f&CheckBlank != 0
	)
	checker := NewChecker()
	checker.Asserts = asserts
	checker.Blank = blank
	checker.SetExclude(map[string]bool{
		fmt.Sprintf("(%s.ErrorMakerInterface).MakeNilError", testPackage): true,
	})
	err := checker.CheckPackages(testPackage)
	uerr, ok := err.(*UncheckedErrors)
	if !ok {
		t.Fatalf("wrong error type returned: %v", err)
	}

	numErrors := len(uncheckedMarkers)
	if blank {
		numErrors += len(blankMarkers)
	}
	if asserts {
		numErrors += len(assertMarkers)
	}

	if len(uerr.Errors) != numErrors {
		t.Errorf("got %d errors, want %d", len(uerr.Errors), numErrors)
	unchecked_loop:
		for k := range uncheckedMarkers {
			for _, e := range uerr.Errors {
				if newMarker(e) == k {
					continue unchecked_loop
				}
			}
			t.Errorf("Expected unchecked at %s", k)
		}
		if blank {
		blank_loop:
			for k := range blankMarkers {
				for _, e := range uerr.Errors {
					if newMarker(e) == k {
						continue blank_loop
					}
				}
				t.Errorf("Expected blank at %s", k)
			}
		}
		if asserts {
		assert_loop:
			for k := range assertMarkers {
				for _, e := range uerr.Errors {
					if newMarker(e) == k {
						continue assert_loop
					}
				}
				t.Errorf("Expected assert at %s", k)
			}
		}
	}

	for i, err := range uerr.Errors {
		m := marker{err.Pos.Filename, err.Pos.Line}
		if !uncheckedMarkers[m] && !blankMarkers[m] && !assertMarkers[m] {
			t.Errorf("%d: unexpected error: %v", i, err)
		}
	}
}
<|endoftext|>


<commit_before>package plugin

//go:generate goderive .

import (
	"crypto/tls"
	"fmt"
	"mime"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"

	"github.com/dcos/dcos-cli/pkg/config"
	"github.com/dcos/dcos-cli/pkg/fsutil"
	"github.com/dcos/dcos-cli/pkg/httpclient"
	"github.com/pelletier/go-toml"
	"github.com/sirupsen/logrus"
	"github.com/spf13/afero"
)

// Manager retrieves the plugins available for the current cluster
// by navigating into the filesystem.
type Manager struct {
	fs      afero.Fs
	logger  *logrus.Logger
	cluster *config.Cluster
}

// NewManager returns a new plugin manager.
func NewManager(fs afero.Fs, logger *logrus.Logger) *Manager {
	return &Manager{
		fs:     fs,
		logger: logger,
	}
}

// InstallOpts are installation options for plugin resources.
type InstallOpts struct {
	// Name specify the name of the plugin.
	Name string

	// Update allows to potentially overwrite an already existing plugin of the same name.
	Update bool

	// PostInstall is a hook which can be invoked after plugin installation.
	// It is invoked right before the plugin is moved to its final location.
	PostInstall func(fs afero.Fs, pluginDir string) error

	path       string
	stagingDir string
}

// Install installs a plugin from a resource.
func (m *Manager) Install(resource string, installOpts *InstallOpts) (err error) {
	// If it's a remote resource, download it first.
	if strings.HasPrefix(resource, "https://") || strings.HasPrefix(resource, "http://") {
		installOpts.path, err = m.downloadPlugin(resource)
		if err != nil {
			return err
		}
	}

	// The staging dir is where the plugin will be constructed
	// before eventually getting moved to its final location.
	installOpts.stagingDir, err = afero.TempDir(m.fs, os.TempDir(), "dcos-cli")
	if err != nil {
		return err
	}

	// Build the plugin into the staging directory.
	err = m.buildPlugin(installOpts)
	if err != nil {
		return err
	}
	return m.installPlugin(installOpts)
}

// SetCluster sets the plugin manager's target cluster.
func (m *Manager) SetCluster(cluster *config.Cluster) {
	m.cluster = cluster
}

// Remove removes a plugin from the filesystem.
func (m *Manager) Remove(name string) error {
	pluginDir := filepath.Join(m.pluginsDir(), name)
	pluginDirExists, err := afero.DirExists(m.fs, pluginDir)
	if err != nil {
		return err
	}
	if !pluginDirExists {
		return fmt.Errorf("'%s' is not a plugin directory", pluginDir)
	}
	return m.fs.RemoveAll(pluginDir)
}

// Plugins returns the plugins associated with the current cluster.
func (m *Manager) Plugins() (plugins []*Plugin) {
	pluginDirs, err := afero.ReadDir(m.fs, m.pluginsDir())
	if err != nil {
		m.logger.Debugf("Couldn't open plugin dir: %s", err)
		return plugins
	}

	for _, pluginDir := range pluginDirs {
		if !pluginDir.IsDir() {
			continue
		}
		plugin, err := m.loadPlugin(pluginDir.Name())
		if err != nil {
			// We don't want to see the CLI failing if a single plugin is malformed.
			// We thus log the error but continue if there is an issue at that step.
			m.logger.Debugf("Couldn't load plugin: %s", err)
			continue
		}
		plugins = append(plugins, plugin)
	}
	return plugins
}

// loadPlugin loads a plugin based on its name.
func (m *Manager) loadPlugin(name string) (*Plugin, error) {
	m.logger.Infof("Loading plugin '%s'...", name)

	plugin := &Plugin{Name: name}
	pluginPath := filepath.Join(m.pluginsDir(), name, "env")
	pluginFilePath := filepath.Join(pluginPath, "plugin.toml")

	if err := m.unmarshalPlugin(plugin, pluginFilePath); err != nil {
		return nil, err
	}

	// Save a deep copy of the plugin once it's loaded from the plugin.toml file.
	persistedPlugin := &Plugin{}
	deriveDeepCopy(persistedPlugin, plugin)

	if len(plugin.Commands) == 0 {
		plugin.Commands = m.findCommands(pluginPath)
	}

	// Normalize plugin commands by putting binary full paths and description summaries.
	for i, cmd := range plugin.Commands {
		if !filepath.IsAbs(cmd.Path) {
			cmd.Path = filepath.Join(pluginPath, cmd.Path)
		}
		if cmd.Description == "" {
			cmd.Description = m.commandDescription(cmd)
		}
		plugin.Commands[i] = cmd
	}

	// Compare the normalized plugin with the saved copy to know whether or not the file should be updated.
	if !reflect.DeepEqual(persistedPlugin, plugin) {
		m.persistPlugin(plugin, pluginFilePath)
	}
	return plugin, nil
}

// findCommands discovers commands in a given directory according to conventions.
// Each command should be contained in a dedicated binary named `dcos-{command}`.
// On Windows it must have the `.exe`` extension.
func (m *Manager) findCommands(pluginDir string) (commands []Command) {
	binDir := filepath.Join(pluginDir, "bin")
	if runtime.GOOS == "windows" {
		// On Windows we check for the `bin` dir. If it doesn't exist we use the legacy `Scripts` dir.
		binDirExists, err := afero.DirExists(m.fs, binDir)
		if err != nil || !binDirExists {
			binDir = filepath.Join(pluginDir, "Scripts")
		}
	}

	m.logger.Debugf("Discovering commands in '%s'...", binDir)

	binaries, err := afero.ReadDir(m.fs, binDir)
	if err != nil {
		m.logger.Debug(err)
		return nil
	}

	for _, binary := range binaries {
		if !strings.HasPrefix(binary.Name(), "dcos-") {
			continue
		}
		cmd := Command{
			Path: filepath.Join(binDir, binary.Name()),
			Name: strings.TrimPrefix(binary.Name(), "dcos-"),
		}
		if runtime.GOOS == "windows" {
			cmd.Name = strings.TrimSuffix(cmd.Name, ".exe")
		}
		m.logger.Debugf("Discovered '%s' command", cmd.Name)
		commands = append(commands, cmd)
	}
	return commands
}

// commandDescription gets the command info summary by invoking the binary with the `--info` flag.
func (m *Manager) commandDescription(cmd Command) (desc string) {
	infoCmd, err := exec.Command(cmd.Path, cmd.Name, "--info").Output()
	if err != nil {
		m.logger.Debugf("Couldn't get info summary for the '%s' command: %s", cmd.Name, err)
	} else {
		desc = strings.TrimSpace(string(infoCmd))
	}
	return desc
}

// unmarshalPlugin unmarshals a `plugin.toml` file into a Plugin structure.
func (m *Manager) unmarshalPlugin(plugin *Plugin, path string) error {
	data, err := afero.ReadFile(m.fs, path)
	if err != nil {
		m.logger.Debugf("Couldn't open plugin.toml: %s", err)
		return nil
	}
	return toml.Unmarshal(data, plugin)
}

// persistPlugin saves a `plugin.toml` file representing the plugin.
func (m *Manager) persistPlugin(plugin *Plugin, path string) {
	pluginTOML, err := toml.Marshal(*plugin)
	if err == nil {
		err = afero.WriteFile(m.fs, path, pluginTOML, 0644)
	}
	if err != nil {
		m.logger.Debug(err)
	}
}

// pluginsDir returns the path to the plugins directory.
func (m *Manager) pluginsDir() string {
	return filepath.Join(m.cluster.Dir(), "subcommands")
}

// downloadPlugin downloads a plugin and returns the path to the temporary file it stored it to.
func (m *Manager) downloadPlugin(url string) (string, error) {
	tmpDir, err := afero.TempDir(m.fs, os.TempDir(), "dcos-cli")
	if err != nil {
		return "", err
	}

	resp, err := m.httpClient(url).Get(url)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	downloadedFilePath := filepath.Join(tmpDir, m.downloadFilename(resp))

	if err := fsutil.CopyReader(m.fs, resp.Body, downloadedFilePath, 0644); err != nil {
		return "", err
	}
	return downloadedFilePath, nil
}

// downloadFilename picks a filename for the resource to download. It first reads the
// `Content-Disposition` header, when not set it defaults to the URL path basename.
func (m *Manager) downloadFilename(resp *http.Response) string {
	contentDisposition := resp.Header.Get("Content-Disposition")
	_, dispositionParams, err := mime.ParseMediaType(contentDisposition)
	if err != nil {
		m.logger.Debug(err)
	} else if filename, ok := dispositionParams["filename"]; ok {
		return filename
	}
	return path.Base(resp.Request.URL.Path)
}

// buildPlugin constructs the plugin structure within a target directory, based on a path to a plugin resource.
func (m *Manager) buildPlugin(installOpts *InstallOpts) error {
	plugin := &Plugin{}

	// Detect the plugin resource media type.
	contentType, err := fsutil.DetectMediaType(m.fs, installOpts.path)
	if err != nil {
		return err
	}

	switch contentType {
	case "application/zip":
		// Unzip the plugin into the staging dir and validate its plugin.toml, if any.
		if err := fsutil.Unzip(m.fs, installOpts.path, installOpts.stagingDir); err != nil {
			return err
		}
		pluginFilePath := filepath.Join(installOpts.stagingDir, "plugin.toml")
		if err := m.unmarshalPlugin(plugin, pluginFilePath); err != nil {
			return err
		}
		if plugin.Name != "" {
			installOpts.Name = plugin.Name
		}

	// The current media type detection mechanism (based on the stdlib) cannot
	// detect binary executables. Thus, we assume the resource is an actual binary
	// when it's not a ZIP archive.
	default:
		// Copy the binary into the staging dir's bin folder,
		binDir := filepath.Join(installOpts.stagingDir, "bin")
		if err := m.fs.MkdirAll(binDir, 0755); err != nil {
			return err
		}
		binPath := filepath.Join(binDir, filepath.Base(installOpts.path))
		err := fsutil.Copy(m.fs, installOpts.path, binPath, 0751)
		if err != nil {
			return err
		}
	}

	// If there is no plugin name, use the resource file basename.
	if installOpts.Name == "" {
		basename := filepath.Base(installOpts.path)
		installOpts.Name = strings.TrimSuffix(basename, filepath.Ext(basename))
	}
	if installOpts.PostInstall != nil {
		return installOpts.PostInstall(m.fs, installOpts.stagingDir)
	}
	return nil
}

// installPlugin installs a plugin from a staging dir into its final location.
// "update" indicates whether an already existing plugin can be overwritten.
func (m *Manager) installPlugin(installOpts *InstallOpts) error {
	dest := filepath.Join(m.pluginsDir(), installOpts.Name, "env")

	if installOpts.Update {
		if err := m.fs.RemoveAll(dest); err != nil {
			return err
		}
	} else {
		pluginDirExists, err := afero.DirExists(m.fs, dest)
		if err != nil {
			m.logger.Debug(err)
		}
		if pluginDirExists {
			return fmt.Errorf("'%s' is already installed", installOpts.Name)
		}
	}

	// Move the plugin folder to the final location.
	if err := m.fs.MkdirAll(filepath.Dir(dest), 0755); err != nil {
		return err
	}
	return m.fs.Rename(installOpts.stagingDir, dest)
}

// httpClient returns the appropriate HTTP client for a given resource.
func (m *Manager) httpClient(url string) *httpclient.Client {
	httpOpts := []httpclient.Option{
		httpclient.Logger(m.logger),
		httpclient.FailOnErrStatus(true),
	}
	if strings.HasPrefix(url, m.cluster.URL()) {
		httpOpts = append(
			httpOpts,
			httpclient.ACSToken(m.cluster.ACSToken()),
			httpclient.TLS(&tls.Config{
				InsecureSkipVerify: m.cluster.TLS().Insecure,
				RootCAs:            m.cluster.TLS().RootCAs,
			}),
		)
	}
	return httpclient.New("", httpOpts...)
}
<commit_msg>Package list fix (#1304)<commit_after>package plugin

//go:generate goderive .

import (
	"crypto/tls"
	"fmt"
	"mime"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"

	"github.com/dcos/dcos-cli/pkg/config"
	"github.com/dcos/dcos-cli/pkg/fsutil"
	"github.com/dcos/dcos-cli/pkg/httpclient"
	"github.com/pelletier/go-toml"
	"github.com/sirupsen/logrus"
	"github.com/spf13/afero"
)

// Manager retrieves the plugins available for the current cluster
// by navigating into the filesystem.
type Manager struct {
	fs      afero.Fs
	logger  *logrus.Logger
	cluster *config.Cluster
}

// NewManager returns a new plugin manager.
func NewManager(fs afero.Fs, logger *logrus.Logger) *Manager {
	return &Manager{
		fs:     fs,
		logger: logger,
	}
}

// InstallOpts are installation options for plugin resources.
type InstallOpts struct {
	// Name specify the name of the plugin.
	Name string

	// Update allows to potentially overwrite an already existing plugin of the same name.
	Update bool

	// PostInstall is a hook which can be invoked after plugin installation.
	// It is invoked right before the plugin is moved to its final location.
	PostInstall func(fs afero.Fs, pluginDir string) error

	path       string
	stagingDir string
}

// Install installs a plugin from a resource.
func (m *Manager) Install(resource string, installOpts *InstallOpts) (err error) {
	// If it's a remote resource, download it first.
	if strings.HasPrefix(resource, "https://") || strings.HasPrefix(resource, "http://") {
		installOpts.path, err = m.downloadPlugin(resource)
		if err != nil {
			return err
		}
	}

	// The staging dir is where the plugin will be constructed
	// before eventually getting moved to its final location.
	installOpts.stagingDir, err = afero.TempDir(m.fs, os.TempDir(), "dcos-cli")
	if err != nil {
		return err
	}

	// Build the plugin into the staging directory.
	err = m.buildPlugin(installOpts)
	if err != nil {
		return err
	}
	return m.installPlugin(installOpts)
}

// SetCluster sets the plugin manager's target cluster.
func (m *Manager) SetCluster(cluster *config.Cluster) {
	m.cluster = cluster
}

// Remove removes a plugin from the filesystem.
func (m *Manager) Remove(name string) error {
	pluginDir := filepath.Join(m.pluginsDir(), name)
	pluginDirExists, err := afero.DirExists(m.fs, pluginDir)
	if err != nil {
		return err
	}
	if !pluginDirExists {
		return fmt.Errorf("'%s' is not a plugin directory", pluginDir)
	}
	return m.fs.RemoveAll(pluginDir)
}

// Plugins returns the plugins associated with the current cluster.
func (m *Manager) Plugins() (plugins []*Plugin) {
	pluginDirs, err := afero.ReadDir(m.fs, m.pluginsDir())
	if err != nil {
		m.logger.Debugf("Couldn't open plugin dir: %s", err)
		return plugins
	}

	for _, pluginDir := range pluginDirs {
		if !pluginDir.IsDir() {
			continue
		}
		plugin, err := m.loadPlugin(pluginDir.Name())
		if err != nil {
			// We don't want to see the CLI failing if a single plugin is malformed.
			// We thus log the error but continue if there is an issue at that step.
			m.logger.Debugf("Couldn't load plugin: %s", err)
			continue
		}
		plugins = append(plugins, plugin)
	}
	return plugins
}

// loadPlugin loads a plugin based on its name.
func (m *Manager) loadPlugin(name string) (*Plugin, error) {
	m.logger.Infof("Loading plugin '%s'...", name)

	plugin := &Plugin{Name: name}
	pluginPath := filepath.Join(m.pluginsDir(), name, "env")
	pluginFilePath := filepath.Join(pluginPath, "plugin.toml")

	if err := m.unmarshalPlugin(plugin, pluginFilePath); err != nil {
		return nil, err
	}

	// Save a deep copy of the plugin once it's loaded from the plugin.toml file.
	persistedPlugin := &Plugin{}
	deriveDeepCopy(persistedPlugin, plugin)

	if len(plugin.Commands) == 0 {
		plugin.Commands = m.findCommands(pluginPath)
	}

	// Normalize plugin commands by putting binary full paths and description summaries.
	for i, cmd := range plugin.Commands {
		if !filepath.IsAbs(cmd.Path) {
			cmd.Path = filepath.Join(pluginPath, cmd.Path)
		}
		if cmd.Description == "" {
			cmd.Description = m.commandDescription(cmd)
		}
		plugin.Commands[i] = cmd
	}

	// Compare the normalized plugin with the saved copy to know whether or not the file should be updated.
	if !reflect.DeepEqual(persistedPlugin, plugin) {
		m.persistPlugin(plugin, pluginFilePath)
	}
	return plugin, nil
}

// findCommands discovers commands in a given directory according to conventions.
// Each command should be contained in a dedicated binary named `dcos-{command}`.
// On Windows it must have the `.exe`` extension.
func (m *Manager) findCommands(pluginDir string) (commands []Command) {
	binDir := filepath.Join(pluginDir, "bin")
	if runtime.GOOS == "windows" {
		// On Windows we check for the `bin` dir. If it doesn't exist we use the legacy `Scripts` dir.
		binDirExists, err := afero.DirExists(m.fs, binDir)
		if err != nil || !binDirExists {
			binDir = filepath.Join(pluginDir, "Scripts")
		}
	}

	m.logger.Debugf("Discovering commands in '%s'...", binDir)

	binaries, err := afero.ReadDir(m.fs, binDir)
	if err != nil {
		m.logger.Debug(err)
		return nil
	}

	for _, binary := range binaries {
		if !strings.HasPrefix(binary.Name(), "dcos-") {
			continue
		}
		cmd := Command{
			Path: filepath.Join(binDir, binary.Name()),
			Name: strings.TrimPrefix(binary.Name(), "dcos-"),
		}
		if runtime.GOOS == "windows" {
			cmd.Name = strings.TrimSuffix(cmd.Name, ".exe")
		}
		m.logger.Debugf("Discovered '%s' command", cmd.Name)
		commands = append(commands, cmd)
	}
	return commands
}

// commandDescription gets the command info summary by invoking the binary with the `--info` flag.
func (m *Manager) commandDescription(cmd Command) (desc string) {
	infoCmd, err := exec.Command(cmd.Path, cmd.Name, "--info").Output()
	if err != nil {
		m.logger.Debugf("Couldn't get info summary for the '%s' command: %s", cmd.Name, err)
	} else {
		desc = strings.TrimSpace(string(infoCmd))
	}
	return desc
}

// unmarshalPlugin unmarshals a `plugin.toml` file into a Plugin structure.
func (m *Manager) unmarshalPlugin(plugin *Plugin, path string) error {
	data, err := afero.ReadFile(m.fs, path)
	if err != nil {
		m.logger.Debugf("Couldn't open plugin.toml: %s", err)
		return nil
	}
	return toml.Unmarshal(data, plugin)
}

// persistPlugin saves a `plugin.toml` file representing the plugin.
func (m *Manager) persistPlugin(plugin *Plugin, path string) {
	pluginTOML, err := toml.Marshal(*plugin)
	if err == nil {
		err = afero.WriteFile(m.fs, path, pluginTOML, 0644)
	}
	if err != nil {
		m.logger.Debug(err)
	}
}

// pluginsDir returns the path to the plugins directory.
func (m *Manager) pluginsDir() string {
	return filepath.Join(m.cluster.Dir(), "subcommands")
}

// downloadPlugin downloads a plugin and returns the path to the temporary file it stored it to.
func (m *Manager) downloadPlugin(url string) (string, error) {
	tmpDir, err := afero.TempDir(m.fs, os.TempDir(), "dcos-cli")
	if err != nil {
		return "", err
	}

	resp, err := m.httpClient(url).Get(url)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	downloadedFilePath := filepath.Join(tmpDir, m.downloadFilename(resp))

	if err := fsutil.CopyReader(m.fs, resp.Body, downloadedFilePath, 0644); err != nil {
		return "", err
	}
	return downloadedFilePath, nil
}

// downloadFilename picks a filename for the resource to download. It first reads the
// `Content-Disposition` header, when not set it defaults to the URL path basename.
func (m *Manager) downloadFilename(resp *http.Response) string {
	contentDisposition := resp.Header.Get("Content-Disposition")
	_, dispositionParams, err := mime.ParseMediaType(contentDisposition)
	if err != nil {
		m.logger.Debug(err)
	} else if filename, ok := dispositionParams["filename"]; ok {
		return filename
	}
	return path.Base(resp.Request.URL.Path)
}

// buildPlugin constructs the plugin structure within a target directory, based on a path to a plugin resource.
func (m *Manager) buildPlugin(installOpts *InstallOpts) error {
	plugin := &Plugin{}

	// Detect the plugin resource media type.
	contentType, err := fsutil.DetectMediaType(m.fs, installOpts.path)
	if err != nil {
		return err
	}

	envDir := filepath.Join(installOpts.stagingDir, "env")

	switch contentType {
	case "application/zip":
		// Unzip the plugin into the staging dir and validate its plugin.toml, if any.
		if err := fsutil.Unzip(m.fs, installOpts.path, envDir); err != nil {
			return err
		}
		pluginFilePath := filepath.Join(envDir, "plugin.toml")
		if err := m.unmarshalPlugin(plugin, pluginFilePath); err != nil {
			return err
		}
		if plugin.Name != "" {
			installOpts.Name = plugin.Name
		}

	// The current media type detection mechanism (based on the stdlib) cannot
	// detect binary executables. Thus, we assume the resource is an actual binary
	// when it's not a ZIP archive.
	default:
		// Copy the binary into the staging dir's bin folder,
		binDir := filepath.Join(envDir, "bin")
		if err := m.fs.MkdirAll(binDir, 0755); err != nil {
			return err
		}
		binPath := filepath.Join(binDir, filepath.Base(installOpts.path))
		err := fsutil.Copy(m.fs, installOpts.path, binPath, 0751)
		if err != nil {
			return err
		}
	}

	// If there is no plugin name, use the resource file basename.
	if installOpts.Name == "" {
		basename := filepath.Base(installOpts.path)
		installOpts.Name = strings.TrimSuffix(basename, filepath.Ext(basename))
	}
	if installOpts.PostInstall != nil {
		return installOpts.PostInstall(m.fs, installOpts.stagingDir)
	}
	return nil
}

// installPlugin installs a plugin from a staging dir into its final location.
// "update" indicates whether an already existing plugin can be overwritten.
func (m *Manager) installPlugin(installOpts *InstallOpts) error {
	dest := filepath.Join(m.pluginsDir(), installOpts.Name)

	if installOpts.Update {
		if err := m.fs.RemoveAll(dest); err != nil {
			return err
		}
	} else {
		pluginDirExists, err := afero.DirExists(m.fs, dest)
		if err != nil {
			m.logger.Debug(err)
		}
		if pluginDirExists {
			return fmt.Errorf("'%s' is already installed", installOpts.Name)
		}
	}

	// Move the plugin folder to the final location.
	if err := m.fs.MkdirAll(filepath.Dir(dest), 0755); err != nil {
		return err
	}
	return m.fs.Rename(installOpts.stagingDir, dest)
}

// httpClient returns the appropriate HTTP client for a given resource.
func (m *Manager) httpClient(url string) *httpclient.Client {
	httpOpts := []httpclient.Option{
		httpclient.Logger(m.logger),
		httpclient.FailOnErrStatus(true),
	}
	if strings.HasPrefix(url, m.cluster.URL()) {
		httpOpts = append(
			httpOpts,
			httpclient.ACSToken(m.cluster.ACSToken()),
			httpclient.TLS(&tls.Config{
				InsecureSkipVerify: m.cluster.TLS().Insecure,
				RootCAs:            m.cluster.TLS().RootCAs,
			}),
		)
	}
	return httpclient.New("", httpOpts...)
}
<|endoftext|>


<commit_before>/*******************************************************************************
*
* Copyright 2017 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, 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 plugins

import (
	"fmt"

	"github.com/gophercloud/gophercloud"
	"github.com/gophercloud/gophercloud/openstack"
	"github.com/gophercloud/gophercloud/openstack/sharedfilesystems/v2/sharenetworks"
	"github.com/gophercloud/gophercloud/pagination"
	"github.com/sapcc/go-bits/logg"
	"github.com/sapcc/limes/pkg/limes"
)

type manilaPlugin struct {
	cfg limes.ServiceConfiguration
}

var manilaResources = []limes.ResourceInfo{
	{
		Name: "share_networks",
		Unit: limes.UnitNone,
	},
	{
		Name: "share_capacity",
		Unit: limes.UnitGibibytes,
	},
	{
		Name: "shares",
		Unit: limes.UnitNone,
	},
	{
		Name: "snapshot_capacity",
		Unit: limes.UnitGibibytes,
	},
	{
		Name: "share_snapshots",
		Unit: limes.UnitNone,
	},
}

func init() {
	limes.RegisterQuotaPlugin(func(c limes.ServiceConfiguration, scrapeSubresources map[string]bool) limes.QuotaPlugin {
		return &manilaPlugin{c}
	})
}

//Init implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) Init(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts) error {
	return nil
}

//ServiceInfo implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) ServiceInfo() limes.ServiceInfo {
	return limes.ServiceInfo{
		Type:        "sharev2",
		ProductName: "manila",
		Area:        "storage",
	}
}

//Resources implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) Resources() []limes.ResourceInfo {
	return manilaResources
}

//Scrape implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) Scrape(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts, clusterID, domainUUID, projectUUID string) (map[string]limes.ResourceData, error) {
	client, err := openstack.NewSharedFileSystemV2(provider, eo)
	if err != nil {
		return nil, err
	}

	var result gophercloud.Result
	var totalShareUsage, totalSnapshotUsage, totalShareNetworksUsage = uint64(0), uint64(0), uint64(0)

	//Get absolute quota limits per project
	url := client.ServiceURL("os-quota-sets", projectUUID)
	_, err = client.Get(url, &result.Body, nil)
	if err != nil {
		return nil, err
	}
	var manilaQuotaData struct {
		QuotaSet struct {
			Gigabytes         int64 `json:"gigabytes"`
			Shares            int64 `json:"shares"`
			SnapshotGigabytes int64 `json:"snapshot_gigabytes"`
			Snapshots         int64 `json:"snapshots"`
			ShareNetworks     int64 `json:"share_networks"`
		} `json:"quota_set"`
	}
	err = result.ExtractInto(&manilaQuotaData)
	if err != nil {
		return nil, err
	}

	//Get usage of shares per project
	url = client.ServiceURL("shares", "detail") + fmt.Sprintf("?project_id=%s&all_tenants=1", projectUUID)
	_, err = client.Get(url, &result.Body, nil)
	if err != nil {
		return nil, err
	}

	var manilaShareUsageData struct {
		Shares []struct {
			Size uint64 `json:"size"`
		} `json:"shares"`
	}
	err = result.ExtractInto(&manilaShareUsageData)
	if err != nil {
		return nil, err
	}

	for _, element := range manilaShareUsageData.Shares {
		totalShareUsage += element.Size
	}

	//Get usage of snapshots per project
	url = client.ServiceURL("snapshots", "detail") + fmt.Sprintf("?project_id=%s&all_tenants=1", projectUUID)
	_, err = client.Get(url, &result.Body, nil)
	if err != nil {
		return nil, err
	}

	var manilaSnapshotUsageData struct {
		Snapshots []struct {
			ShareSize uint64 `json:"share_size"`
		} `json:"snapshots"`
	}
	err = result.ExtractInto(&manilaSnapshotUsageData)
	if err != nil {
		return nil, err
	}

	for _, element := range manilaSnapshotUsageData.Snapshots {
		totalSnapshotUsage += element.ShareSize
	}

	//Get usage of shared networks
	pages := 0
	sharenetworks.ListDetail(client, sharenetworks.ListOpts{ProjectID: projectUUID}).EachPage(func(page pagination.Page) (bool, error) {
		pages++
		sn, err := sharenetworks.ExtractShareNetworks(page)
		if err != nil {
			return false, err
		}
		totalShareNetworksUsage = uint64(len(sn))
		return true, nil
	})

	logg.Debug("Scraped quota and usage for service: sharev2.")

	return map[string]limes.ResourceData{
		"shares": {
			Quota: manilaQuotaData.QuotaSet.Shares,
			Usage: uint64(len(manilaShareUsageData.Shares)),
		},
		"share_snapshots": {
			Quota: manilaQuotaData.QuotaSet.Snapshots,
			Usage: uint64(len(manilaSnapshotUsageData.Snapshots)),
		},
		"share_networks": {
			Quota: manilaQuotaData.QuotaSet.ShareNetworks,
			Usage: uint64(totalShareNetworksUsage),
		},
		"share_capacity": {
			Quota: manilaQuotaData.QuotaSet.Gigabytes,
			Usage: uint64(totalShareUsage),
		},
		"snapshot_capacity": {
			Quota: manilaQuotaData.QuotaSet.SnapshotGigabytes,
			Usage: uint64(totalSnapshotUsage),
		},
	}, err
}

//SetQuota implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) SetQuota(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts, clusterID, domainUUID, projectUUID string, quotas map[string]uint64) error {
	client, err := openstack.NewSharedFileSystemV2(provider, eo)
	if err != nil {
		return err
	}

	requestData := map[string]map[string]uint64{
		"quota_set": {
			"gigabytes":          quotas["share_capacity"],
			"snapshots":          quotas["share_snapshots"],
			"snapshot_gigabytes": quotas["snapshot_capacity"],
			"shares":             quotas["shares"],
			"share_networks":     quotas["share_networks"],
		},
	}

	url := client.ServiceURL("os-quota-sets", projectUUID)
	_, err = client.Put(url, requestData, nil, &gophercloud.RequestOpts{OkCodes: []int{200}})

	return err
}
<commit_msg>add API pagination to Manila quota plugin<commit_after>/*******************************************************************************
*
* Copyright 2017 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, 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 plugins

import (
	"fmt"

	"github.com/gophercloud/gophercloud"
	"github.com/gophercloud/gophercloud/openstack"
	"github.com/gophercloud/gophercloud/openstack/sharedfilesystems/v2/sharenetworks"
	"github.com/gophercloud/gophercloud/pagination"
	"github.com/sapcc/go-bits/logg"
	"github.com/sapcc/limes/pkg/limes"
)

type manilaPlugin struct {
	cfg limes.ServiceConfiguration
}

var manilaResources = []limes.ResourceInfo{
	{
		Name: "share_networks",
		Unit: limes.UnitNone,
	},
	{
		Name: "share_capacity",
		Unit: limes.UnitGibibytes,
	},
	{
		Name: "shares",
		Unit: limes.UnitNone,
	},
	{
		Name: "snapshot_capacity",
		Unit: limes.UnitGibibytes,
	},
	{
		Name: "share_snapshots",
		Unit: limes.UnitNone,
	},
}

func init() {
	limes.RegisterQuotaPlugin(func(c limes.ServiceConfiguration, scrapeSubresources map[string]bool) limes.QuotaPlugin {
		return &manilaPlugin{c}
	})
}

//Init implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) Init(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts) error {
	return nil
}

//ServiceInfo implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) ServiceInfo() limes.ServiceInfo {
	return limes.ServiceInfo{
		Type:        "sharev2",
		ProductName: "manila",
		Area:        "storage",
	}
}

//Resources implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) Resources() []limes.ResourceInfo {
	return manilaResources
}

//Scrape implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) Scrape(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts, clusterID, domainUUID, projectUUID string) (map[string]limes.ResourceData, error) {
	client, err := openstack.NewSharedFileSystemV2(provider, eo)
	if err != nil {
		return nil, err
	}

	var result gophercloud.Result
	var totalShareUsage, totalSnapshotUsage, totalShareNetworksUsage = uint64(0), uint64(0), uint64(0)

	//Get absolute quota limits per project
	url := client.ServiceURL("os-quota-sets", projectUUID)
	_, err = client.Get(url, &result.Body, nil)
	if err != nil {
		return nil, err
	}
	var manilaQuotaData struct {
		QuotaSet struct {
			Gigabytes         int64 `json:"gigabytes"`
			Shares            int64 `json:"shares"`
			SnapshotGigabytes int64 `json:"snapshot_gigabytes"`
			Snapshots         int64 `json:"snapshots"`
			ShareNetworks     int64 `json:"share_networks"`
		} `json:"quota_set"`
	}
	err = result.ExtractInto(&manilaQuotaData)
	if err != nil {
		return nil, err
	}

	shares, err := manilaGetShares(client, projectUUID)
	if err != nil {
		return nil, err
	}
	for _, share := range shares {
		totalShareUsage += share.Size
	}

	//Get usage of snapshots per project
	snapshots, err := manilaGetSnapshots(client, projectUUID)
	if err != nil {
		return nil, err
	}
	for _, snapshot := range snapshots {
		totalSnapshotUsage += snapshot.ShareSize
	}

	//Get usage of shared networks
	sharenetworks.ListDetail(client, sharenetworks.ListOpts{ProjectID: projectUUID}).EachPage(func(page pagination.Page) (bool, error) {
		sn, err := sharenetworks.ExtractShareNetworks(page)
		if err != nil {
			return false, err
		}
		totalShareNetworksUsage = uint64(len(sn))
		return true, nil
	})

	logg.Debug("Scraped quota and usage for service: sharev2.")

	return map[string]limes.ResourceData{
		"shares": {
			Quota: manilaQuotaData.QuotaSet.Shares,
			Usage: uint64(len(shares)),
		},
		"share_snapshots": {
			Quota: manilaQuotaData.QuotaSet.Snapshots,
			Usage: uint64(len(snapshots)),
		},
		"share_networks": {
			Quota: manilaQuotaData.QuotaSet.ShareNetworks,
			Usage: uint64(totalShareNetworksUsage),
		},
		"share_capacity": {
			Quota: manilaQuotaData.QuotaSet.Gigabytes,
			Usage: uint64(totalShareUsage),
		},
		"snapshot_capacity": {
			Quota: manilaQuotaData.QuotaSet.SnapshotGigabytes,
			Usage: uint64(totalSnapshotUsage),
		},
	}, err
}

//SetQuota implements the limes.QuotaPlugin interface.
func (p *manilaPlugin) SetQuota(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts, clusterID, domainUUID, projectUUID string, quotas map[string]uint64) error {
	client, err := openstack.NewSharedFileSystemV2(provider, eo)
	if err != nil {
		return err
	}

	requestData := map[string]map[string]uint64{
		"quota_set": {
			"gigabytes":          quotas["share_capacity"],
			"snapshots":          quotas["share_snapshots"],
			"snapshot_gigabytes": quotas["snapshot_capacity"],
			"shares":             quotas["shares"],
			"share_networks":     quotas["share_networks"],
		},
	}

	url := client.ServiceURL("os-quota-sets", projectUUID)
	_, err = client.Put(url, requestData, nil, &gophercloud.RequestOpts{OkCodes: []int{200}})

	return err
}

////////////////////////////////////////////////////////////////////////////////

type manilaShare struct {
	ID   string `json:"id"`
	Size uint64 `json:"size"`
}

func manilaGetShares(client *gophercloud.ServiceClient, projectUUID string) (result []manilaShare, err error) {
	page := 0
	pageSize := 250

	for {
		url := client.ServiceURL("shares", "detail") + fmt.Sprintf("?project_id=%s&all_tenants=1&limit=%d&offset=%d", projectUUID, pageSize, page*pageSize)
		var r gophercloud.Result
		_, err = client.Get(url, &r.Body, nil)
		if err != nil {
			return nil, err
		}

		var data struct {
			Shares []manilaShare `json:"shares"`
		}
		err = r.ExtractInto(&data)
		if err != nil {
			return nil, err
		}

		if len(data.Shares) > 0 {
			result = append(result, data.Shares...)
			page++
		} else {
			//last page reached
			return
		}
	}
}

type manilaSnapshot struct {
	ID        string `json:"id"`
	ShareSize uint64 `json:"share_size"`
}

func manilaGetSnapshots(client *gophercloud.ServiceClient, projectUUID string) (result []manilaSnapshot, err error) {
	page := 0
	pageSize := 250

	for {
		url := client.ServiceURL("snapshots", "detail") + fmt.Sprintf("?project_id=%s&all_tenants=1&limit=%d&offset=%d", projectUUID, pageSize, page*pageSize)
		var r gophercloud.Result
		_, err = client.Get(url, &r.Body, nil)
		if err != nil {
			return nil, err
		}

		var data struct {
			Snapshots []manilaSnapshot `json:"snapshots"`
		}
		err = r.ExtractInto(&data)
		if err != nil {
			return nil, err
		}

		if len(data.Snapshots) > 0 {
			result = append(result, data.Snapshots...)
			page++
		} else {
			//last page reached
			return
		}
	}
}
<|endoftext|>


<commit_before>package tenant

import (
	"github.com/google/uuid"
	"github.com/maurofran/kit/assert"
	"github.com/maurofran/kit/domain"
)

// Tenant is the aggregate root object used to manage a tenant.
type Tenant struct {
	*domain.AggregateRoot `bson:"-"`
	ID                    ID            `bson:"_id"`
	Name                  string        `bson:"name"`
	Description           string        `bson:"description"`
	Active                bool          `bson:"active"`
	Invitations           []*Invitation `bson:"invitations,omitempty"`
}

// New is used to create a new tenant instance.
// It returns an erro if any of supplied parameter does not comply with restrictions.
func New(id ID, name, description string, active bool) (*Tenant, error) {
	if err := assert.NotZero(id, "id"); err != nil {
		return nil, err
	}
	if err := assert.NotEmpty(name, "name"); err != nil {
		return nil, err
	}
	if err := assert.MaxLength(name, 70, "name"); err != nil {
		return nil, err
	}
	if err := assert.NotEmpty(description, "description"); err != nil {
		return nil, err
	}
	tenant := &Tenant{
		AggregateRoot: &domain.AggregateRoot{},
		ID:            id,
		Name:          name,
		Description:   description,
		Active:        active,
	}
	tenant.RegisterEvent(provisioned(id, name))
	return tenant, nil
}

// Activate this tenant if not active.
func (t *Tenant) Activate() {
	if !t.Active {
		t.Active = true
		t.RegisterEvent(activated(t.ID))
	}
}

// Deactivate this tenant if active.
func (t *Tenant) Deactivate() {
	if t.Active {
		t.Active = false
		t.RegisterEvent(deactivated(t.ID))
	}
}

// AllAvailableInvitations will retrieve a slice of available invitation descriptors.
func (t *Tenant) AllAvailableInvitations() ([]InvitationDescriptor, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return nil, err
	}
	var result []InvitationDescriptor
	for _, invitation := range t.Invitations {
		if invitation.IsAvailable() {
			result = append(result, invitation.toDescriptor(t.ID))
		}
	}
	return result, nil
}

// AllUnavailableInvitations will retrieve a slice of unavailable invitation descriptors.
func (t *Tenant) AllUnavailableInvitations() ([]InvitationDescriptor, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return nil, err
	}
	var result []InvitationDescriptor
	for _, invitation := range t.Invitations {
		if !invitation.IsAvailable() {
			result = append(result, invitation.toDescriptor(t.ID))
		}
	}
	return result, nil
}

// IsInvitationAvailableThrough will check if an invitation is available.
func (t *Tenant) IsInvitationAvailableThrough(invitationIdentifier string) (bool, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return false, err
	}
	idx, _ := t.invitationFor(invitationIdentifier)
	return idx >= 0, nil
}

// OfferInvitation will offer a new invitation and return to the caller.
func (t *Tenant) OfferInvitation(description string) (*Invitation, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return nil, err
	}
	exists, err := t.IsInvitationAvailableThrough(description)
	if err != nil {
		return nil, err
	}
	if err := assert.StateNot(exists, "description exists"); err != nil {
		return nil, err
	}
	invitation, err := newInvitation(uuid.New().String(), description)
	if err != nil {
		return nil, err
	}
	invitations := append(t.Invitations, invitation)
	t.Invitations = invitations
	return invitation, nil
}

// WithdrawInvitation will remove the invitation identified by supplied identifier
func (t *Tenant) WithdrawInvitation(identifier string) error {
	if err := assert.State(t.Active, "active"); err != nil {
		return err
	}
	if idx, _ := t.invitationFor(identifier); idx >= 0 {
		var inv []*Invitation
		copy(inv, t.Invitations)
		inv[idx] = inv[len(inv)-1]
		inv[len(inv)-1] = nil
		t.Invitations = inv[:len(inv)-1]
	}
	return nil
}

// GetInvitation will find the invitation for supplied identifier.
func (t *Tenant) GetInvitation(identifier string) (*Invitation, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return nil, err
	}
	_, invitation := t.invitationFor(identifier)
	return invitation, nil
}

func (t *Tenant) invitationFor(identifier string) (int, *Invitation) {
	for idx, invitation := range t.Invitations {
		if invitation.IsIdentifiedBy(identifier) {
			return idx, invitation
		}
	}
	return -1, nil
}
<commit_msg>Renamed ID bson to tenantId.<commit_after>package tenant

import (
	"github.com/google/uuid"
	"github.com/maurofran/kit/assert"
	"github.com/maurofran/kit/domain"
)

// Tenant is the aggregate root object used to manage a tenant.
type Tenant struct {
	*domain.AggregateRoot `bson:"-"`
	ID                    ID            `bson:"tenantId"`
	Name                  string        `bson:"name"`
	Description           string        `bson:"description"`
	Active                bool          `bson:"active"`
	Invitations           []*Invitation `bson:"invitations,omitempty"`
}

// New is used to create a new tenant instance.
// It returns an erro if any of supplied parameter does not comply with restrictions.
func New(id ID, name, description string, active bool) (*Tenant, error) {
	if err := assert.NotZero(id, "id"); err != nil {
		return nil, err
	}
	if err := assert.NotEmpty(name, "name"); err != nil {
		return nil, err
	}
	if err := assert.MaxLength(name, 70, "name"); err != nil {
		return nil, err
	}
	if err := assert.NotEmpty(description, "description"); err != nil {
		return nil, err
	}
	tenant := &Tenant{
		AggregateRoot: &domain.AggregateRoot{},
		ID:            id,
		Name:          name,
		Description:   description,
		Active:        active,
	}
	tenant.RegisterEvent(provisioned(id, name))
	return tenant, nil
}

// Activate this tenant if not active.
func (t *Tenant) Activate() {
	if !t.Active {
		t.Active = true
		t.RegisterEvent(activated(t.ID))
	}
}

// Deactivate this tenant if active.
func (t *Tenant) Deactivate() {
	if t.Active {
		t.Active = false
		t.RegisterEvent(deactivated(t.ID))
	}
}

// AllAvailableInvitations will retrieve a slice of available invitation descriptors.
func (t *Tenant) AllAvailableInvitations() ([]InvitationDescriptor, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return nil, err
	}
	var result []InvitationDescriptor
	for _, invitation := range t.Invitations {
		if invitation.IsAvailable() {
			result = append(result, invitation.toDescriptor(t.ID))
		}
	}
	return result, nil
}

// AllUnavailableInvitations will retrieve a slice of unavailable invitation descriptors.
func (t *Tenant) AllUnavailableInvitations() ([]InvitationDescriptor, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return nil, err
	}
	var result []InvitationDescriptor
	for _, invitation := range t.Invitations {
		if !invitation.IsAvailable() {
			result = append(result, invitation.toDescriptor(t.ID))
		}
	}
	return result, nil
}

// IsInvitationAvailableThrough will check if an invitation is available.
func (t *Tenant) IsInvitationAvailableThrough(invitationIdentifier string) (bool, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return false, err
	}
	idx, _ := t.invitationFor(invitationIdentifier)
	return idx >= 0, nil
}

// OfferInvitation will offer a new invitation and return to the caller.
func (t *Tenant) OfferInvitation(description string) (*Invitation, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return nil, err
	}
	exists, err := t.IsInvitationAvailableThrough(description)
	if err != nil {
		return nil, err
	}
	if err := assert.StateNot(exists, "description exists"); err != nil {
		return nil, err
	}
	invitation, err := newInvitation(uuid.New().String(), description)
	if err != nil {
		return nil, err
	}
	invitations := append(t.Invitations, invitation)
	t.Invitations = invitations
	return invitation, nil
}

// WithdrawInvitation will remove the invitation identified by supplied identifier
func (t *Tenant) WithdrawInvitation(identifier string) error {
	if err := assert.State(t.Active, "active"); err != nil {
		return err
	}
	if idx, _ := t.invitationFor(identifier); idx >= 0 {
		var inv []*Invitation
		copy(inv, t.Invitations)
		inv[idx] = inv[len(inv)-1]
		inv[len(inv)-1] = nil
		t.Invitations = inv[:len(inv)-1]
	}
	return nil
}

// GetInvitation will find the invitation for supplied identifier.
func (t *Tenant) GetInvitation(identifier string) (*Invitation, error) {
	if err := assert.State(t.Active, "active"); err != nil {
		return nil, err
	}
	_, invitation := t.invitationFor(identifier)
	return invitation, nil
}

func (t *Tenant) invitationFor(identifier string) (int, *Invitation) {
	for idx, invitation := range t.Invitations {
		if invitation.IsIdentifiedBy(identifier) {
			return idx, invitation
		}
	}
	return -1, nil
}
<|endoftext|>


<commit_before>package workspace

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/GetStream/vg/internal/utils"
	"github.com/pkg/errors"
)

func (ws *Workspace) InstallLocalPackagePersistently(pkg string, localPath string) error {
	err := ws.InstallLocalPackage(pkg, localPath)
	if err != nil {
		return err
	}

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Persistent = true

	fmt.Fprintf(os.Stderr, "Persisting the local install for %q\n", pkg)
	return ws.SaveSettings(settings)
}

func (ws *Workspace) InstallLocalPackage(pkg string, localPath string) error {
	pkgDir := filepath.Join(path.Split(pkg))
	target := filepath.Join(ws.Src(), pkgDir)

	err := ws.Uninstall(pkg, os.Stderr)
	if err != nil {
		return err
	}

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg] = &localInstall{
		Path: localPath,
	}

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	hasBindfs := true

	_, err = exec.LookPath("bindfs")
	if err != nil {
		execErr, ok := err.(*exec.Error)
		if !ok || execErr.Err != exec.ErrNotFound {
			return err
		}

		hasBindfs = false
	}

	if hasBindfs {
		err = ws.installLocalPackageWithBindfs(pkg, localPath, target)
	} else {
		err = ws.installLocalPackageWithSymlink(pkg, localPath, target)
	}

	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Successful = true
	return ws.SaveSettings(settings)

}

func (ws *Workspace) installLocalPackageWithBindfs(pkg, src, target string) error {
	_, _ = fmt.Fprintf(os.Stderr, "Installing local sources at %q in workspace as %q using bindfs\n", src, pkg)

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Bindfs = true

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	err = os.MkdirAll(target, 0755)
	if err != nil {
		return errors.WithStack(err)
	}

	cmd := exec.Command("bindfs")
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stderr
	err = cmd.Run()
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (ws *Workspace) installLocalPackageWithSymlink(pkg, src, target string) error {
	_, _ = fmt.Fprintf(os.Stderr, "Installing local sources at %q in workspace as %q using symbolic links\n", src, pkg)

	err := os.MkdirAll(filepath.Dir(target), 0755)
	if err != nil {
		return errors.WithStack(err)
	}

	err = os.RemoveAll(target)
	if err != nil {
		return errors.WithStack(err)
	}

	return errors.WithStack(os.Symlink(src, target))
}

func (ws *Workspace) InstallPersistentLocalPackages() error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	for pkg, install := range settings.LocalInstalls {
		if !install.Persistent {
			continue
		}
		err := ws.InstallLocalPackagePersistently(pkg, install.Path)
		if err != nil {
			return err
		}
	}
	return nil
}

func (ws *Workspace) Uninstall(pkg string, logWriter io.Writer) error {
	pkgDir := utils.PkgToDir(pkg)
	pkgSrc := filepath.Join(ws.Src(), pkgDir)
	_, err := os.Stat(pkgSrc)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return errors.WithStack(err)
	}

	fmt.Fprintf(logWriter, "Uninstalling %q from workspace\n", pkg)
	fmt.Fprintf(logWriter, "  Removing sources at %q\n", pkgSrc)
	err = os.RemoveAll(pkgSrc)
	if err != nil {
		return errors.Wrapf(err, "Couldn't remove package src %q", ws.Name())
	}

	// Remove possible LocalInstall entry
	settings, err := ws.Settings()
	if err != nil {
		return err
	}
	install, ok := settings.LocalInstalls[pkg]
	if ok && !install.Persistent {
		fmt.Fprintf(logWriter, "  Removing %q from locally installed packages\n", pkg)
		delete(settings.LocalInstalls, pkg)
	}

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}
	pkgInstalledDirs, err := filepath.Glob(filepath.Join(ws.Pkg(), "*", pkgDir))
	if err != nil {
		return errors.Wrapf(err, "Something went wrong when globbing files for %q", pkg)
	}

	for _, path := range pkgInstalledDirs {
		fmt.Fprintf(logWriter, "  Removing %q\n", path)

		err = os.RemoveAll(path)
		if err != nil {
			return errors.Wrapf(err, "Couldn't remove installed package files for %q", pkg)
		}
	}

	pkgInstalledFiles, err := filepath.Glob(filepath.Join(ws.Pkg(), "*", pkgDir+".a"))
	if err != nil {
		return errors.Wrapf(err, "Something went wrong when globbing files for %q", pkg)
	}

	for _, path := range pkgInstalledFiles {
		fmt.Fprintf(logWriter, "  Removing %q\n", path)

		err = os.RemoveAll(path)
		if err != nil {
			return errors.Wrapf(err, "Couldn't remove installed package files for %q", pkg)
		}
	}
	return nil
}

func (ws *Workspace) ClearSrc() error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	err = os.RemoveAll(ws.ensureMarker())
	if err != nil {
		return errors.WithStack(err)
	}

	for pkg := range settings.LocalInstalls {
		err := ws.Uninstall(pkg, os.Stdout)
		if err != nil {
			return err
		}
	}

	return errors.WithStack(os.RemoveAll(ws.Src()))

}

func (ws *Workspace) UnpersistLocalInstall(pkg string) error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	if install, ok := settings.LocalInstalls[pkg]; ok {
		fmt.Printf("Removing %q from persistent local installs\n", pkg)
		install.Persistent = false

		err = ws.SaveSettings(settings)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Actually use bindfs to mount local installs and unmount them when uninstalling<commit_after>package workspace

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/GetStream/vg/internal/utils"
	"github.com/pkg/errors"
)

func (ws *Workspace) InstallLocalPackagePersistently(pkg string, localPath string) error {
	err := ws.InstallLocalPackage(pkg, localPath)
	if err != nil {
		return err
	}

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Persistent = true

	fmt.Fprintf(os.Stderr, "Persisting the local install for %q\n", pkg)
	return ws.SaveSettings(settings)
}

func (ws *Workspace) InstallLocalPackage(pkg string, localPath string) error {
	pkgDir := filepath.Join(path.Split(pkg))
	target := filepath.Join(ws.Src(), pkgDir)

	err := ws.Uninstall(pkg, os.Stderr)
	if err != nil {
		return err
	}

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg] = &localInstall{
		Path: localPath,
	}

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	hasBindfs := true

	_, err = exec.LookPath("bindfs")
	if err != nil {
		execErr, ok := err.(*exec.Error)
		if !ok || execErr.Err != exec.ErrNotFound {
			return err
		}

		hasBindfs = false
	}

	if hasBindfs {
		err = ws.installLocalPackageWithBindfs(pkg, localPath, target)
	} else {
		err = ws.installLocalPackageWithSymlink(pkg, localPath, target)
	}

	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Successful = true
	return ws.SaveSettings(settings)

}

func (ws *Workspace) installLocalPackageWithBindfs(pkg, src, target string) error {
	_, _ = fmt.Fprintf(os.Stderr, "Installing local sources at %q in workspace as %q using bindfs\n", src, pkg)

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Bindfs = true

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	err = os.MkdirAll(target, 0755)
	if err != nil {
		return errors.WithStack(err)
	}

	cmd := exec.Command("bindfs", "-r", "--no-allow-other", src, target)
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stderr
	return errors.WithStack(cmd.Run())
}

func (ws *Workspace) installLocalPackageWithSymlink(pkg, src, target string) error {
	_, _ = fmt.Fprintf(os.Stderr, "Installing local sources at %q in workspace as %q using symbolic links\n", src, pkg)

	err := os.MkdirAll(filepath.Dir(target), 0755)
	if err != nil {
		return errors.WithStack(err)
	}

	err = os.RemoveAll(target)
	if err != nil {
		return errors.WithStack(err)
	}

	return errors.WithStack(os.Symlink(src, target))
}

func (ws *Workspace) InstallPersistentLocalPackages() error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	for pkg, install := range settings.LocalInstalls {
		if !install.Persistent {
			continue
		}
		err := ws.InstallLocalPackagePersistently(pkg, install.Path)
		if err != nil {
			return err
		}
	}
	return nil
}

func (ws *Workspace) Uninstall(pkg string, logWriter io.Writer) error {
	pkgDir := utils.PkgToDir(pkg)
	pkgSrc := filepath.Join(ws.Src(), pkgDir)
	_, err := os.Stat(pkgSrc)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return errors.WithStack(err)
	}

	fmt.Fprintf(logWriter, "Uninstalling %q from workspace\n", pkg)
	// Check if locally installed
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	install, localInstalled := settings.LocalInstalls[pkg]

	if localInstalled && install.Bindfs {
		fmt.Fprintf(logWriter, "  Unmounting bindfs mount at %q\n", pkgSrc)
		stderrBuff := &bytes.Buffer{}
		outputBuff := &bytes.Buffer{}

		cmd := exec.Command("fusermount", "-u", pkgSrc)
		cmd.Stderr = io.MultiWriter(stderrBuff, outputBuff)
		cmd.Stdout = outputBuff

		err := cmd.Run()
		if err != nil {
			notMountedOutput := fmt.Sprintf("fusermount: entry for %s not found", pkgSrc)
			if !strings.HasPrefix(stderrBuff.String(), notMountedOutput) {
				io.Copy(outputBuff, os.Stderr)
				return errors.WithStack(err)
			}
		}

	}

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	fmt.Fprintf(logWriter, "  Removing sources at %q\n", pkgSrc)
	err = os.RemoveAll(pkgSrc)
	if err != nil {
		return errors.Wrapf(err, "Couldn't remove package src %q", ws.Name())
	}

	if localInstalled && !install.Persistent {
		fmt.Fprintf(logWriter, "  Removing %q from locally installed packages\n", pkg)
		delete(settings.LocalInstalls, pkg)
		err = ws.SaveSettings(settings)
		if err != nil {
			return err
		}
	}

	pkgInstalledDirs, err := filepath.Glob(filepath.Join(ws.Pkg(), "*", pkgDir))
	if err != nil {
		return errors.Wrapf(err, "Something went wrong when globbing files for %q", pkg)
	}

	for _, path := range pkgInstalledDirs {
		fmt.Fprintf(logWriter, "  Removing %q\n", path)

		err = os.RemoveAll(path)
		if err != nil {
			return errors.Wrapf(err, "Couldn't remove installed package files for %q", pkg)
		}
	}

	pkgInstalledFiles, err := filepath.Glob(filepath.Join(ws.Pkg(), "*", pkgDir+".a"))
	if err != nil {
		return errors.Wrapf(err, "Something went wrong when globbing files for %q", pkg)
	}

	for _, path := range pkgInstalledFiles {
		fmt.Fprintf(logWriter, "  Removing %q\n", path)

		err = os.RemoveAll(path)
		if err != nil {
			return errors.Wrapf(err, "Couldn't remove installed package files for %q", pkg)
		}
	}
	return nil
}

func (ws *Workspace) ClearSrc() error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	err = os.RemoveAll(ws.ensureMarker())
	if err != nil {
		return errors.WithStack(err)
	}

	for pkg := range settings.LocalInstalls {
		err := ws.Uninstall(pkg, os.Stdout)
		if err != nil {
			return err
		}
	}

	return errors.WithStack(os.RemoveAll(ws.Src()))

}

func (ws *Workspace) UnpersistLocalInstall(pkg string) error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	if install, ok := settings.LocalInstalls[pkg]; ok {
		fmt.Printf("Removing %q from persistent local installs\n", pkg)
		install.Persistent = false

		err = ws.SaveSettings(settings)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package bot

import (
	"errors"

	"github.com/ChimeraCoder/anaconda"
	"github.com/erbridge/gotwit/callback"
	"github.com/erbridge/gotwit/twitter"
)

type (
	Bot struct {
		screenName     string
		consumerConfig twitter.ConsumerConfig
		client         twitter.Client
		callbacks      map[callback.Type]map[int]callback.Callback
		nextCallbackId int
	}
)

var (
	activeBot *Bot
)

func New(screenName string, consumerConfig twitter.ConsumerConfig, accessConfig twitter.AccessConfig) Bot {
	b := Bot{
		screenName:     screenName,
		consumerConfig: consumerConfig,
		callbacks:      make(map[callback.Type]map[int]callback.Callback),
	}
	b.client = twitter.NewClient(accessConfig, b.handle)
	return b
}

func (b *Bot) Start() error {
	if activeBot == nil {
		activeBot = b
	} else if activeBot == b {
		return errors.New("This bot is already started")
	} else {
		return errors.New("Only one bot can be active at once")
	}

	twitter.BindConsumer(b.consumerConfig)
	b.client.Start()

	return nil
}

func (b *Bot) Stop() error {
	if activeBot == b {
		activeBot = nil
	} else {
		return errors.New("This bot hasn't been started")
	}

	twitter.UnbindConsumer()
	b.client.Stop()

	return nil
}

func (b *Bot) PostTweet(message string) {
	b.client.PostTweet(message)
}

func (b *Bot) RegisterCallback(t callback.Type, cb callback.Callback) (id int) {
	if _, ok := b.callbacks[t]; !ok {
		b.callbacks[t] = make(map[int]callback.Callback)
	}

	id = b.nextCallbackId
	b.callbacks[t][id] = cb
	b.nextCallbackId++
	return
}

func (b *Bot) UnregisterCallback(t callback.Type, id int) {
	delete(b.callbacks[t], id)
}

func (b *Bot) triggerCallback(t callback.Type, tweet anaconda.Tweet) {
	for _, cb := range b.callbacks[t] {
		cb(tweet)
	}
}

func (b *Bot) handle(tweet anaconda.Tweet) {
	if tweet.RetweetedStatus != nil {
		b.triggerCallback(callback.Retweet, tweet)
		return
	}

	if tweet.InReplyToScreenName == b.screenName {
		b.triggerCallback(callback.Reply, tweet)
		return
	}

	for _, entity := range tweet.Entities.User_mentions {
		if entity.Screen_name == b.screenName {
			b.triggerCallback(callback.Mention, tweet)
			return
		}
	}
}
<commit_msg>Make Bot.PostTweet return errors.<commit_after>package bot

import (
	"errors"

	"github.com/ChimeraCoder/anaconda"
	"github.com/erbridge/gotwit/callback"
	"github.com/erbridge/gotwit/twitter"
)

type (
	Bot struct {
		screenName     string
		consumerConfig twitter.ConsumerConfig
		client         twitter.Client
		callbacks      map[callback.Type]map[int]callback.Callback
		nextCallbackId int
	}
)

var (
	activeBot *Bot
)

func New(screenName string, consumerConfig twitter.ConsumerConfig, accessConfig twitter.AccessConfig) Bot {
	b := Bot{
		screenName:     screenName,
		consumerConfig: consumerConfig,
		callbacks:      make(map[callback.Type]map[int]callback.Callback),
	}
	b.client = twitter.NewClient(accessConfig, b.handle)
	return b
}

func (b *Bot) Start() error {
	if activeBot == nil {
		activeBot = b
	} else if activeBot == b {
		return errors.New("This bot is already started")
	} else {
		return errors.New("Only one bot can be active at once")
	}

	twitter.BindConsumer(b.consumerConfig)
	b.client.Start()

	return nil
}

func (b *Bot) Stop() error {
	if activeBot == b {
		activeBot = nil
	} else {
		return errors.New("This bot hasn't been started")
	}

	twitter.UnbindConsumer()
	b.client.Stop()

	return nil
}

func (b *Bot) PostTweet(message string) error {
	return b.client.PostTweet(message)
}

func (b *Bot) RegisterCallback(t callback.Type, cb callback.Callback) (id int) {
	if _, ok := b.callbacks[t]; !ok {
		b.callbacks[t] = make(map[int]callback.Callback)
	}

	id = b.nextCallbackId
	b.callbacks[t][id] = cb
	b.nextCallbackId++
	return
}

func (b *Bot) UnregisterCallback(t callback.Type, id int) {
	delete(b.callbacks[t], id)
}

func (b *Bot) triggerCallback(t callback.Type, tweet anaconda.Tweet) {
	for _, cb := range b.callbacks[t] {
		cb(tweet)
	}
}

func (b *Bot) handle(tweet anaconda.Tweet) {
	if tweet.RetweetedStatus != nil {
		b.triggerCallback(callback.Retweet, tweet)
		return
	}

	if tweet.InReplyToScreenName == b.screenName {
		b.triggerCallback(callback.Reply, tweet)
		return
	}

	for _, entity := range tweet.Entities.User_mentions {
		if entity.Screen_name == b.screenName {
			b.triggerCallback(callback.Mention, tweet)
			return
		}
	}
}
<|endoftext|>


<commit_before>package twitch

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net/textproto"
	"strings"
	"sync/atomic"
	"time"
)

const (
	// ircTwitch constant for twitch irc chat address
	ircTwitch = "irc.chat.twitch.tv:443"
)

// User data you receive from tmi
type User struct {
	Username    string
	DisplayName string
	UserType    string
	Color       string
	Badges      map[string]int
}

// Message data you receive from tmi
type Message struct {
	Type   msgType
	Time   time.Time
	Action bool
	Emotes []*Emote
	Tags   map[string]string
	Text   string
}

// Client client to control your connection and attach callbacks
type Client struct {
	IrcAddress            string
	ircUser               string
	ircToken              string
	connection            *tls.Conn
	connActive            tAtomBool
	onNewMessage          func(channel string, user User, message Message)
	onNewRoomstateMessage func(channel string, user User, message Message)
	onNewClearchatMessage func(channel string, user User, message Message)
}

// NewClient to create a new client
func NewClient(username, oauth string) *Client {
	return &Client{
		ircUser:    username,
		ircToken:   oauth,
		IrcAddress: ircTwitch,
	}
}

// OnNewMessage attach callback to new standard chat messages
func (c *Client) OnNewMessage(callback func(channel string, user User, message Message)) {
	c.onNewMessage = callback
}

// OnNewRoomstateMessage attach callback to new messages such as submode enabled
func (c *Client) OnNewRoomstateMessage(callback func(channel string, user User, message Message)) {
	c.onNewRoomstateMessage = callback
}

// OnNewClearchatMessage attach callback to new messages such as timeouts
func (c *Client) OnNewClearchatMessage(callback func(channel string, user User, message Message)) {
	c.onNewClearchatMessage = callback
}

// Say write something in a chat
func (c *Client) Say(channel, text string) {
	c.send(fmt.Sprintf("PRIVMSG #%s :%s", channel, text))
}

// Join enter a twitch channel to read more messages
func (c *Client) Join(channel string) {
	go c.send(fmt.Sprintf("JOIN #%s", channel))
}

// Connect connect the client to the irc server
func (c *Client) Connect() error {
	var conf *tls.Config
	// This means we are connecting to "localhost". Disable certificate chain check
	if strings.HasPrefix(c.IrcAddress, ":") {
		conf = &tls.Config{
			InsecureSkipVerify: true,
		}
	} else {
		conf = &tls.Config{}
	}
	for {
		conn, err := tls.Dial("tcp", c.IrcAddress, conf)
		c.connection = conn
		if err != nil {
			return err
		}

		go c.setupConnection()

		err = c.readConnection(conn)
		if err != nil {
			time.Sleep(time.Millisecond * 200)
			continue
		}
	}
}

func (c *Client) readConnection(conn *tls.Conn) error {
	reader := bufio.NewReader(conn)
	tp := textproto.NewReader(reader)
	for {
		line, err := tp.ReadLine()
		if err != nil {
			return err
		}
		messages := strings.Split(line, "\r\n")
		for _, msg := range messages {
			if !c.connActive.get() && strings.Contains(msg, ":tmi.twitch.tv 001") {
				c.connActive.set(true)
			}
			c.handleLine(msg)
		}
	}
}

func (c *Client) setupConnection() {
	c.connection.Write([]byte("PASS " + c.ircToken + "\r\n"))
	c.connection.Write([]byte("NICK " + c.ircUser + "\r\n"))
	c.connection.Write([]byte("CAP REQ :twitch.tv/tags\r\n"))
	c.connection.Write([]byte("CAP REQ :twitch.tv/commands\r\n"))
}

func (c *Client) send(line string) {
	for i := 0; i < 1000; i++ {
		if !c.connActive.get() {
			time.Sleep(time.Millisecond * 2)
			continue
		}
		c.connection.Write([]byte(line + "\r\n"))
		return
	}
}

func (c *Client) handleLine(line string) {
	if strings.HasPrefix(line, "PING") {
		c.send(strings.Replace(line, "PING", "PONG", 1))
	}
	if strings.HasPrefix(line, "@") {
		message := parseMessage(line)

		Channel := message.Channel

		User := &User{
			Username:    message.Username,
			DisplayName: message.DisplayName,
			UserType:    message.UserType,
			Color:       message.Color,
			Badges:      message.Badges,
		}

		clientMessage := &Message{
			Type:   message.Type,
			Time:   message.Time,
			Action: message.Action,
			Emotes: message.Emotes,
			Tags:   message.Tags,
			Text:   message.Text,
		}

		switch message.Type {
		case PRIVMSG:
			if c.onNewMessage != nil {
				c.onNewMessage(Channel, *User, *clientMessage)
			}
		case ROOMSTATE:
			if c.onNewRoomstateMessage != nil {
				c.onNewRoomstateMessage(Channel, *User, *clientMessage)
			}
		case CLEARCHAT:
			if c.onNewClearchatMessage != nil {
				c.onNewClearchatMessage(Channel, *User, *clientMessage)
			}
		}
	}
}

// tAtomBool atomic bool for writing/reading across threads
type tAtomBool struct{ flag int32 }

func (b *tAtomBool) set(value bool) {
	var i int32
	if value {
		i = 1
	}
	atomic.StoreInt32(&(b.flag), int32(i))
}

func (b *tAtomBool) get() bool {
	if atomic.LoadInt32(&(b.flag)) != 0 {
		return true
	}
	return false
}
<commit_msg>adding whisper functionality<commit_after>package twitch

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net/textproto"
	"strings"
	"sync/atomic"
	"time"
)

const (
	// ircTwitch constant for twitch irc chat address
	ircTwitch = "irc.chat.twitch.tv:443"
)

// User data you receive from tmi
type User struct {
	Username    string
	DisplayName string
	UserType    string
	Color       string
	Badges      map[string]int
}

// Message data you receive from tmi
type Message struct {
	Type   msgType
	Time   time.Time
	Action bool
	Emotes []*Emote
	Tags   map[string]string
	Text   string
}

// Client client to control your connection and attach callbacks
type Client struct {
	IrcAddress            string
	ircUser               string
	ircToken              string
	connection            *tls.Conn
	connActive            tAtomBool
	onNewMessage          func(channel string, user User, message Message)
	onNewRoomstateMessage func(channel string, user User, message Message)
	onNewClearchatMessage func(channel string, user User, message Message)
}

// NewClient to create a new client
func NewClient(username, oauth string) *Client {
	return &Client{
		ircUser:    username,
		ircToken:   oauth,
		IrcAddress: ircTwitch,
	}
}

// OnNewMessage attach callback to new standard chat messages
func (c *Client) OnNewMessage(callback func(channel string, user User, message Message)) {
	c.onNewMessage = callback
}

// OnNewRoomstateMessage attach callback to new messages such as submode enabled
func (c *Client) OnNewRoomstateMessage(callback func(channel string, user User, message Message)) {
	c.onNewRoomstateMessage = callback
}

// OnNewClearchatMessage attach callback to new messages such as timeouts
func (c *Client) OnNewClearchatMessage(callback func(channel string, user User, message Message)) {
	c.onNewClearchatMessage = callback
}

// Say write something in a chat
func (c *Client) Say(channel, text string) {
	c.send(fmt.Sprintf("PRIVMSG #%s :%s", channel, text))
}

// Join enter a twitch channel to read more messages
func (c *Client) Join(channel string) {
	go c.send(fmt.Sprintf("JOIN #%s", channel))
}

// Connect connect the client to the irc server
func (c *Client) Connect() error {
	var conf *tls.Config
	// This means we are connecting to "localhost". Disable certificate chain check
	if strings.HasPrefix(c.IrcAddress, ":") {
		conf = &tls.Config{
			InsecureSkipVerify: true,
		}
	} else {
		conf = &tls.Config{}
	}
	for {
		conn, err := tls.Dial("tcp", c.IrcAddress, conf)
		c.connection = conn
		if err != nil {
			return err
		}

		go c.setupConnection()

		err = c.readConnection(conn)
		if err != nil {
			time.Sleep(time.Millisecond * 200)
			continue
		}
	}
}

func (c *Client) readConnection(conn *tls.Conn) error {
	reader := bufio.NewReader(conn)
	tp := textproto.NewReader(reader)
	for {
		line, err := tp.ReadLine()
		if err != nil {
			return err
		}
		messages := strings.Split(line, "\r\n")
		for _, msg := range messages {
			if !c.connActive.get() && strings.Contains(msg, ":tmi.twitch.tv 001") {
				c.connActive.set(true)
			}
			c.handleLine(msg)
		}
	}
}

func (c *Client) setupConnection() {
	c.connection.Write([]byte("PASS " + c.ircToken + "\r\n"))
	c.connection.Write([]byte("NICK " + c.ircUser + "\r\n"))
	c.connection.Write([]byte("CAP REQ :twitch.tv/tags\r\n"))
	c.connection.Write([]byte("CAP REQ :twitch.tv/commands\r\n"))
}

func (c *Client) send(line string) {
	for i := 0; i < 1000; i++ {
		if !c.connActive.get() {
			time.Sleep(time.Millisecond * 2)
			continue
		}
		c.connection.Write([]byte(line + "\r\n"))
		return
	}
}

func (c *Client) handleLine(line string) {
	if strings.HasPrefix(line, "PING") {
		c.send(strings.Replace(line, "PING", "PONG", 1))
	}
	if strings.HasPrefix(line, "@") {
		message := parseMessage(line)

		Channel := message.Channel

		User := &User{
			Username:    message.Username,
			DisplayName: message.DisplayName,
			UserType:    message.UserType,
			Color:       message.Color,
			Badges:      message.Badges,
		}

		clientMessage := &Message{
			Type:   message.Type,
			Time:   message.Time,
			Action: message.Action,
			Emotes: message.Emotes,
			Tags:   message.Tags,
			Text:   message.Text,
		}

		switch message.Type {
		case WHISPER:
			if c.onNewMessage != nil {
                                c.onNewMessage(Channel, *User, *clientMessage)
                        }
		case PRIVMSG:
			if c.onNewMessage != nil {
				c.onNewMessage(Channel, *User, *clientMessage)
			}
		case ROOMSTATE:
			if c.onNewRoomstateMessage != nil {
				c.onNewRoomstateMessage(Channel, *User, *clientMessage)
			}
		case CLEARCHAT:
			if c.onNewClearchatMessage != nil {
				c.onNewClearchatMessage(Channel, *User, *clientMessage)
			}
		}
	}
}

// tAtomBool atomic bool for writing/reading across threads
type tAtomBool struct{ flag int32 }

func (b *tAtomBool) set(value bool) {
	var i int32
	if value {
		i = 1
	}
	atomic.StoreInt32(&(b.flag), int32(i))
}

func (b *tAtomBool) get() bool {
	if atomic.LoadInt32(&(b.flag)) != 0 {
		return true
	}
	return false
}
<|endoftext|>