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

package net

import (
	"flag"
	"fmt"
	"io"
	"os"
	"regexp"
	"runtime"
	"testing"
	"time"
)

func newLocalListener(t *testing.T) Listener {
	ln, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		ln, err = Listen("tcp6", "[::1]:0")
	}
	if err != nil {
		t.Fatal(err)
	}
	return ln
}

func TestDialTimeout(t *testing.T) {
	ln := newLocalListener(t)
	defer ln.Close()

	errc := make(chan error)

	numConns := listenerBacklog + 10

	// TODO(bradfitz): It's hard to test this in a portable
	// way. This is unfortunate, but works for now.
	switch runtime.GOOS {
	case "linux":
		// The kernel will start accepting TCP connections before userspace
		// gets a chance to not accept them, so fire off a bunch to fill up
		// the kernel's backlog.  Then we test we get a failure after that.
		for i := 0; i < numConns; i++ {
			go func() {
				_, err := DialTimeout("tcp", ln.Addr().String(), 200*time.Millisecond)
				errc <- err
			}()
		}
	case "darwin", "windows":
		// At least OS X 10.7 seems to accept any number of
		// connections, ignoring listen's backlog, so resort
		// to connecting to a hopefully-dead 127/8 address.
		// Same for windows.
		//
		// Use an IANA reserved port (49151) instead of 80, because
		// on our 386 builder, this Dial succeeds, connecting
		// to an IIS web server somewhere.  The data center
		// or VM or firewall must be stealing the TCP connection.
		//
		// IANA Service Name and Transport Protocol Port Number Registry
		// <http://www.iana.org/assignments/service-names-port-numbers/service-names-port-numbers.xml>
		go func() {
			c, err := DialTimeout("tcp", "127.0.71.111:49151", 200*time.Millisecond)
			if err == nil {
				err = fmt.Errorf("unexpected: connected to %s!", c.RemoteAddr())
				c.Close()
			}
			errc <- err
		}()
	default:
		// TODO(bradfitz):
		// OpenBSD may have a reject route to 127/8 except 127.0.0.1/32
		// by default. FreeBSD likely works, but is untested.
		// TODO(rsc):
		// The timeout never happens on Windows.  Why?  Issue 3016.
		t.Skipf("skipping test on %q; untested.", runtime.GOOS)
	}

	connected := 0
	for {
		select {
		case <-time.After(15 * time.Second):
			t.Fatal("too slow")
		case err := <-errc:
			if err == nil {
				connected++
				if connected == numConns {
					t.Fatal("all connections connected; expected some to time out")
				}
			} else {
				terr, ok := err.(timeout)
				if !ok {
					t.Fatalf("got error %q; want error with timeout interface", err)
				}
				if !terr.Timeout() {
					t.Fatalf("got error %q; not a timeout", err)
				}
				// Pass. We saw a timeout error.
				return
			}
		}
	}
}

func TestSelfConnect(t *testing.T) {
	if runtime.GOOS == "windows" {
		// TODO(brainman): do not know why it hangs.
		t.Skip("skipping known-broken test on windows")
	}
	// Test that Dial does not honor self-connects.
	// See the comment in DialTCP.

	// Find a port that would be used as a local address.
	l, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	c, err := Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	addr := c.LocalAddr().String()
	c.Close()
	l.Close()

	// Try to connect to that address repeatedly.
	n := 100000
	if testing.Short() {
		n = 1000
	}
	switch runtime.GOOS {
	case "darwin", "freebsd", "netbsd", "openbsd", "plan9", "windows":
		// Non-Linux systems take a long time to figure
		// out that there is nothing listening on localhost.
		n = 100
	}
	for i := 0; i < n; i++ {
		c, err := Dial("tcp", addr)
		if err == nil {
			c.Close()
			t.Errorf("#%d: Dial %q succeeded", i, addr)
		}
	}
}

var runErrorTest = flag.Bool("run_error_test", false, "let TestDialError check for dns errors")

type DialErrorTest struct {
	Net     string
	Raddr   string
	Pattern string
}

var dialErrorTests = []DialErrorTest{
	{
		"datakit", "mh/astro/r70",
		"dial datakit mh/astro/r70: unknown network datakit",
	},
	{
		"tcp", "127.0.0.1:☺",
		"dial tcp 127.0.0.1:☺: unknown port tcp/☺",
	},
	{
		"tcp", "no-such-name.google.com.:80",
		"dial tcp no-such-name.google.com.:80: lookup no-such-name.google.com.( on .*)?: no (.*)",
	},
	{
		"tcp", "no-such-name.no-such-top-level-domain.:80",
		"dial tcp no-such-name.no-such-top-level-domain.:80: lookup no-such-name.no-such-top-level-domain.( on .*)?: no (.*)",
	},
	{
		"tcp", "no-such-name:80",
		`dial tcp no-such-name:80: lookup no-such-name\.(.*\.)?( on .*)?: no (.*)`,
	},
	{
		"tcp", "mh/astro/r70:http",
		"dial tcp mh/astro/r70:http: lookup mh/astro/r70: invalid domain name",
	},
	{
		"unix", "/etc/file-not-found",
		"dial unix /etc/file-not-found: no such file or directory",
	},
	{
		"unix", "/etc/",
		"dial unix /etc/: (permission denied|socket operation on non-socket|connection refused)",
	},
	{
		"unixpacket", "/etc/file-not-found",
		"dial unixpacket /etc/file-not-found: no such file or directory",
	},
	{
		"unixpacket", "/etc/",
		"dial unixpacket /etc/: (permission denied|socket operation on non-socket|connection refused)",
	},
}

var duplicateErrorPattern = `dial (.*) dial (.*)`

func TestDialError(t *testing.T) {
	if !*runErrorTest {
		t.Logf("test disabled; use -run_error_test to enable")
		return
	}
	for i, tt := range dialErrorTests {
		c, err := Dial(tt.Net, tt.Raddr)
		if c != nil {
			c.Close()
		}
		if err == nil {
			t.Errorf("#%d: nil error, want match for %#q", i, tt.Pattern)
			continue
		}
		s := err.Error()
		match, _ := regexp.MatchString(tt.Pattern, s)
		if !match {
			t.Errorf("#%d: %q, want match for %#q", i, s, tt.Pattern)
		}
		match, _ = regexp.MatchString(duplicateErrorPattern, s)
		if match {
			t.Errorf("#%d: %q, duplicate error return from Dial", i, s)
		}
	}
}

func TestDialTimeoutFDLeak(t *testing.T) {
	if runtime.GOOS != "linux" {
		// TODO(bradfitz): test on other platforms
		t.Skipf("skipping test on %q", runtime.GOOS)
	}

	ln := newLocalListener(t)
	defer ln.Close()

	type connErr struct {
		conn Conn
		err  error
	}
	dials := listenerBacklog + 100
	// used to be listenerBacklog + 5, but was found to be unreliable, issue 4384.
	maxGoodConnect := listenerBacklog + runtime.NumCPU()*10
	resc := make(chan connErr)
	for i := 0; i < dials; i++ {
		go func() {
			conn, err := DialTimeout("tcp", ln.Addr().String(), 500*time.Millisecond)
			resc <- connErr{conn, err}
		}()
	}

	var firstErr string
	var ngood int
	var toClose []io.Closer
	for i := 0; i < dials; i++ {
		ce := <-resc
		if ce.err == nil {
			ngood++
			if ngood > maxGoodConnect {
				t.Errorf("%d good connects; expected at most %d", ngood, maxGoodConnect)
			}
			toClose = append(toClose, ce.conn)
			continue
		}
		err := ce.err
		if firstErr == "" {
			firstErr = err.Error()
		} else if err.Error() != firstErr {
			t.Fatalf("inconsistent error messages: first was %q, then later %q", firstErr, err)
		}
	}
	for _, c := range toClose {
		c.Close()
	}
	for i := 0; i < 100; i++ {
		if got := numFD(); got < dials {
			// Test passes.
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	if got := numFD(); got >= dials {
		t.Errorf("num fds after %d timeouts = %d; want <%d", dials, got, dials)
	}
}

func numFD() int {
	if runtime.GOOS == "linux" {
		f, err := os.Open("/proc/self/fd")
		if err != nil {
			panic(err)
		}
		defer f.Close()
		names, err := f.Readdirnames(0)
		if err != nil {
			panic(err)
		}
		return len(names)
	}
	// All tests using this should be skipped anyway, but:
	panic("numFDs not implemented on " + runtime.GOOS)
}
<commit_msg>net: add test for Dial and Listen arguments<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"flag"
	"fmt"
	"io"
	"os"
	"reflect"
	"regexp"
	"runtime"
	"testing"
	"time"
)

func newLocalListener(t *testing.T) Listener {
	ln, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		ln, err = Listen("tcp6", "[::1]:0")
	}
	if err != nil {
		t.Fatal(err)
	}
	return ln
}

func TestDialTimeout(t *testing.T) {
	ln := newLocalListener(t)
	defer ln.Close()

	errc := make(chan error)

	numConns := listenerBacklog + 10

	// TODO(bradfitz): It's hard to test this in a portable
	// way. This is unfortunate, but works for now.
	switch runtime.GOOS {
	case "linux":
		// The kernel will start accepting TCP connections before userspace
		// gets a chance to not accept them, so fire off a bunch to fill up
		// the kernel's backlog.  Then we test we get a failure after that.
		for i := 0; i < numConns; i++ {
			go func() {
				_, err := DialTimeout("tcp", ln.Addr().String(), 200*time.Millisecond)
				errc <- err
			}()
		}
	case "darwin", "windows":
		// At least OS X 10.7 seems to accept any number of
		// connections, ignoring listen's backlog, so resort
		// to connecting to a hopefully-dead 127/8 address.
		// Same for windows.
		//
		// Use an IANA reserved port (49151) instead of 80, because
		// on our 386 builder, this Dial succeeds, connecting
		// to an IIS web server somewhere.  The data center
		// or VM or firewall must be stealing the TCP connection.
		//
		// IANA Service Name and Transport Protocol Port Number Registry
		// <http://www.iana.org/assignments/service-names-port-numbers/service-names-port-numbers.xml>
		go func() {
			c, err := DialTimeout("tcp", "127.0.71.111:49151", 200*time.Millisecond)
			if err == nil {
				err = fmt.Errorf("unexpected: connected to %s!", c.RemoteAddr())
				c.Close()
			}
			errc <- err
		}()
	default:
		// TODO(bradfitz):
		// OpenBSD may have a reject route to 127/8 except 127.0.0.1/32
		// by default. FreeBSD likely works, but is untested.
		// TODO(rsc):
		// The timeout never happens on Windows.  Why?  Issue 3016.
		t.Skipf("skipping test on %q; untested.", runtime.GOOS)
	}

	connected := 0
	for {
		select {
		case <-time.After(15 * time.Second):
			t.Fatal("too slow")
		case err := <-errc:
			if err == nil {
				connected++
				if connected == numConns {
					t.Fatal("all connections connected; expected some to time out")
				}
			} else {
				terr, ok := err.(timeout)
				if !ok {
					t.Fatalf("got error %q; want error with timeout interface", err)
				}
				if !terr.Timeout() {
					t.Fatalf("got error %q; not a timeout", err)
				}
				// Pass. We saw a timeout error.
				return
			}
		}
	}
}

func TestSelfConnect(t *testing.T) {
	if runtime.GOOS == "windows" {
		// TODO(brainman): do not know why it hangs.
		t.Skip("skipping known-broken test on windows")
	}
	// Test that Dial does not honor self-connects.
	// See the comment in DialTCP.

	// Find a port that would be used as a local address.
	l, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	c, err := Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	addr := c.LocalAddr().String()
	c.Close()
	l.Close()

	// Try to connect to that address repeatedly.
	n := 100000
	if testing.Short() {
		n = 1000
	}
	switch runtime.GOOS {
	case "darwin", "freebsd", "netbsd", "openbsd", "plan9", "windows":
		// Non-Linux systems take a long time to figure
		// out that there is nothing listening on localhost.
		n = 100
	}
	for i := 0; i < n; i++ {
		c, err := Dial("tcp", addr)
		if err == nil {
			c.Close()
			t.Errorf("#%d: Dial %q succeeded", i, addr)
		}
	}
}

var runErrorTest = flag.Bool("run_error_test", false, "let TestDialError check for dns errors")

type DialErrorTest struct {
	Net     string
	Raddr   string
	Pattern string
}

var dialErrorTests = []DialErrorTest{
	{
		"datakit", "mh/astro/r70",
		"dial datakit mh/astro/r70: unknown network datakit",
	},
	{
		"tcp", "127.0.0.1:☺",
		"dial tcp 127.0.0.1:☺: unknown port tcp/☺",
	},
	{
		"tcp", "no-such-name.google.com.:80",
		"dial tcp no-such-name.google.com.:80: lookup no-such-name.google.com.( on .*)?: no (.*)",
	},
	{
		"tcp", "no-such-name.no-such-top-level-domain.:80",
		"dial tcp no-such-name.no-such-top-level-domain.:80: lookup no-such-name.no-such-top-level-domain.( on .*)?: no (.*)",
	},
	{
		"tcp", "no-such-name:80",
		`dial tcp no-such-name:80: lookup no-such-name\.(.*\.)?( on .*)?: no (.*)`,
	},
	{
		"tcp", "mh/astro/r70:http",
		"dial tcp mh/astro/r70:http: lookup mh/astro/r70: invalid domain name",
	},
	{
		"unix", "/etc/file-not-found",
		"dial unix /etc/file-not-found: no such file or directory",
	},
	{
		"unix", "/etc/",
		"dial unix /etc/: (permission denied|socket operation on non-socket|connection refused)",
	},
	{
		"unixpacket", "/etc/file-not-found",
		"dial unixpacket /etc/file-not-found: no such file or directory",
	},
	{
		"unixpacket", "/etc/",
		"dial unixpacket /etc/: (permission denied|socket operation on non-socket|connection refused)",
	},
}

var duplicateErrorPattern = `dial (.*) dial (.*)`

func TestDialError(t *testing.T) {
	if !*runErrorTest {
		t.Logf("test disabled; use -run_error_test to enable")
		return
	}
	for i, tt := range dialErrorTests {
		c, err := Dial(tt.Net, tt.Raddr)
		if c != nil {
			c.Close()
		}
		if err == nil {
			t.Errorf("#%d: nil error, want match for %#q", i, tt.Pattern)
			continue
		}
		s := err.Error()
		match, _ := regexp.MatchString(tt.Pattern, s)
		if !match {
			t.Errorf("#%d: %q, want match for %#q", i, s, tt.Pattern)
		}
		match, _ = regexp.MatchString(duplicateErrorPattern, s)
		if match {
			t.Errorf("#%d: %q, duplicate error return from Dial", i, s)
		}
	}
}

var invalidDialAndListenArgTests = []struct {
	net  string
	addr string
	err  error
}{
	{"foo", "bar", &OpError{Op: "dial", Net: "foo", Addr: nil, Err: UnknownNetworkError("foo")}},
	{"baz", "", &OpError{Op: "listen", Net: "baz", Addr: nil, Err: UnknownNetworkError("baz")}},
	{"tcp", "", &OpError{Op: "dial", Net: "tcp", Addr: nil, Err: errMissingAddress}},
}

func TestInvalidDialAndListenArgs(t *testing.T) {
	for _, tt := range invalidDialAndListenArgTests {
		var err error
		switch tt.err.(*OpError).Op {
		case "dial":
			_, err = Dial(tt.net, tt.addr)
		case "listen":
			_, err = Listen(tt.net, tt.addr)
		}
		if !reflect.DeepEqual(tt.err, err) {
			t.Fatalf("got %#v; expected %#v", err, tt.err)
		}
	}
}

func TestDialTimeoutFDLeak(t *testing.T) {
	if runtime.GOOS != "linux" {
		// TODO(bradfitz): test on other platforms
		t.Skipf("skipping test on %q", runtime.GOOS)
	}

	ln := newLocalListener(t)
	defer ln.Close()

	type connErr struct {
		conn Conn
		err  error
	}
	dials := listenerBacklog + 100
	// used to be listenerBacklog + 5, but was found to be unreliable, issue 4384.
	maxGoodConnect := listenerBacklog + runtime.NumCPU()*10
	resc := make(chan connErr)
	for i := 0; i < dials; i++ {
		go func() {
			conn, err := DialTimeout("tcp", ln.Addr().String(), 500*time.Millisecond)
			resc <- connErr{conn, err}
		}()
	}

	var firstErr string
	var ngood int
	var toClose []io.Closer
	for i := 0; i < dials; i++ {
		ce := <-resc
		if ce.err == nil {
			ngood++
			if ngood > maxGoodConnect {
				t.Errorf("%d good connects; expected at most %d", ngood, maxGoodConnect)
			}
			toClose = append(toClose, ce.conn)
			continue
		}
		err := ce.err
		if firstErr == "" {
			firstErr = err.Error()
		} else if err.Error() != firstErr {
			t.Fatalf("inconsistent error messages: first was %q, then later %q", firstErr, err)
		}
	}
	for _, c := range toClose {
		c.Close()
	}
	for i := 0; i < 100; i++ {
		if got := numFD(); got < dials {
			// Test passes.
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	if got := numFD(); got >= dials {
		t.Errorf("num fds after %d timeouts = %d; want <%d", dials, got, dials)
	}
}

func numFD() int {
	if runtime.GOOS == "linux" {
		f, err := os.Open("/proc/self/fd")
		if err != nil {
			panic(err)
		}
		defer f.Close()
		names, err := f.Readdirnames(0)
		if err != nil {
			panic(err)
		}
		return len(names)
	}
	// All tests using this should be skipped anyway, but:
	panic("numFDs not implemented on " + runtime.GOOS)
}
<|endoftext|>


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

package net

import (
	"testing"
)

type portTest struct {
	netw string
	name string
	port int
	ok   bool
}

var porttests = []portTest{
	portTest{"tcp", "echo", 7, true},
	portTest{"tcp", "discard", 9, true},
	portTest{"tcp", "systat", 11, true},
	portTest{"tcp", "daytime", 13, true},
	portTest{"tcp", "chargen", 19, true},
	portTest{"tcp", "ftp-data", 20, true},
	portTest{"tcp", "ftp", 21, true},
	portTest{"tcp", "ssh", 22, true},
	portTest{"tcp", "telnet", 23, true},
	portTest{"tcp", "smtp", 25, true},
	portTest{"tcp", "time", 37, true},
	portTest{"tcp", "domain", 53, true},
	portTest{"tcp", "gopher", 70, true},
	portTest{"tcp", "finger", 79, true},
	portTest{"tcp", "http", 80, true},

	portTest{"udp", "echo", 7, true},
	portTest{"udp", "tacacs", 49, true},
	portTest{"udp", "tftp", 69, true},
	portTest{"udp", "bootpc", 68, true},
	portTest{"udp", "bootps", 67, true},
	portTest{"udp", "domain", 53, true},
	portTest{"udp", "ntp", 123, true},
	portTest{"udp", "snmp", 161, true},
	portTest{"udp", "syslog", 514, true},
	portTest{"udp", "nfs", 2049, true},

	portTest{"--badnet--", "zzz", 0, false},
	portTest{"tcp", "--badport--", 0, false},
}

func TestLookupPort(t *testing.T) {
	for i := 0; i < len(porttests); i++ {
		tt := porttests[i]
		if port, err := LookupPort(tt.netw, tt.name); port != tt.port || (err == nil) != tt.ok {
			t.Errorf("LookupPort(%q, %q) = %v, %s; want %v",
				tt.netw, tt.name, port, err, tt.port)
		}
	}
}
<commit_msg>net: drop non-RHEL-4.6 ports from test<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"testing"
)

type portTest struct {
	netw string
	name string
	port int
	ok   bool
}

var porttests = []portTest{
	portTest{"tcp", "echo", 7, true},
	portTest{"tcp", "discard", 9, true},
	portTest{"tcp", "systat", 11, true},
	portTest{"tcp", "daytime", 13, true},
	portTest{"tcp", "chargen", 19, true},
	portTest{"tcp", "ftp-data", 20, true},
	portTest{"tcp", "ftp", 21, true},
	portTest{"tcp", "ssh", 22, true},
	portTest{"tcp", "telnet", 23, true},
	portTest{"tcp", "smtp", 25, true},
	portTest{"tcp", "time", 37, true},
	portTest{"tcp", "domain", 53, true},
	portTest{"tcp", "gopher", 70, true},
	portTest{"tcp", "finger", 79, true},
	portTest{"tcp", "http", 80, true},

	portTest{"udp", "echo", 7, true},
	portTest{"udp", "tftp", 69, true},
	portTest{"udp", "bootpc", 68, true},
	portTest{"udp", "bootps", 67, true},
	portTest{"udp", "domain", 53, true},
	portTest{"udp", "ntp", 123, true},
	portTest{"udp", "snmp", 161, true},
	portTest{"udp", "syslog", 514, true},

	portTest{"--badnet--", "zzz", 0, false},
	portTest{"tcp", "--badport--", 0, false},
}

func TestLookupPort(t *testing.T) {
	for i := 0; i < len(porttests); i++ {
		tt := porttests[i]
		if port, err := LookupPort(tt.netw, tt.name); port != tt.port || (err == nil) != tt.ok {
			t.Errorf("LookupPort(%q, %q) = %v, %s; want %v",
				tt.netw, tt.name, port, err, tt.port)
		}
	}
}
<|endoftext|>


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

/*
	Package unsafe contains operations that step around the type safety of Go programs.
*/
package unsafe

// ArbitraryType is here for the purposes of documentation only and is not actually
// part of the unsafe package.  It represents the type of an arbitrary Go expression.
type ArbitraryType int

// Pointer represents a pointer to an arbitrary type.  There are three special operations
// available for type Pointer that are not available for other types.
//	1) A pointer value of any type can be converted to a Pointer.
//	2) A Pointer can be converted to a pointer value of any type.
//	3) A uintptr can be converted to a Pointer.
//	4) A Pointer can be converted to a uintptr.
// Pointer therefore allows a program to defeat the type system and read and write
// arbitrary memory. It should be used with extreme care.
type Pointer *ArbitraryType

// Sizeof returns the size in bytes occupied by the value v.  The size is that of the
// "top level" of the value only.  For instance, if v is a slice, it returns the size of
// the slice descriptor, not the size of the memory referenced by the slice.
func Sizeof(v ArbitraryType) uintptr

// Offsetof returns the offset within the struct of the field represented by v,
// which must be of the form structValue.field.  In other words, it returns the
// number of bytes between the start of the struct and the start of the field.
func Offsetof(v ArbitraryType) uintptr

// Alignof returns the alignment of the value v.  It is the maximum value m such
// that the address of a variable with the type of v will always be zero mod m.
// If v is of the form structValue.field, it returns the alignment of field f within struct object obj.
func Alignof(v ArbitraryType) uintptr
<commit_msg>unsafe: fix a typo<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
	Package unsafe contains operations that step around the type safety of Go programs.
*/
package unsafe

// ArbitraryType is here for the purposes of documentation only and is not actually
// part of the unsafe package.  It represents the type of an arbitrary Go expression.
type ArbitraryType int

// Pointer represents a pointer to an arbitrary type.  There are four special operations
// available for type Pointer that are not available for other types.
//	1) A pointer value of any type can be converted to a Pointer.
//	2) A Pointer can be converted to a pointer value of any type.
//	3) A uintptr can be converted to a Pointer.
//	4) A Pointer can be converted to a uintptr.
// Pointer therefore allows a program to defeat the type system and read and write
// arbitrary memory. It should be used with extreme care.
type Pointer *ArbitraryType

// Sizeof returns the size in bytes occupied by the value v.  The size is that of the
// "top level" of the value only.  For instance, if v is a slice, it returns the size of
// the slice descriptor, not the size of the memory referenced by the slice.
func Sizeof(v ArbitraryType) uintptr

// Offsetof returns the offset within the struct of the field represented by v,
// which must be of the form structValue.field.  In other words, it returns the
// number of bytes between the start of the struct and the start of the field.
func Offsetof(v ArbitraryType) uintptr

// Alignof returns the alignment of the value v.  It is the maximum value m such
// that the address of a variable with the type of v will always be zero mod m.
// If v is of the form structValue.field, it returns the alignment of field f within struct object obj.
func Alignof(v ArbitraryType) uintptr
<|endoftext|>


<commit_before>package zk

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"
	"time"
)

type TestServer struct {
	Port int
	Path string
	Srv  *Server
}

type TestCluster struct {
	Path    string
	Servers []TestServer
}

func StartTestCluster(size int, stdout, stderr io.Writer) (*TestCluster, error) {
	tmpPath, err := ioutil.TempDir("", "gozk")
	if err != nil {
		return nil, err
	}
	success := false
	startPort := int(rand.Int31n(6000) + 10000)
	cluster := &TestCluster{Path: tmpPath}
	defer func() {
		if !success {
			cluster.Stop()
		}
	}()
	for serverN := 0; serverN < size; serverN++ {
		srvPath := filepath.Join(tmpPath, fmt.Sprintf("srv%d", serverN))
		if err := os.Mkdir(srvPath, 0700); err != nil {
			return nil, err
		}
		port := startPort + serverN*3
		cfg := ServerConfig{
			ClientPort: port,
			DataDir:    srvPath,
		}
		for i := 0; i < size; i++ {
			cfg.Servers = append(cfg.Servers, ServerConfigServer{
				ID:                 i + 1,
				Host:               "127.0.0.1",
				PeerPort:           startPort + i*3 + 1,
				LeaderElectionPort: startPort + i*3 + 2,
			})
		}
		cfgPath := filepath.Join(srvPath, "zoo.cfg")
		fi, err := os.Create(cfgPath)
		if err != nil {
			return nil, err
		}
		err = cfg.Marshall(fi)
		fi.Close()
		if err != nil {
			return nil, err
		}

		fi, err = os.Create(filepath.Join(srvPath, "myid"))
		if err != nil {
			return nil, err
		}
		_, err = fmt.Fprintf(fi, "%d\n", serverN+1)
		fi.Close()
		if err != nil {
			return nil, err
		}

		srv := &Server{
			ConfigPath: cfgPath,
			Stdout:     stdout,
			Stderr:     stderr,
		}
		if err := srv.Start(); err != nil {
			return nil, err
		}
		cluster.Servers = append(cluster.Servers, TestServer{
			Path: srvPath,
			Port: cfg.ClientPort,
			Srv:  srv,
		})
	}
	success = true
	time.Sleep(time.Second) // Give the server time to become active. Should probably actually attempt to connect to verify.
	return cluster, nil
}

func (ts *TestCluster) Connect(idx int) (*Conn, error) {
	zk, _, err := Connect([]string{fmt.Sprintf("127.0.0.1:%d", ts.Servers[idx].Port)}, time.Second*15)
	return zk, err
}

func (ts *TestCluster) ConnectAll() (*Conn, <-chan Event, error) {
	return ts.ConnectAllTimeout(time.Second * 15)
}

func (ts *TestCluster) ConnectAllTimeout(sessionTimeout time.Duration) (*Conn, <-chan Event, error) {
	hosts := make([]string, len(ts.Servers))
	for i, srv := range ts.Servers {
		hosts[i] = fmt.Sprintf("127.0.0.1:%d", srv.Port)
	}
	zk, ch, err := Connect(hosts, sessionTimeout)
	return zk, ch, err
}

func (ts *TestCluster) Stop() error {
	for _, srv := range ts.Servers {
		srv.Srv.Stop()
	}
	defer os.RemoveAll(ts.Path)
	return nil
}
<commit_msg>More careful checks that test cluster is up/down before/after tests<commit_after>package zk

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"
	"time"
)

type TestServer struct {
	Port int
	Path string
	Srv  *Server
}

type TestCluster struct {
	Path    string
	Servers []TestServer
}

func StartTestCluster(size int, stdout, stderr io.Writer) (*TestCluster, error) {
	tmpPath, err := ioutil.TempDir("", "gozk")
	if err != nil {
		return nil, err
	}
	success := false
	startPort := int(rand.Int31n(6000) + 10000)
	cluster := &TestCluster{Path: tmpPath}
	defer func() {
		if !success {
			cluster.Stop()
		}
	}()
	for serverN := 0; serverN < size; serverN++ {
		srvPath := filepath.Join(tmpPath, fmt.Sprintf("srv%d", serverN))
		if err := os.Mkdir(srvPath, 0700); err != nil {
			return nil, err
		}
		port := startPort + serverN*3
		cfg := ServerConfig{
			ClientPort: port,
			DataDir:    srvPath,
		}
		for i := 0; i < size; i++ {
			cfg.Servers = append(cfg.Servers, ServerConfigServer{
				ID:                 i + 1,
				Host:               "127.0.0.1",
				PeerPort:           startPort + i*3 + 1,
				LeaderElectionPort: startPort + i*3 + 2,
			})
		}
		cfgPath := filepath.Join(srvPath, "zoo.cfg")
		fi, err := os.Create(cfgPath)
		if err != nil {
			return nil, err
		}
		err = cfg.Marshall(fi)
		fi.Close()
		if err != nil {
			return nil, err
		}

		fi, err = os.Create(filepath.Join(srvPath, "myid"))
		if err != nil {
			return nil, err
		}
		_, err = fmt.Fprintf(fi, "%d\n", serverN+1)
		fi.Close()
		if err != nil {
			return nil, err
		}

		srv := &Server{
			ConfigPath: cfgPath,
			Stdout:     stdout,
			Stderr:     stderr,
		}
		if err := srv.Start(); err != nil {
			return nil, err
		}
		cluster.Servers = append(cluster.Servers, TestServer{
			Path: srvPath,
			Port: cfg.ClientPort,
			Srv:  srv,
		})
	}
	if err := cluster.waitForStart(5, time.Second); err != nil {
		return nil, err
	}
	success = true
	return cluster, nil
}

func (ts *TestCluster) Connect(idx int) (*Conn, error) {
	zk, _, err := Connect([]string{fmt.Sprintf("127.0.0.1:%d", ts.Servers[idx].Port)}, time.Second*15)
	return zk, err
}

func (ts *TestCluster) ConnectAll() (*Conn, <-chan Event, error) {
	return ts.ConnectAllTimeout(time.Second * 15)
}

func (ts *TestCluster) ConnectAllTimeout(sessionTimeout time.Duration) (*Conn, <-chan Event, error) {
	hosts := make([]string, len(ts.Servers))
	for i, srv := range ts.Servers {
		hosts[i] = fmt.Sprintf("127.0.0.1:%d", srv.Port)
	}
	zk, ch, err := Connect(hosts, sessionTimeout)
	return zk, ch, err
}

func (ts *TestCluster) Stop() error {
	for _, srv := range ts.Servers {
		srv.Srv.Stop()
	}
	defer os.RemoveAll(ts.Path)
	return ts.waitForStop(5, 1*time.Second)
}

// block until the cluster is up
func (ts *TestCluster) waitForStart(maxRetry int, interval time.Duration) error {
	// verify that the servers are up with SRVR
	serverAddrs := make([]string, len(ts.Servers))
	for i, s := range ts.Servers {
		serverAddrs[i] = fmt.Sprintf("127.0.0.1:%d", s.Port)
	}

	for i := 0; i < maxRetry; i++ {
		_, ok := FLWSrvr(serverAddrs, time.Second)
		if ok {
			return nil
		}
		time.Sleep(interval)
	}
	return fmt.Errorf("unable to verify health of servers!")
}

// block until the cluster is down
func (ts *TestCluster) waitForStop(maxRetry int, interval time.Duration) error {
	// verify that the servers are up with RUOK
	serverAddrs := make([]string, len(ts.Servers))
	for i, s := range ts.Servers {
		serverAddrs[i] = fmt.Sprintf("127.0.0.1:%d", s.Port)
	}

	var success bool
	for i := 0; i < maxRetry && !success; i++ {
		success = true
		for _, ok := range FLWRuok(serverAddrs, time.Second) {
			if ok {
				success = false
			}
		}
		if !success {
			time.Sleep(interval)
		}
	}
	if !success {
		return fmt.Errorf("unable to verify servers are down!")
	}
	return nil
}
<|endoftext|>


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

package cmd

import (
	"errors"
	"fmt"
	"time"

	"github.com/buger/goterm"
	"github.com/ligato/cn-infra/servicelabel"
	"github.com/ligato/vpp-agent/cmd/agentctl/utils"
	"github.com/spf13/cobra"
)

var watchCommand = &cobra.Command{
	Use:     "watch [agent-label-filter(s)] [interface-filter(s)]",
	Aliases: []string{"w", "wa"},
	Short:   "Show real-time traffic data for vpp interfaces",
	Long: `
'watch' displays a table with periodically updated interface statistics
counters. Data is printed for every agent whose configuration contains
at least one interface with statistics (in case no filters are applied).

Data can be filtered using one or more agent's microservice labels (only
data for specified labels are shown) and, in case microservice label is
defined, it is possible to filter interfaces using the second argument
[interface-filter(s)] for that agent. It is also possible to filter multiple
interfaces, however, it is not allowed to filter interfaces if more than
one agent is specified (see examples)

Note that agent's configuration data is stored in Etcd by a 3rd party
orchestrator. Agent's state data is periodically updated in Etcd by agents
themselves every 10 seconds. `,
	Example: `  Watch all agent with all interfaces containing statistics:
    $ agentctl watch
  Watch statistics for the specific agent:
    $ agentctl watch vpp
  Watch statistics for multiple agents. Note: in this case is not possible
  to filter interfaces:
    $ agentctl watch vpp1,vpp2
    $ agentctl watch "vpp1, vpp2"
    $ agentctl watch 'vpp1, vpp2'
  Watch statistics for the specific interface:
    $ agentctl watch vpp interface
  Watch statistics for multiple interfaces:
    $ agentctl watch vpp interface1,interface2
    $ agentctl watch vpp "interface1, interface2"
    $ agentctl watch vpp 'interface1,interface2'
  Show shortened output (in_packets, out_packets and drop only). Can be
  arbitrarily combined with filters.
    $ agentctl watch --short
  Filter all interfaces/statistic types with full zero values. Can be arbitrarily
  combined with filters. Note: if used together with --short, it is
  possible to see full zero output because there could be fields which are
  non-zero and not shown in shortened output:
    $ agentctl watch --active`,
	Run: watchFunc,
}

var (
	showShort  bool
	showActive bool
)

func init() {
	RootCmd.AddCommand(watchCommand)
	watchCommand.Flags().BoolVarP(&showShort, "short", "s", false,
		"show shortened list of statistics for every vpp/interface")
	watchCommand.Flags().BoolVarP(&showActive, "active", "a", false,
		"only vpp/interfaces with non-zero statistics will be displayed")
}

func watchFunc(cmd *cobra.Command, args []string) {
	if len(args) > 2 {
		utils.ExitWithError(utils.ExitBadArgs, errors.New("Too many arguments, use 'show <label> <interface>'"))
	}
	// Clear current screen
	goterm.Clear()
	for {
		// Get console dimensions
		width := goterm.Width()
		height := goterm.Height()

		if width != 0 && height != 0 {
			goterm.Clear()
		}
		// This hooks watcher to the top of the console window
		goterm.MoveCursor(1, 1)
		table := goterm.NewTable(0, 10, 2, ' ', 0)

		if width < 145 && !showShort {
			fmt.Fprintf(table, "%s", utils.NoSpace)
		} else {

			etcdDump := getLatestEtcdData()

			if showShort {
				_, err := etcdDump.PrintDataAsTable(table, args, showShort, showActive)
				if err != nil {
					utils.ExitWithError(utils.ExitError, errors.New(err.Error()))
				}
			} else {
				_, err := etcdDump.PrintDataAsTable(table, args, showShort, showActive)
				if err != nil {
					utils.ExitWithError(utils.ExitError, errors.New(err.Error()))
				}
			}
		}

		goterm.Println(table)
		goterm.Flush()
		fmt.Println("Press Ctrl-C to exit watcher ...")
		// ETCD interface statistics are updated every 10 seconds (interval set directly in VPP)
		time.Sleep(time.Second)
	}
}

// Read latest ETCD data to keep counters updated
func getLatestEtcdData() utils.EtcdDump {
	// Get data broker
	dataBroker, err := utils.GetDbForAllAgents(globalFlags.Endpoints)
	if err != nil {
		utils.ExitWithError(utils.ExitError, errors.New("Failed to connect to Etcd - "+err.Error()))
	}
	// Read agent prefixes
	keyIter, err := dataBroker.ListKeys(servicelabel.GetAllAgentsPrefix())
	if err != nil {
		utils.ExitWithError(utils.ExitError, errors.New("Failed to get keys - "+err.Error()))
	}
	etcdDump := utils.NewEtcdDump()
	for {
		if key, _, done := keyIter.GetNext(); !done {
			if _, err = etcdDump.ReadDataFromDb(dataBroker, key, nil, nil); err != nil {
				utils.ExitWithError(utils.ExitError, err)
			}
			continue
		}
		break
	}
	return etcdDump
}
<commit_msg>ODPM-555: agentctl's watch cmd: enforce data flush and screen reset during periodic updates<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"errors"
	"fmt"
	"time"

	"strings"

	"github.com/buger/goterm"
	"github.com/ligato/cn-infra/servicelabel"
	"github.com/ligato/vpp-agent/cmd/agentctl/utils"
	"github.com/spf13/cobra"
)

var watchCommand = &cobra.Command{
	Use:     "watch [agent-label-filter(s)] [interface-filter(s)]",
	Aliases: []string{"w", "wa"},
	Short:   "Show real-time traffic data for vpp interfaces",
	Long: `
'watch' displays a table with periodically updated interface statistics
counters. Data is printed for every agent whose configuration contains
at least one interface with statistics (in case no filters are applied).

Data can be filtered using one or more agent's microservice labels (only
data for specified labels are shown) and, in case microservice label is
defined, it is possible to filter interfaces using the second argument
[interface-filter(s)] for that agent. It is also possible to filter multiple
interfaces, however, it is not allowed to filter interfaces if more than
one agent is specified (see examples)

Note that agent's configuration data is stored in Etcd by a 3rd party
orchestrator. Agent's state data is periodically updated in Etcd by agents
themselves every 10 seconds. `,
	Example: `  Watch all agent with all interfaces containing statistics:
    $ agentctl watch
  Watch statistics for the specific agent:
    $ agentctl watch vpp
  Watch statistics for multiple agents. Note: in this case is not possible
  to filter interfaces:
    $ agentctl watch vpp1,vpp2
    $ agentctl watch "vpp1, vpp2"
    $ agentctl watch 'vpp1, vpp2'
  Watch statistics for the specific interface:
    $ agentctl watch vpp interface
  Watch statistics for multiple interfaces:
    $ agentctl watch vpp interface1,interface2
    $ agentctl watch vpp "interface1, interface2"
    $ agentctl watch vpp 'interface1,interface2'
  Show shortened output (in_packets, out_packets and drop only). Can be
  arbitrarily combined with filters.
    $ agentctl watch --short
  Filter all interfaces/statistic types with full zero values. Can be arbitrarily
  combined with filters. Note: if used together with --short, it is
  possible to see full zero output because there could be fields which are
  non-zero and not shown in shortened output:
    $ agentctl watch --active`,
	Run: watchFunc,
}

var (
	showShort  bool
	showActive bool
)

func init() {
	RootCmd.AddCommand(watchCommand)
	watchCommand.Flags().BoolVarP(&showShort, "short", "s", false,
		"show shortened list of statistics for every vpp/interface")
	watchCommand.Flags().BoolVarP(&showActive, "active", "a", false,
		"only vpp/interfaces with non-zero statistics will be displayed")
}

func watchFunc(cmd *cobra.Command, args []string) {
	if len(args) > 2 {
		utils.ExitWithError(utils.ExitBadArgs, errors.New("Too many arguments, use 'show <label> <interface>'"))
	}
	// Clear current screen
	goterm.Clear()
	for {
		// Get console dimensions
		width := goterm.Width()
		height := goterm.Height()

		if width != 0 && height != 0 {
			goterm.Clear()
		}
		// This hooks watcher to the top of the console window
		goterm.MoveCursor(1, 1)
		table := goterm.NewTable(0, 10, 2, ' ', 0)

		if width < 145 && !showShort {
			fmt.Fprintf(table, "%s", utils.NoSpace)
		} else {

			etcdDump := getLatestEtcdData()

			if showShort {
				_, err := etcdDump.PrintDataAsTable(table, args, showShort, showActive)
				if err != nil {
					utils.ExitWithError(utils.ExitError, errors.New(err.Error()))
				}
			} else {
				_, err := etcdDump.PrintDataAsTable(table, args, showShort, showActive)
				if err != nil {
					utils.ExitWithError(utils.ExitError, errors.New(err.Error()))
				}
			}
		}

		goterm.Println(table)
		// Flush buffer and ensure that it will not overflow the screen
		for idx, str := range strings.Split(goterm.Screen.String(), "\n") {
			if idx > goterm.Height() {
				return
			}
			goterm.Output.WriteString(str + "\n")
		}
		// Write buffered data
		err := goterm.Output.Flush()
		if err != nil {
			fmt.Errorf("%v", err)
		}
		// Reset buffer to be empty
		goterm.Screen.Reset()

		fmt.Println("Press Ctrl-C to exit watcher ...")

		// ETCD interface statistics are updated every 10 seconds (interval set directly in VPP)
		time.Sleep(time.Second)
	}
}

// Read latest ETCD data to keep counters updated
func getLatestEtcdData() utils.EtcdDump {
	// Get data broker
	dataBroker, err := utils.GetDbForAllAgents(globalFlags.Endpoints)
	if err != nil {
		utils.ExitWithError(utils.ExitError, errors.New("Failed to connect to Etcd - "+err.Error()))
	}
	// Read agent prefixes
	keyIter, err := dataBroker.ListKeys(servicelabel.GetAllAgentsPrefix())
	if err != nil {
		utils.ExitWithError(utils.ExitError, errors.New("Failed to get keys - "+err.Error()))
	}
	etcdDump := utils.NewEtcdDump()
	for {
		if key, _, done := keyIter.GetNext(); !done {
			if _, err = etcdDump.ReadDataFromDb(dataBroker, key, nil, nil); err != nil {
				utils.ExitWithError(utils.ExitError, err)
			}
			continue
		}
		break
	}
	return etcdDump
}
<|endoftext|>


<commit_before>package msgpack_test

import (
	"bytes"
	"encoding/binary"
	"encoding/csv"
	"encoding/gob"
	"encoding/json"
	"math"
	"testing"
	"time"

	"github.com/vmihailenco/msgpack"
)

func benchmarkEncodeDecode(b *testing.B, src, dst interface{}) {
	var buf bytes.Buffer
	dec := msgpack.NewDecoder(&buf)
	enc := msgpack.NewEncoder(&buf)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if err := enc.Encode(src); err != nil {
			b.Fatal(err)
		}
		if err := dec.Decode(dst); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkBool(b *testing.B) {
	var dst bool
	benchmarkEncodeDecode(b, true, &dst)
}

func BenchmarkInt0(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, 1, &dst)
}

func BenchmarkInt1(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, -33, &dst)
}

func BenchmarkInt2(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, 128, &dst)
}

func BenchmarkInt4(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, 32768, &dst)
}

func BenchmarkInt8(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, int64(2147483648), &dst)
}

func BenchmarkInt32(b *testing.B) {
	var dst int32
	benchmarkEncodeDecode(b, int32(0), &dst)
}

func BenchmarkInt0Binary(b *testing.B) {
	var buf bytes.Buffer
	var out int32
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := binary.Write(&buf, binary.BigEndian, int32(1)); err != nil {
			b.Fatal(err)
		}
		if err := binary.Read(&buf, binary.BigEndian, &out); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkTime(b *testing.B) {
	var dst time.Time
	benchmarkEncodeDecode(b, time.Now(), &dst)
}

func BenchmarkDuration(b *testing.B) {
	var dst time.Duration
	benchmarkEncodeDecode(b, time.Hour, &dst)
}

func BenchmarkByteSlice(b *testing.B) {
	src := make([]byte, 1024)
	var dst []byte
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkByteArray(b *testing.B) {
	var src [1024]byte
	var dst [1024]byte
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkMapStringString(b *testing.B) {
	src := map[string]string{
		"hello": "world",
		"foo":   "bar",
	}
	var dst map[string]string
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkMapStringStringPtr(b *testing.B) {
	src := map[string]string{
		"hello": "world",
		"foo":   "bar",
	}
	var dst map[string]string
	dstptr := &dst
	benchmarkEncodeDecode(b, src, &dstptr)
}

func BenchmarkMapStringInterface(b *testing.B) {
	src := map[string]interface{}{
		"hello": "world",
		"foo":   "bar",
	}
	var dst map[string]interface{}
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkMapIntInt(b *testing.B) {
	src := map[int]int{
		1: 10,
		2: 20,
	}
	var dst map[int]int
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkStringSlice(b *testing.B) {
	src := []string{"hello", "world"}
	var dst []string
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkStringSlicePtr(b *testing.B) {
	src := []string{"hello", "world"}
	var dst []string
	dstptr := &dst
	benchmarkEncodeDecode(b, src, &dstptr)
}

type benchmarkStruct struct {
	Name      string
	Age       int
	Colors    []string
	Data      []byte
	CreatedAt time.Time
	UpdatedAt time.Time
}

type benchmarkStruct2 struct {
	Name      string
	Age       int
	Colors    []string
	Data      []byte
	CreatedAt time.Time
	UpdatedAt time.Time
}

var _ msgpack.CustomEncoder = (*benchmarkStruct2)(nil)
var _ msgpack.CustomDecoder = (*benchmarkStruct2)(nil)

func (s *benchmarkStruct2) EncodeMsgpack(enc *msgpack.Encoder) error {
	return enc.Encode(
		s.Name,
		s.Colors,
		s.Age,
		s.Data,
		s.CreatedAt,
		s.UpdatedAt,
	)
}

func (s *benchmarkStruct2) DecodeMsgpack(dec *msgpack.Decoder) error {
	return dec.Decode(
		&s.Name,
		&s.Colors,
		&s.Age,
		&s.Data,
		&s.CreatedAt,
		&s.UpdatedAt,
	)
}

func structForBenchmark() *benchmarkStruct {
	return &benchmarkStruct{
		Name:      "Hello World",
		Colors:    []string{"red", "orange", "yellow", "green", "blue", "violet"},
		Age:       math.MaxInt32,
		Data:      make([]byte, 1024),
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}
}

func structForBenchmark2() *benchmarkStruct2 {
	return &benchmarkStruct2{
		Name:      "Hello World",
		Colors:    []string{"red", "orange", "yellow", "green", "blue", "violet"},
		Age:       math.MaxInt32,
		Data:      make([]byte, 1024),
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}
}

func BenchmarkStructVmihailencoMsgpack(b *testing.B) {
	in := structForBenchmark()
	out := new(benchmarkStruct)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		buf, err := msgpack.Marshal(in)
		if err != nil {
			b.Fatal(err)
		}

		err = msgpack.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStructMarshal(b *testing.B) {
	in := structForBenchmark()

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		_, err := msgpack.Marshal(in)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStructUnmarshal(b *testing.B) {
	in := structForBenchmark()
	buf, err := msgpack.Marshal(in)
	if err != nil {
		b.Fatal(err)
	}
	out := new(benchmarkStruct)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		err = msgpack.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStructManual(b *testing.B) {
	in := structForBenchmark2()
	out := new(benchmarkStruct2)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		buf, err := msgpack.Marshal(in)
		if err != nil {
			b.Fatal(err)
		}

		err = msgpack.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStructJSON(b *testing.B) {
	in := structForBenchmark()
	out := new(benchmarkStruct)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		buf, err := json.Marshal(in)
		if err != nil {
			b.Fatal(err)
		}

		err = json.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStructGOB(b *testing.B) {
	in := structForBenchmark()
	out := new(benchmarkStruct)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		buf := new(bytes.Buffer)
		enc := gob.NewEncoder(buf)
		dec := gob.NewDecoder(buf)

		if err := enc.Encode(in); err != nil {
			b.Fatal(err)
		}

		if err := dec.Decode(out); err != nil {
			b.Fatal(err)
		}
	}
}

type benchmarkStructPartially struct {
	Name string
	Age  int
}

func BenchmarkStructUnmarshalPartially(b *testing.B) {
	in := structForBenchmark()
	buf, err := msgpack.Marshal(in)
	if err != nil {
		b.Fatal(err)
	}
	out := new(benchmarkStructPartially)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		err = msgpack.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkCSV(b *testing.B) {
	for i := 0; i < b.N; i++ {
		record := []string{"1", "hello", "world"}

		buf := new(bytes.Buffer)
		r := csv.NewReader(buf)
		w := csv.NewWriter(buf)

		if err := w.Write(record); err != nil {
			b.Fatal(err)
		}
		w.Flush()
		if _, err := r.Read(); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkCSVMsgpack(b *testing.B) {
	for i := 0; i < b.N; i++ {
		var num int
		var hello, world string

		buf := new(bytes.Buffer)
		enc := msgpack.NewEncoder(buf)
		dec := msgpack.NewDecoder(buf)

		if err := enc.Encode(1, "hello", "world"); err != nil {
			b.Fatal(err)
		}
		if err := dec.Decode(&num, &hello, &world); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkQuery(b *testing.B) {
	var records []map[string]interface{}
	for i := 0; i < 1000; i++ {
		record := map[string]interface{}{
			"id":    i,
			"attrs": map[string]interface{}{"phone": i},
		}
		records = append(records, record)
	}

	bs, err := msgpack.Marshal(records)
	if err != nil {
		b.Fatal(err)
	}

	dec := msgpack.NewDecoder(bytes.NewBuffer(bs))

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		dec.Reset(bytes.NewBuffer(bs))

		values, err := dec.Query("10.attrs.phone")
		if err != nil {
			b.Fatal(err)
		}
		if values[0].(int8) != 10 {
			b.Fatalf("%v != %d", values[0], 10)
		}
	}
}
<commit_msg>Remove some benchmarks<commit_after>package msgpack_test

import (
	"bytes"
	"math"
	"testing"
	"time"

	"github.com/vmihailenco/msgpack"
)

func benchmarkEncodeDecode(b *testing.B, src, dst interface{}) {
	var buf bytes.Buffer
	dec := msgpack.NewDecoder(&buf)
	enc := msgpack.NewEncoder(&buf)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if err := enc.Encode(src); err != nil {
			b.Fatal(err)
		}
		if err := dec.Decode(dst); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkBool(b *testing.B) {
	var dst bool
	benchmarkEncodeDecode(b, true, &dst)
}

func BenchmarkInt0(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, 1, &dst)
}

func BenchmarkInt1(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, -33, &dst)
}

func BenchmarkInt2(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, 128, &dst)
}

func BenchmarkInt4(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, 32768, &dst)
}

func BenchmarkInt8(b *testing.B) {
	var dst int
	benchmarkEncodeDecode(b, int64(2147483648), &dst)
}

func BenchmarkInt32(b *testing.B) {
	var dst int32
	benchmarkEncodeDecode(b, int32(0), &dst)
}

func BenchmarkTime(b *testing.B) {
	var dst time.Time
	benchmarkEncodeDecode(b, time.Now(), &dst)
}

func BenchmarkDuration(b *testing.B) {
	var dst time.Duration
	benchmarkEncodeDecode(b, time.Hour, &dst)
}

func BenchmarkByteSlice(b *testing.B) {
	src := make([]byte, 1024)
	var dst []byte
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkByteArray(b *testing.B) {
	var src [1024]byte
	var dst [1024]byte
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkMapStringString(b *testing.B) {
	src := map[string]string{
		"hello": "world",
		"foo":   "bar",
	}
	var dst map[string]string
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkMapStringStringPtr(b *testing.B) {
	src := map[string]string{
		"hello": "world",
		"foo":   "bar",
	}
	var dst map[string]string
	dstptr := &dst
	benchmarkEncodeDecode(b, src, &dstptr)
}

func BenchmarkMapStringInterface(b *testing.B) {
	src := map[string]interface{}{
		"hello": "world",
		"foo":   "bar",
	}
	var dst map[string]interface{}
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkMapIntInt(b *testing.B) {
	src := map[int]int{
		1: 10,
		2: 20,
	}
	var dst map[int]int
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkStringSlice(b *testing.B) {
	src := []string{"hello", "world"}
	var dst []string
	benchmarkEncodeDecode(b, src, &dst)
}

func BenchmarkStringSlicePtr(b *testing.B) {
	src := []string{"hello", "world"}
	var dst []string
	dstptr := &dst
	benchmarkEncodeDecode(b, src, &dstptr)
}

type benchmarkStruct struct {
	Name      string
	Age       int
	Colors    []string
	Data      []byte
	CreatedAt time.Time
	UpdatedAt time.Time
}

type benchmarkStruct2 struct {
	Name      string
	Age       int
	Colors    []string
	Data      []byte
	CreatedAt time.Time
	UpdatedAt time.Time
}

var _ msgpack.CustomEncoder = (*benchmarkStruct2)(nil)
var _ msgpack.CustomDecoder = (*benchmarkStruct2)(nil)

func (s *benchmarkStruct2) EncodeMsgpack(enc *msgpack.Encoder) error {
	return enc.Encode(
		s.Name,
		s.Colors,
		s.Age,
		s.Data,
		s.CreatedAt,
		s.UpdatedAt,
	)
}

func (s *benchmarkStruct2) DecodeMsgpack(dec *msgpack.Decoder) error {
	return dec.Decode(
		&s.Name,
		&s.Colors,
		&s.Age,
		&s.Data,
		&s.CreatedAt,
		&s.UpdatedAt,
	)
}

func structForBenchmark() *benchmarkStruct {
	return &benchmarkStruct{
		Name:      "Hello World",
		Colors:    []string{"red", "orange", "yellow", "green", "blue", "violet"},
		Age:       math.MaxInt32,
		Data:      make([]byte, 1024),
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}
}

func structForBenchmark2() *benchmarkStruct2 {
	return &benchmarkStruct2{
		Name:      "Hello World",
		Colors:    []string{"red", "orange", "yellow", "green", "blue", "violet"},
		Age:       math.MaxInt32,
		Data:      make([]byte, 1024),
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}
}

func BenchmarkStructVmihailencoMsgpack(b *testing.B) {
	in := structForBenchmark()
	out := new(benchmarkStruct)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		buf, err := msgpack.Marshal(in)
		if err != nil {
			b.Fatal(err)
		}

		err = msgpack.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStructMarshal(b *testing.B) {
	in := structForBenchmark()

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		_, err := msgpack.Marshal(in)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStructUnmarshal(b *testing.B) {
	in := structForBenchmark()
	buf, err := msgpack.Marshal(in)
	if err != nil {
		b.Fatal(err)
	}
	out := new(benchmarkStruct)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		err = msgpack.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStructManual(b *testing.B) {
	in := structForBenchmark2()
	out := new(benchmarkStruct2)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		buf, err := msgpack.Marshal(in)
		if err != nil {
			b.Fatal(err)
		}

		err = msgpack.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

type benchmarkStructPartially struct {
	Name string
	Age  int
}

func BenchmarkStructUnmarshalPartially(b *testing.B) {
	in := structForBenchmark()
	buf, err := msgpack.Marshal(in)
	if err != nil {
		b.Fatal(err)
	}
	out := new(benchmarkStructPartially)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		err = msgpack.Unmarshal(buf, out)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkQuery(b *testing.B) {
	var records []map[string]interface{}
	for i := 0; i < 1000; i++ {
		record := map[string]interface{}{
			"id":    i,
			"attrs": map[string]interface{}{"phone": i},
		}
		records = append(records, record)
	}

	bs, err := msgpack.Marshal(records)
	if err != nil {
		b.Fatal(err)
	}

	dec := msgpack.NewDecoder(bytes.NewBuffer(bs))

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		dec.Reset(bytes.NewBuffer(bs))

		values, err := dec.Query("10.attrs.phone")
		if err != nil {
			b.Fatal(err)
		}
		if values[0].(int8) != 10 {
			b.Fatalf("%v != %d", values[0], 10)
		}
	}
}
<|endoftext|>


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

package defaultrolemanager

import (
	"strings"
	"sync"

	"github.com/casbin/casbin/v2/errors"
	"github.com/casbin/casbin/v2/log"
	"github.com/casbin/casbin/v2/rbac"
)

const defaultDomain string = "casbin::default"

type MatchingFunc func(arg1, arg2 string) bool

// RoleManager provides a default implementation for the RoleManager interface
type RoleManager struct {
	allDomains         *sync.Map
	maxHierarchyLevel  int
	hasPattern         bool
	matchingFunc       MatchingFunc
	hasDomainPattern   bool
	domainMatchingFunc MatchingFunc

	logger log.Logger
}

// NewRoleManager is the constructor for creating an instance of the
// default RoleManager implementation.
func NewRoleManager(maxHierarchyLevel int) rbac.RoleManager {
	rm := RoleManager{}
	rm.allDomains = &sync.Map{}
	rm.allDomains.LoadOrStore(defaultDomain, new(Roles))
	rm.maxHierarchyLevel = maxHierarchyLevel
	rm.hasPattern = false
	rm.hasDomainPattern = false

	rm.SetLogger(&log.DefaultLogger{})

	return &rm
}

// AddMatchingFunc support use pattern in g
func (rm *RoleManager) AddMatchingFunc(name string, fn MatchingFunc) {
	rm.hasPattern = true
	rm.matchingFunc = fn
}

// AddDomainMatchingFunc support use domain pattern in g
func (rm *RoleManager) AddDomainMatchingFunc(name string, fn MatchingFunc) {
	rm.hasDomainPattern = true
	rm.domainMatchingFunc = fn
}

// SetLogger sets role manager's logger.
func (rm *RoleManager) SetLogger(logger log.Logger) {
	rm.logger = logger
}

func (rm *RoleManager) generateTempRoles(domain string) *Roles {
	rm.allDomains.LoadOrStore(domain, new(Roles))

	patternDomain := []string{domain}
	if rm.hasDomainPattern {
		rm.allDomains.Range(func(key, value interface{}) bool {
			if rm.domainMatchingFunc(domain, key.(string)) {
				patternDomain = append(patternDomain, key.(string))
			}
			return true
		})
	}

	allRoles := new(Roles)
	for _, domain := range patternDomain {
		value, _ := rm.allDomains.LoadOrStore(domain, new(Roles))
		value.(*Roles).Range(func(key, value interface{}) bool {
			role2 := value.(*Role)
			role1 := allRoles.createRole(role2.name, rm.matchingFunc)
			for _, v := range role2.roles {
				role3 := allRoles.createRole(v.name, rm.matchingFunc)
				role1.addRole(role3)
			}

			return true
		})
	}

	return allRoles
}

// Clear clears all stored data and resets the role manager to the initial state.
func (rm *RoleManager) Clear() error {
	rm.allDomains = &sync.Map{}
	rm.allDomains.LoadOrStore(defaultDomain, new(Roles))
	return nil
}

// AddLink adds the inheritance link between role: name1 and role: name2.
// aka role: name1 inherits role: name2.
func (rm *RoleManager) AddLink(name1 string, name2 string, domain ...string) error {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return errors.ERR_DOMAIN_PARAMETER
	}

	allRoles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))

	role1, _ := allRoles.(*Roles).LoadOrStore(name1, newRole(name1))
	role2, _ := allRoles.(*Roles).LoadOrStore(name2, newRole(name2))
	role1.(*Role).addRole(role2.(*Role))
	return nil
}

// DeleteLink deletes the inheritance link between role: name1 and role: name2.
// aka role: name1 does not inherit role: name2 any more.
func (rm *RoleManager) DeleteLink(name1 string, name2 string, domain ...string) error {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return errors.ERR_DOMAIN_PARAMETER
	}

	allRoles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))

	_, ok1 := allRoles.(*Roles).Load(name1)
	_, ok2 := allRoles.(*Roles).Load(name1)
	if !ok1 || !ok2 {
		return errors.ERR_NAMES12_NOT_FOUND
	}

	role1, _ := allRoles.(*Roles).LoadOrStore(name1, newRole(name1))
	role2, _ := allRoles.(*Roles).LoadOrStore(name2, newRole(name2))
	role1.(*Role).deleteRole(role2.(*Role))
	return nil
}

// HasLink determines whether role: name1 inherits role: name2.
func (rm *RoleManager) HasLink(name1 string, name2 string, domain ...string) (bool, error) {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return false, errors.ERR_DOMAIN_PARAMETER
	}

	if name1 == name2 {
		return true, nil
	}

	var allRoles *Roles
	if rm.hasDomainPattern || rm.hasPattern {
		allRoles = rm.generateTempRoles(domain[0])
	} else {
		roles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))
		allRoles = roles.(*Roles)
	}

	if !allRoles.hasRole(name1, rm.matchingFunc) || !allRoles.hasRole(name2, rm.matchingFunc) {
		return false, nil
	}

	role1 := allRoles.createRole(name1, rm.matchingFunc)
	return role1.hasRole(name2, rm.maxHierarchyLevel), nil
}

// GetRoles gets the roles that a subject inherits.
func (rm *RoleManager) GetRoles(name string, domain ...string) ([]string, error) {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return nil, errors.ERR_DOMAIN_PARAMETER
	}

	var allRoles *Roles
	if rm.hasDomainPattern || rm.hasPattern {
		allRoles = rm.generateTempRoles(domain[0])
	} else {
		roles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))
		allRoles = roles.(*Roles)
	}

	if !allRoles.hasRole(name, rm.matchingFunc) {
		return []string{}, nil
	}

	roles := allRoles.createRole(name, rm.matchingFunc).getRoles()

	return roles, nil
}

// GetUsers gets the users that inherits a subject.
// domain is an unreferenced parameter here, may be used in other implementations.
func (rm *RoleManager) GetUsers(name string, domain ...string) ([]string, error) {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return nil, errors.ERR_DOMAIN_PARAMETER
	}

	var allRoles *Roles
	if rm.hasDomainPattern || rm.hasPattern {
		allRoles = rm.generateTempRoles(domain[0])
	} else {
		roles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))
		allRoles = roles.(*Roles)
	}

	if !allRoles.hasRole(name, rm.domainMatchingFunc) {
		return nil, errors.ERR_NAME_NOT_FOUND
	}

	names := []string{}
	allRoles.Range(func(_, value interface{}) bool {
		role := value.(*Role)
		if role.hasDirectRole(name) {
			names = append(names, role.name)
		}
		return true
	})

	return names, nil
}

// PrintRoles prints all the roles to log.
func (rm *RoleManager) PrintRoles() error {
	if !(rm.logger).IsEnabled() {
		return nil
	}

	var roles []string
	rm.allDomains.Range(func(_, value interface{}) bool {
		value.(*Roles).Range(func(_, value interface{}) bool {
			if text := value.(*Role).toString(); text != "" {
				roles = append(roles, text)
			}
			return true
		})

		return true
	})

	rm.logger.LogRole(roles)
	return nil
}

// Roles represents all roles in a domain
type Roles struct {
	sync.Map
}

func (roles *Roles) hasRole(name string, matchingFunc MatchingFunc) bool {
	var ok bool
	if matchingFunc != nil {
		roles.Range(func(key, value interface{}) bool {
			if matchingFunc(name, key.(string)) {
				ok = true
			}
			return true
		})
	} else {
		_, ok = roles.Load(name)
	}

	return ok
}

func (roles *Roles) createRole(name string, matchingFunc MatchingFunc) *Role {
	role, _ := roles.LoadOrStore(name, newRole(name))

	if matchingFunc != nil {
		roles.Range(func(key, value interface{}) bool {
			if matchingFunc(name, key.(string)) && name != key.(string) {
				// Add new role to matching role
				role1, _ := roles.LoadOrStore(key.(string), newRole(key.(string)))
				role.(*Role).addRole(role1.(*Role))
			}
			return true
		})
	}

	return role.(*Role)
}

// Role represents the data structure for a role in RBAC.
type Role struct {
	name  string
	roles []*Role
}

func newRole(name string) *Role {
	r := Role{}
	r.name = name
	return &r
}

func (r *Role) addRole(role *Role) {
	for _, rr := range r.roles {
		if rr.name == role.name {
			return
		}
	}

	r.roles = append(r.roles, role)
}

func (r *Role) deleteRole(role *Role) {
	for i, rr := range r.roles {
		if rr.name == role.name {
			r.roles = append(r.roles[:i], r.roles[i+1:]...)
			return
		}
	}
}

func (r *Role) hasRole(name string, hierarchyLevel int) bool {
	if r.name == name {
		return true
	}

	if hierarchyLevel <= 0 {
		return false
	}

	for _, role := range r.roles {
		if role.hasRole(name, hierarchyLevel-1) {
			return true
		}
	}
	return false
}

func (r *Role) hasDirectRole(name string) bool {
	for _, role := range r.roles {
		if role.name == name {
			return true
		}
	}

	return false
}

func (r *Role) toString() string {
	if len(r.roles) == 0 {
		return ""
	}

	var sb strings.Builder
	sb.WriteString(r.name)
	sb.WriteString(" < ")
	if len(r.roles) != 1 {
		sb.WriteString("(")
	}

	for i, role := range r.roles {
		if i == 0 {
			sb.WriteString(role.name)
		} else {
			sb.WriteString(", ")
			sb.WriteString(role.name)
		}
	}

	if len(r.roles) != 1 {
		sb.WriteString(")")
	}

	return sb.String()
}

func (r *Role) getRoles() []string {
	names := []string{}
	for _, role := range r.roles {
		names = append(names, role.name)
	}
	return names
}
<commit_msg>fix: Remove duplicate domains in generateTempRoles()<commit_after>// Copyright 2017 The casbin Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package defaultrolemanager

import (
	"strings"
	"sync"

	"github.com/casbin/casbin/v2/errors"
	"github.com/casbin/casbin/v2/log"
	"github.com/casbin/casbin/v2/rbac"
)

const defaultDomain string = "casbin::default"

type MatchingFunc func(arg1, arg2 string) bool

// RoleManager provides a default implementation for the RoleManager interface
type RoleManager struct {
	allDomains         *sync.Map
	maxHierarchyLevel  int
	hasPattern         bool
	matchingFunc       MatchingFunc
	hasDomainPattern   bool
	domainMatchingFunc MatchingFunc

	logger log.Logger
}

// NewRoleManager is the constructor for creating an instance of the
// default RoleManager implementation.
func NewRoleManager(maxHierarchyLevel int) rbac.RoleManager {
	rm := RoleManager{}
	rm.allDomains = &sync.Map{}
	rm.allDomains.LoadOrStore(defaultDomain, new(Roles))
	rm.maxHierarchyLevel = maxHierarchyLevel
	rm.hasPattern = false
	rm.hasDomainPattern = false

	rm.SetLogger(&log.DefaultLogger{})

	return &rm
}

// AddMatchingFunc support use pattern in g
func (rm *RoleManager) AddMatchingFunc(name string, fn MatchingFunc) {
	rm.hasPattern = true
	rm.matchingFunc = fn
}

// AddDomainMatchingFunc support use domain pattern in g
func (rm *RoleManager) AddDomainMatchingFunc(name string, fn MatchingFunc) {
	rm.hasDomainPattern = true
	rm.domainMatchingFunc = fn
}

// SetLogger sets role manager's logger.
func (rm *RoleManager) SetLogger(logger log.Logger) {
	rm.logger = logger
}

func (rm *RoleManager) generateTempRoles(domain string) *Roles {
	rm.allDomains.LoadOrStore(domain, new(Roles))

	patternDomain := []string{domain}
	if rm.hasDomainPattern {
		rm.allDomains.Range(func(key, value interface{}) bool {
			if domain != key.(string) && rm.domainMatchingFunc(domain, key.(string)) {
				patternDomain = append(patternDomain, key.(string))
			}
			return true
		})
	}

	allRoles := new(Roles)
	for _, domain := range patternDomain {
		value, _ := rm.allDomains.LoadOrStore(domain, new(Roles))
		value.(*Roles).Range(func(key, value interface{}) bool {
			role2 := value.(*Role)
			role1 := allRoles.createRole(role2.name, rm.matchingFunc)
			for _, v := range role2.roles {
				role3 := allRoles.createRole(v.name, rm.matchingFunc)
				role1.addRole(role3)
			}

			return true
		})
	}

	return allRoles
}

// Clear clears all stored data and resets the role manager to the initial state.
func (rm *RoleManager) Clear() error {
	rm.allDomains = &sync.Map{}
	rm.allDomains.LoadOrStore(defaultDomain, new(Roles))
	return nil
}

// AddLink adds the inheritance link between role: name1 and role: name2.
// aka role: name1 inherits role: name2.
func (rm *RoleManager) AddLink(name1 string, name2 string, domain ...string) error {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return errors.ERR_DOMAIN_PARAMETER
	}

	allRoles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))

	role1, _ := allRoles.(*Roles).LoadOrStore(name1, newRole(name1))
	role2, _ := allRoles.(*Roles).LoadOrStore(name2, newRole(name2))
	role1.(*Role).addRole(role2.(*Role))
	return nil
}

// DeleteLink deletes the inheritance link between role: name1 and role: name2.
// aka role: name1 does not inherit role: name2 any more.
func (rm *RoleManager) DeleteLink(name1 string, name2 string, domain ...string) error {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return errors.ERR_DOMAIN_PARAMETER
	}

	allRoles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))

	_, ok1 := allRoles.(*Roles).Load(name1)
	_, ok2 := allRoles.(*Roles).Load(name1)
	if !ok1 || !ok2 {
		return errors.ERR_NAMES12_NOT_FOUND
	}

	role1, _ := allRoles.(*Roles).LoadOrStore(name1, newRole(name1))
	role2, _ := allRoles.(*Roles).LoadOrStore(name2, newRole(name2))
	role1.(*Role).deleteRole(role2.(*Role))
	return nil
}

// HasLink determines whether role: name1 inherits role: name2.
func (rm *RoleManager) HasLink(name1 string, name2 string, domain ...string) (bool, error) {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return false, errors.ERR_DOMAIN_PARAMETER
	}

	if name1 == name2 {
		return true, nil
	}

	var allRoles *Roles
	if rm.hasDomainPattern || rm.hasPattern {
		allRoles = rm.generateTempRoles(domain[0])
	} else {
		roles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))
		allRoles = roles.(*Roles)
	}

	if !allRoles.hasRole(name1, rm.matchingFunc) || !allRoles.hasRole(name2, rm.matchingFunc) {
		return false, nil
	}

	role1 := allRoles.createRole(name1, rm.matchingFunc)
	return role1.hasRole(name2, rm.maxHierarchyLevel), nil
}

// GetRoles gets the roles that a subject inherits.
func (rm *RoleManager) GetRoles(name string, domain ...string) ([]string, error) {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return nil, errors.ERR_DOMAIN_PARAMETER
	}

	var allRoles *Roles
	if rm.hasDomainPattern || rm.hasPattern {
		allRoles = rm.generateTempRoles(domain[0])
	} else {
		roles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))
		allRoles = roles.(*Roles)
	}

	if !allRoles.hasRole(name, rm.matchingFunc) {
		return []string{}, nil
	}

	roles := allRoles.createRole(name, rm.matchingFunc).getRoles()

	return roles, nil
}

// GetUsers gets the users that inherits a subject.
// domain is an unreferenced parameter here, may be used in other implementations.
func (rm *RoleManager) GetUsers(name string, domain ...string) ([]string, error) {
	if len(domain) == 0 {
		domain = []string{defaultDomain}
	} else if len(domain) > 1 {
		return nil, errors.ERR_DOMAIN_PARAMETER
	}

	var allRoles *Roles
	if rm.hasDomainPattern || rm.hasPattern {
		allRoles = rm.generateTempRoles(domain[0])
	} else {
		roles, _ := rm.allDomains.LoadOrStore(domain[0], new(Roles))
		allRoles = roles.(*Roles)
	}

	if !allRoles.hasRole(name, rm.domainMatchingFunc) {
		return nil, errors.ERR_NAME_NOT_FOUND
	}

	names := []string{}
	allRoles.Range(func(_, value interface{}) bool {
		role := value.(*Role)
		if role.hasDirectRole(name) {
			names = append(names, role.name)
		}
		return true
	})

	return names, nil
}

// PrintRoles prints all the roles to log.
func (rm *RoleManager) PrintRoles() error {
	if !(rm.logger).IsEnabled() {
		return nil
	}

	var roles []string
	rm.allDomains.Range(func(_, value interface{}) bool {
		value.(*Roles).Range(func(_, value interface{}) bool {
			if text := value.(*Role).toString(); text != "" {
				roles = append(roles, text)
			}
			return true
		})

		return true
	})

	rm.logger.LogRole(roles)
	return nil
}

// Roles represents all roles in a domain
type Roles struct {
	sync.Map
}

func (roles *Roles) hasRole(name string, matchingFunc MatchingFunc) bool {
	var ok bool
	if matchingFunc != nil {
		roles.Range(func(key, value interface{}) bool {
			if matchingFunc(name, key.(string)) {
				ok = true
			}
			return true
		})
	} else {
		_, ok = roles.Load(name)
	}

	return ok
}

func (roles *Roles) createRole(name string, matchingFunc MatchingFunc) *Role {
	role, _ := roles.LoadOrStore(name, newRole(name))

	if matchingFunc != nil {
		roles.Range(func(key, value interface{}) bool {
			if matchingFunc(name, key.(string)) && name != key.(string) {
				// Add new role to matching role
				role1, _ := roles.LoadOrStore(key.(string), newRole(key.(string)))
				role.(*Role).addRole(role1.(*Role))
			}
			return true
		})
	}

	return role.(*Role)
}

// Role represents the data structure for a role in RBAC.
type Role struct {
	name  string
	roles []*Role
}

func newRole(name string) *Role {
	r := Role{}
	r.name = name
	return &r
}

func (r *Role) addRole(role *Role) {
	for _, rr := range r.roles {
		if rr.name == role.name {
			return
		}
	}

	r.roles = append(r.roles, role)
}

func (r *Role) deleteRole(role *Role) {
	for i, rr := range r.roles {
		if rr.name == role.name {
			r.roles = append(r.roles[:i], r.roles[i+1:]...)
			return
		}
	}
}

func (r *Role) hasRole(name string, hierarchyLevel int) bool {
	if r.name == name {
		return true
	}

	if hierarchyLevel <= 0 {
		return false
	}

	for _, role := range r.roles {
		if role.hasRole(name, hierarchyLevel-1) {
			return true
		}
	}
	return false
}

func (r *Role) hasDirectRole(name string) bool {
	for _, role := range r.roles {
		if role.name == name {
			return true
		}
	}

	return false
}

func (r *Role) toString() string {
	if len(r.roles) == 0 {
		return ""
	}

	var sb strings.Builder
	sb.WriteString(r.name)
	sb.WriteString(" < ")
	if len(r.roles) != 1 {
		sb.WriteString("(")
	}

	for i, role := range r.roles {
		if i == 0 {
			sb.WriteString(role.name)
		} else {
			sb.WriteString(", ")
			sb.WriteString(role.name)
		}
	}

	if len(r.roles) != 1 {
		sb.WriteString(")")
	}

	return sb.String()
}

func (r *Role) getRoles() []string {
	names := []string{}
	for _, role := range r.roles {
		names = append(names, role.name)
	}
	return names
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"sync"
	"html/template"
	"net/http"
	"net/http/fcgi"
	"log"
	"os/exec"
	"encoding/json"
	"bytes"
	"strings"
)

const (
	PS = string(os.PathSeparator)
	WWWDIR = "/usr/local/sambactl/"
	SRCURL = "https://github.com/Knorkebrot/sambactl"
)

var (
	m sync.Mutex = sync.Mutex{}
	cwd string = ""
	fileHandler http.Handler
)

type TplData struct {
	SrcURL		string
	Errors		[]string
	Users		[]User
}

type User struct {
	Username	string
}

func main() {
	err := os.Chdir(WWWDIR)
	if err != nil {
		log.Fatalln("Chdir:", err)
	}
	err = os.Mkdir(WWWDIR + PS + "htdocs", 0755)
	if err != nil && !os.IsExist(err) {
		log.Fatalln("Mkdir:", err)
	}

	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalln("Hostname:", err)
	}
	log.Println("Running in", WWWDIR, "- http://" + hostname + ":80/")

	fileHandler = http.FileServer(http.Dir(WWWDIR + PS + "htdocs"))
	http.HandleFunc("/", http.HandlerFunc(indexHandler))
	http.HandleFunc("/user/", http.HandlerFunc(handler))

	err = fcgi.Serve(nil, nil)
	/*err = http.ListenAndServe(":80", nil)*/
	if err != nil {
		log.Fatalln("Serve:", err)
	}
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	if strings.HasSuffix(path, "/index.html") {
		http.Redirect(w, r, "/", http.StatusMovedPermanently)
		return
	}
	if path != "/" {
		fileHandler.ServeHTTP(w, r)
		return
	}

	tpl, err := template.ParseFiles("htdocs/index.html")
	if err != nil {
		log.Println("template.ParseFiles:", err)
	}
	data := TplData{}
	data.SrcURL = SRCURL

	errors := make([]string, 0, 10)

	users, err := getUsers()
	if err != nil {
		errors = append(errors, "User konnten nicht geladen werden.")
	}
	data.Users = users

	// dummy
	//users = make([]User, 0, 1)
	//data.Users = append(users, User{"Bildbearbeitung"})

	data.Errors = errors
	tpl.Execute(w, data)
}

func handler(w http.ResponseWriter, r *http.Request) {
	username := r.FormValue("username")
	passwd := r.FormValue("password")
	del := r.FormValue("del")
	found := false

	users, err := getUsers()
	if err != nil {
		http.Error(w, "User loading failed:" + err.Error(), http.StatusInternalServerError)
		return
	}
	for _, u := range users {
		if u.Username == username {
			found = true
			break
		}
	}
	if found {
		if len(del) != 0 {
			err = delUser(username)
		} else {
			err = changePassword(username, passwd)
		}
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	} else {
		err = addUser(username, passwd)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte("\"ok\"\n"))
}

func getUsers() ([]User, error) {
	cmd := exec.Command("/usr/bin/sudo", "/usr/local/bin/sambactl-worker", "list")
	var out bytes.Buffer
	cmd.Stdout = &out
	m.Lock()
	err := cmd.Run()
	m.Unlock()
	if err != nil {
		return nil, err
	}
	var users []User
	err = json.Unmarshal(out.Bytes(), &users)
	if err != nil {
		return nil, err
	}
	return users, nil
}

func changePassword(username, newpasswd string) error {
	cmd := exec.Command("/usr/bin/sudo", "/usr/local/bin/sambactl-worker", "passwd", username, newpasswd)
	var out bytes.Buffer
	cmd.Stdout = &out
	m.Lock()
	err := cmd.Run()
	m.Unlock()
	if err != nil {
		return err
	}
	return nil
}

func addUser(username, passwd string) error {
	cmd := exec.Command("/usr/bin/sudo", "/usr/local/bin/sambactl-worker", "adduser", username, passwd)
	var out bytes.Buffer
	cmd.Stdout = &out
	m.Lock()
	err := cmd.Run()
	m.Unlock()
	if err != nil {
		return err
	}
	return nil
}

func delUser(username string) error {
	cmd := exec.Command("/usr/bin/sudo", "/usr/local/bin/sambactl-worker", "deluser", username)
	var out bytes.Buffer
	cmd.Stdout = &out
	m.Lock()
	err := cmd.Run()
	m.Unlock()
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>remove unused var, somehow the compiler doesnt complain if its in the var block<commit_after>package main

import (
	"os"
	"sync"
	"html/template"
	"net/http"
	"net/http/fcgi"
	"log"
	"os/exec"
	"encoding/json"
	"bytes"
	"strings"
)

const (
	PS = string(os.PathSeparator)
	WWWDIR = "/usr/local/sambactl/"
	SRCURL = "https://github.com/Knorkebrot/sambactl"
)

var (
	m sync.Mutex = sync.Mutex{}
	fileHandler http.Handler
)

type TplData struct {
	SrcURL		string
	Errors		[]string
	Users		[]User
}

type User struct {
	Username	string
}

func main() {
	err := os.Chdir(WWWDIR)
	if err != nil {
		log.Fatalln("Chdir:", err)
	}
	err = os.Mkdir(WWWDIR + PS + "htdocs", 0755)
	if err != nil && !os.IsExist(err) {
		log.Fatalln("Mkdir:", err)
	}

	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalln("Hostname:", err)
	}
	log.Println("Running in", WWWDIR, "- http://" + hostname + ":80/")

	fileHandler = http.FileServer(http.Dir(WWWDIR + PS + "htdocs"))
	http.HandleFunc("/", http.HandlerFunc(indexHandler))
	http.HandleFunc("/user/", http.HandlerFunc(handler))

	err = fcgi.Serve(nil, nil)
	/*err = http.ListenAndServe(":80", nil)*/
	if err != nil {
		log.Fatalln("Serve:", err)
	}
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	if strings.HasSuffix(path, "/index.html") {
		http.Redirect(w, r, "/", http.StatusMovedPermanently)
		return
	}
	if path != "/" {
		fileHandler.ServeHTTP(w, r)
		return
	}

	tpl, err := template.ParseFiles("htdocs/index.html")
	if err != nil {
		log.Println("template.ParseFiles:", err)
	}
	data := TplData{}
	data.SrcURL = SRCURL

	errors := make([]string, 0, 10)

	users, err := getUsers()
	if err != nil {
		errors = append(errors, "User konnten nicht geladen werden.")
	}
	data.Users = users

	// dummy
	//users = make([]User, 0, 1)
	//data.Users = append(users, User{"Bildbearbeitung"})

	data.Errors = errors
	tpl.Execute(w, data)
}

func handler(w http.ResponseWriter, r *http.Request) {
	username := r.FormValue("username")
	passwd := r.FormValue("password")
	del := r.FormValue("del")
	found := false

	users, err := getUsers()
	if err != nil {
		http.Error(w, "User loading failed:" + err.Error(), http.StatusInternalServerError)
		return
	}
	for _, u := range users {
		if u.Username == username {
			found = true
			break
		}
	}
	if found {
		if len(del) != 0 {
			err = delUser(username)
		} else {
			err = changePassword(username, passwd)
		}
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	} else {
		err = addUser(username, passwd)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte("\"ok\"\n"))
}

func getUsers() ([]User, error) {
	cmd := exec.Command("/usr/bin/sudo", "/usr/local/bin/sambactl-worker", "list")
	var out bytes.Buffer
	cmd.Stdout = &out
	m.Lock()
	err := cmd.Run()
	m.Unlock()
	if err != nil {
		return nil, err
	}
	var users []User
	err = json.Unmarshal(out.Bytes(), &users)
	if err != nil {
		return nil, err
	}
	return users, nil
}

func changePassword(username, newpasswd string) error {
	cmd := exec.Command("/usr/bin/sudo", "/usr/local/bin/sambactl-worker", "passwd", username, newpasswd)
	var out bytes.Buffer
	cmd.Stdout = &out
	m.Lock()
	err := cmd.Run()
	m.Unlock()
	if err != nil {
		return err
	}
	return nil
}

func addUser(username, passwd string) error {
	cmd := exec.Command("/usr/bin/sudo", "/usr/local/bin/sambactl-worker", "adduser", username, passwd)
	var out bytes.Buffer
	cmd.Stdout = &out
	m.Lock()
	err := cmd.Run()
	m.Unlock()
	if err != nil {
		return err
	}
	return nil
}

func delUser(username string) error {
	cmd := exec.Command("/usr/bin/sudo", "/usr/local/bin/sambactl-worker", "deluser", username)
	var out bytes.Buffer
	cmd.Stdout = &out
	m.Lock()
	err := cmd.Run()
	m.Unlock()
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/ernado/selectel/storage"
	"github.com/jwaldrip/odin/cli"
	"log"
	"os"
)

const (
	envKey       = storage.EnvKey
	envUser      = storage.EnvUser
	envContainer = "SELECTEL_CONTAINER"
)

var (
	client    = cli.New("0.1", "Selectel storage command line client", connect)
	user, key string
	container string
	api       storage.API
	debug     bool
)

func init() {
	client.DefineBoolFlagVar(&debug, "debug", false, "debug mode")
	client.DefineStringFlag("key", "", fmt.Sprintf("selectel storage key (%s)", envKey))
	client.DefineStringFlag("user", "", fmt.Sprintf("selectel storage user (%s)", envUser))
	client.DefineStringFlag("container", "", fmt.Sprintf("default container (%s)", envContainer))
	infoCommand := client.DefineSubCommand("info", "print information about storage/container/object", info)
	infoCommand.DefineStringFlag("type", "storage", "storage, container or object")
}

func readFlag(c cli.Command, name, env string) string {
	if len(os.Getenv(env)) > 0 {
		return os.Getenv(env)
	}
	return c.Flag(name).String()
}

func blank(s string) bool {
	return len(s) == 0
}

// connect reads credentials and performs auth
func connect(c cli.Command) {
	key = readFlag(c, "key", envKey)
	user = readFlag(c, "user", envUser)
	container = readFlag(c, "container", envContainer)

	// checking for blank credentials
	if blank(key) || blank(user) {
		log.Fatal(storage.ErrorBadCredentials)
	}

	// connencting to api
	var err error
	api, err = storage.New(user, key)
	if err != nil {
		log.Fatal(err)
	}
	api.Debug(debug)
}

// info prints information about storage
func info(c cli.Command) {
	connect(c.Parent())
	var (
		containerName = container
		objectName    string
		data          interface{}
		err           error
		arglen        = len(c.Args())
		command       = c.Flag("type").String()
	)

	defer func() {
		if err != nil {
			log.Fatal(err)
		}
		if blank(containerName) || command == "storage" {
			data = api.Info()
		} else {
			containerApi := api.Container(containerName)
			if blank(objectName) {
				data, err = containerApi.Info()
			} else {
				data, err = containerApi.Object(objectName).Info()
			}
		}
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("%+v\n", data)
	}()

	if arglen > 0 {
		if command == "container" {
			containerName = c.Arg(0).String()
			return
		}
		command = "object"
		if !blank(containerName) && arglen == 1 {
			objectName = c.Arg(0).String()
			return
		}
		if arglen == 2 {
			containerName = c.Arg(0).String()
			objectName = c.Arg(1).String()
			return
		}
	}
	if command == "container" && !blank(containerName) {
		return
	}
	if command == "storage" {
		return
	}
	err = errors.New("Not enough arguments")
}

func main() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("Recovered", r)
		}
	}()
	client.Start()
}
<commit_msg>updated<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/ernado/selectel/storage"
	"github.com/jwaldrip/odin/cli"
	"github.com/olekukonko/tablewriter"
	"io"
	"log"
	"os"
)

const (
	envKey       = storage.EnvKey
	envUser      = storage.EnvUser
	envContainer = "SELECTEL_CONTAINER"
)

var (
	client         = cli.New("0.1", "Selectel storage command line client", connect)
	user, key      string
	container      string
	api            storage.API
	debug          bool
	errorNotEnough = errors.New("Not enought arguments")
)

func init() {
	client.DefineBoolFlagVar(&debug, "debug", false, "debug mode")
	client.DefineStringFlag("key", "", fmt.Sprintf("selectel storage key (%s)", envKey))
	client.DefineStringFlag("user", "", fmt.Sprintf("selectel storage user (%s)", envUser))
	client.DefineStringFlag("container", "", fmt.Sprintf("default container (%s)", envContainer))
	infoCommand := client.DefineSubCommand("info", "print information about storage/container/object", wrap(info))
	infoCommand.DefineStringFlag("type", "storage", "storage, container or object")
	client.DefineSubCommand("upload", "upload object to container", wrap(upload))
	listCommand := client.DefineSubCommand("list", "list objects in container/storage", wrap(list))
	listCommand.DefineStringFlag("type", "storage", "storage or container")
	downloadCommand := client.DefineSubCommand("download", "download object from container", wrap(download))
	downloadCommand.DefineStringFlag("path", "", "destination path")
}

func readFlag(c cli.Command, name, env string) string {
	if len(os.Getenv(env)) > 0 {
		return os.Getenv(env)
	}
	return c.Flag(name).String()
}

func blank(s string) bool {
	return len(s) == 0
}

// connect reads credentials and performs auth
func connect(c cli.Command) {
	key = readFlag(c, "key", envKey)
	user = readFlag(c, "user", envUser)
	container = readFlag(c, "container", envContainer)

	// checking for blank credentials
	if blank(key) || blank(user) {
		log.Fatal(storage.ErrorBadCredentials)
	}

	// connencting to api
	var err error
	api, err = storage.New(user, key)
	if err != nil {
		log.Fatal(err)
	}
	api.Debug(debug)
}

func wrap(callback func(cli.Command)) func(cli.Command) {
	return func(c cli.Command) {
		connect(c.Parent())
		callback(c)
	}
}

// info prints information about storage
func info(c cli.Command) {
	var (
		containerName = container
		objectName    string
		data          interface{}
		err           error
		arglen        = len(c.Args())
		command       = c.Flag("type").String()
	)

	defer func() {
		if err != nil {
			log.Fatal(err)
		}
		if blank(containerName) || command == "storage" {
			data = api.Info()
		} else {
			containerApi := api.Container(containerName)
			if blank(objectName) {
				data, err = containerApi.Info()
			} else {
				data, err = containerApi.Object(objectName).Info()
			}
		}
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("%+v\n", data)
	}()

	if arglen > 0 {
		if command == "container" {
			containerName = c.Arg(0).String()
			return
		}
		command = "object"
		if !blank(containerName) && arglen == 1 {
			objectName = c.Arg(0).String()
			return
		}
		if arglen == 2 {
			containerName = c.Arg(0).String()
			objectName = c.Arg(1).String()
			return
		}
	}
	if command == "container" && !blank(containerName) {
		return
	}
	if command == "storage" {
		return
	}
	err = errorNotEnough
}

func upload(c cli.Command) {
	var path string
	switch len(c.Args()) {
	case 1:
		path = c.Arg(0).String()
	case 2:
		container = c.Arg(0).String()
		path = c.Arg(1).String()
	}
	if blank(container) || blank(path) {
		log.Fatal(errorNotEnough)
	}
	if err := api.Container(container).UploadFile(path); err != nil {
		log.Fatal(err)
	}
	fmt.Printf("uploaded to %s\n", container)
}

func list(c cli.Command) {
	var (
		arglen = len(c.Args())
		table  = tablewriter.NewWriter(os.Stdout)
	)
	if arglen == 0 && (blank(container) || c.Flag("type").String() == "storage") {
		containers, err := api.ContainersInfo()
		if err != nil {
			log.Fatal(err)
		}
		table.SetHeader([]string{"Name", "Objects", "Type"})
		for _, cont := range containers {
			v := []string{cont.Name, fmt.Sprint(cont.ObjectCount), cont.Type}
			table.Append(v)
		}
		table.Render()
		return
	}
	if arglen == 1 {
		container = c.Arg(0).String()
	}
	if blank(container) {
		log.Fatal(errorNotEnough)
	}
	objects, err := api.Container(container).ObjectsInfo()
	if err != nil {
		log.Fatal(err)
	}
	table.SetHeader([]string{"Name", "Size", "Downloaded"})
	for _, object := range objects {
		v := []string{object.Name, fmt.Sprint(object.Size), fmt.Sprint(object.Downloaded)}
		table.Append(v)
	}
	table.Render()
}

func download(c cli.Command) {
	var (
		arglen     = len(c.Args())
		objectName string
		path       = c.Flag("path").String()
	)
	switch arglen {
	case 1:
		objectName = c.Arg(0).String()
	case 2:
		objectName = c.Arg(1).String()
		container = c.Arg(0).String()
	}
	if blank(container) || blank(objectName) {
		log.Fatal(errorNotEnough)
	}
	if blank(path) {
		path = objectName
	}
	reader, err := api.Container(container).Object(objectName).GetReader()
	if err != nil {
		log.Fatal(err)
	}
	defer reader.Close()
	fmt.Printf("downloading %s->%s from %s\n", objectName, path, container)
	f, err := os.Create(path)
	if err != nil {
		log.Fatal(err)
	}
	n, err := io.Copy(f, reader)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("downloaded %s, %d bytes\n", objectName, n)
}

func main() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("Recovered", r)
		}
	}()
	client.Start()
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"regexp"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
)

// return a slice of s3 tags associated with the given s3 bucket. Essentially
// s3.GetBucketTagging, except returns an empty slice instead of an error when
// there are no tags.
func getTagSetS3Bucket(conn *s3.S3, bucket string) ([]*s3.Tag, error) {
	resp, err := conn.GetBucketTagging(&s3.GetBucketTaggingInput{
		Bucket: aws.String(bucket),
	})
	if err != nil {
		if isAWSErr(err, "NoSuchTagSet", "") {
			return nil, nil
		}
		return nil, err
	}

	return resp.TagSet, nil
}

// setTags is a helper to set the tags for a resource. It expects the
// tags field to be named "tags"
func setTagsS3Bucket(conn *s3.S3, d *schema.ResourceData) error {
	if d.HasChange("tags") {
		oraw, nraw := d.GetChange("tags")
		o := oraw.(map[string]interface{})
		n := nraw.(map[string]interface{})

		// Get any existing system tags.
		var sysTagSet []*s3.Tag
		oTagSet, err := getTagSetS3Bucket(conn, d.Get("bucket").(string))
		if err != nil {
			return fmt.Errorf("error getting S3 bucket tags: %s", err)
		}
		for _, tag := range oTagSet {
			if tagIgnoredS3(tag) {
				sysTagSet = append(sysTagSet, tag)
			}
		}

		if len(n)+len(sysTagSet) > 0 {
			// The bucket's tag set must include any system tags that Terraform ignores.
			nTagSet := append(tagsFromMapS3(n), sysTagSet...)

			req := &s3.PutBucketTaggingInput{
				Bucket: aws.String(d.Get("bucket").(string)),
				Tagging: &s3.Tagging{
					TagSet: nTagSet,
				},
			}
			if _, err := RetryOnAwsCodes([]string{"NoSuchBucket", "OperationAborted"}, func() (interface{}, error) {
				return conn.PutBucketTagging(req)
			}); err != nil {
				return fmt.Errorf("error setting S3 bucket tags: %s", err)
			}
		} else if len(o) > 0 && len(sysTagSet) == 0 {
			req := &s3.DeleteBucketTaggingInput{
				Bucket: aws.String(d.Get("bucket").(string)),
			}
			if _, err := RetryOnAwsCodes([]string{"NoSuchBucket", "OperationAborted"}, func() (interface{}, error) {
				return conn.DeleteBucketTagging(req)
			}); err != nil {
				return fmt.Errorf("error deleting S3 bucket tags: %s", err)
			}
		}
	}

	return nil
}

func getTagsS3Object(conn *s3.S3, d *schema.ResourceData) error {
	resp, err := retryOnAwsCode(s3.ErrCodeNoSuchKey, func() (interface{}, error) {
		return conn.GetObjectTagging(&s3.GetObjectTaggingInput{
			Bucket: aws.String(d.Get("bucket").(string)),
			Key:    aws.String(d.Get("key").(string)),
		})
	})
	if err != nil {
		return err
	}

	if err := d.Set("tags", tagsToMapS3(resp.(*s3.GetObjectTaggingOutput).TagSet)); err != nil {
		return err
	}

	return nil
}

func setTagsS3Object(conn *s3.S3, d *schema.ResourceData) error {
	if d.HasChange("tags") {
		oraw, nraw := d.GetChange("tags")
		o := oraw.(map[string]interface{})
		n := nraw.(map[string]interface{})

		// Set tags
		if len(o) > 0 {
			if _, err := conn.DeleteObjectTagging(&s3.DeleteObjectTaggingInput{
				Bucket: aws.String(d.Get("bucket").(string)),
				Key:    aws.String(d.Get("key").(string)),
			}); err != nil {
				return err
			}
		}
		if len(n) > 0 {
			if _, err := conn.PutObjectTagging(&s3.PutObjectTaggingInput{
				Bucket: aws.String(d.Get("bucket").(string)),
				Key:    aws.String(d.Get("key").(string)),
				Tagging: &s3.Tagging{
					TagSet: tagsFromMapS3(n),
				},
			}); err != nil {
				return err
			}
		}
	}

	return nil
}

// tagsFromMap returns the tags for the given map of data.
func tagsFromMapS3(m map[string]interface{}) []*s3.Tag {
	result := make([]*s3.Tag, 0, len(m))
	for k, v := range m {
		t := &s3.Tag{
			Key:   aws.String(k),
			Value: aws.String(v.(string)),
		}
		if !tagIgnoredS3(t) {
			result = append(result, t)
		}
	}

	return result
}

// tagsToMap turns the list of tags into a map.
func tagsToMapS3(ts []*s3.Tag) map[string]string {
	result := make(map[string]string)
	for _, t := range ts {
		if !tagIgnoredS3(t) {
			result[*t.Key] = *t.Value
		}
	}

	return result
}

// compare a tag against a list of strings and checks if it should
// be ignored or not
func tagIgnoredS3(t *s3.Tag) bool {
	filter := []string{"^aws:"}
	for _, v := range filter {
		log.Printf("[DEBUG] Matching %v with %v\n", v, *t.Key)
		r, _ := regexp.MatchString(v, *t.Key)
		if r {
			log.Printf("[DEBUG] Found AWS specific tag %s (val: %s), ignoring.\n", *t.Key, *t.Value)
			return true
		}
	}
	return false
}
<commit_msg>resource/aws_s3_bucket: Retry GetBucketTagging API calls on NoSuchBucket errors due to eventual consistency<commit_after>package aws

import (
	"fmt"
	"log"
	"regexp"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
)

// return a slice of s3 tags associated with the given s3 bucket. Essentially
// s3.GetBucketTagging, except returns an empty slice instead of an error when
// there are no tags.
func getTagSetS3Bucket(conn *s3.S3, bucket string) ([]*s3.Tag, error) {
	input := &s3.GetBucketTaggingInput{
		Bucket: aws.String(bucket),
	}

	// Retry due to S3 eventual consistency
	outputRaw, err := retryOnAwsCode(s3.ErrCodeNoSuchBucket, func() (interface{}, error) {
		return conn.GetBucketTagging(input)
	})

	// S3 API Reference (https://docs.aws.amazon.com/AmazonS3/latest/API/API_GetBucketTagging.html)
	// lists the special error as NoSuchTagSetError, however the existing logic used NoSuchTagSet
	// and the AWS Go SDK has neither as a constant.
	if isAWSErr(err, "NoSuchTagSet", "") {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	var tagSet []*s3.Tag
	if output, ok := outputRaw.(*s3.GetBucketTaggingOutput); ok {
		tagSet = output.TagSet
	}

	return tagSet, nil
}

// setTags is a helper to set the tags for a resource. It expects the
// tags field to be named "tags"
func setTagsS3Bucket(conn *s3.S3, d *schema.ResourceData) error {
	if d.HasChange("tags") {
		oraw, nraw := d.GetChange("tags")
		o := oraw.(map[string]interface{})
		n := nraw.(map[string]interface{})

		// Get any existing system tags.
		var sysTagSet []*s3.Tag
		oTagSet, err := getTagSetS3Bucket(conn, d.Get("bucket").(string))
		if err != nil {
			return fmt.Errorf("error getting S3 bucket tags: %s", err)
		}
		for _, tag := range oTagSet {
			if tagIgnoredS3(tag) {
				sysTagSet = append(sysTagSet, tag)
			}
		}

		if len(n)+len(sysTagSet) > 0 {
			// The bucket's tag set must include any system tags that Terraform ignores.
			nTagSet := append(tagsFromMapS3(n), sysTagSet...)

			req := &s3.PutBucketTaggingInput{
				Bucket: aws.String(d.Get("bucket").(string)),
				Tagging: &s3.Tagging{
					TagSet: nTagSet,
				},
			}
			if _, err := RetryOnAwsCodes([]string{"NoSuchBucket", "OperationAborted"}, func() (interface{}, error) {
				return conn.PutBucketTagging(req)
			}); err != nil {
				return fmt.Errorf("error setting S3 bucket tags: %s", err)
			}
		} else if len(o) > 0 && len(sysTagSet) == 0 {
			req := &s3.DeleteBucketTaggingInput{
				Bucket: aws.String(d.Get("bucket").(string)),
			}
			if _, err := RetryOnAwsCodes([]string{"NoSuchBucket", "OperationAborted"}, func() (interface{}, error) {
				return conn.DeleteBucketTagging(req)
			}); err != nil {
				return fmt.Errorf("error deleting S3 bucket tags: %s", err)
			}
		}
	}

	return nil
}

func getTagsS3Object(conn *s3.S3, d *schema.ResourceData) error {
	resp, err := retryOnAwsCode(s3.ErrCodeNoSuchKey, func() (interface{}, error) {
		return conn.GetObjectTagging(&s3.GetObjectTaggingInput{
			Bucket: aws.String(d.Get("bucket").(string)),
			Key:    aws.String(d.Get("key").(string)),
		})
	})
	if err != nil {
		return err
	}

	if err := d.Set("tags", tagsToMapS3(resp.(*s3.GetObjectTaggingOutput).TagSet)); err != nil {
		return err
	}

	return nil
}

func setTagsS3Object(conn *s3.S3, d *schema.ResourceData) error {
	if d.HasChange("tags") {
		oraw, nraw := d.GetChange("tags")
		o := oraw.(map[string]interface{})
		n := nraw.(map[string]interface{})

		// Set tags
		if len(o) > 0 {
			if _, err := conn.DeleteObjectTagging(&s3.DeleteObjectTaggingInput{
				Bucket: aws.String(d.Get("bucket").(string)),
				Key:    aws.String(d.Get("key").(string)),
			}); err != nil {
				return err
			}
		}
		if len(n) > 0 {
			if _, err := conn.PutObjectTagging(&s3.PutObjectTaggingInput{
				Bucket: aws.String(d.Get("bucket").(string)),
				Key:    aws.String(d.Get("key").(string)),
				Tagging: &s3.Tagging{
					TagSet: tagsFromMapS3(n),
				},
			}); err != nil {
				return err
			}
		}
	}

	return nil
}

// tagsFromMap returns the tags for the given map of data.
func tagsFromMapS3(m map[string]interface{}) []*s3.Tag {
	result := make([]*s3.Tag, 0, len(m))
	for k, v := range m {
		t := &s3.Tag{
			Key:   aws.String(k),
			Value: aws.String(v.(string)),
		}
		if !tagIgnoredS3(t) {
			result = append(result, t)
		}
	}

	return result
}

// tagsToMap turns the list of tags into a map.
func tagsToMapS3(ts []*s3.Tag) map[string]string {
	result := make(map[string]string)
	for _, t := range ts {
		if !tagIgnoredS3(t) {
			result[*t.Key] = *t.Value
		}
	}

	return result
}

// compare a tag against a list of strings and checks if it should
// be ignored or not
func tagIgnoredS3(t *s3.Tag) bool {
	filter := []string{"^aws:"}
	for _, v := range filter {
		log.Printf("[DEBUG] Matching %v with %v\n", v, *t.Key)
		r, _ := regexp.MatchString(v, *t.Key)
		if r {
			log.Printf("[DEBUG] Found AWS specific tag %s (val: %s), ignoring.\n", *t.Key, *t.Value)
			return true
		}
	}
	return false
}
<|endoftext|>