<commit_before>// Middleware for keeping track of users, login states and permissions.
package permissionsql

import (
	"github.com/xyproto/pinterface"
	"net/http"
	"strings"
)

// The structure that keeps track of the permissions for various path prefixes
type Permissions struct {
	state              *UserState
	adminPathPrefixes  []string
	userPathPrefixes   []string
	publicPathPrefixes []string
	rootIsPublic       bool
	denied             http.HandlerFunc
}

const (
	// Version number. Stable API within major version numbers.
	Version = 2.0
)

// Initialize a Permissions struct with all the default settings.
// This will also connect to the database host at port 3306.
func New() (*Permissions, error) {
	state, err := NewUserStateSimple()
	if err != nil {
		return nil, err
	}
	return NewPermissions(state), nil
}

// Initialize a Permissions struct with a database connection string
func NewWithConf(connectionString string) (*Permissions, error) {
	state, err := NewUserState(connectionString, true)
	if err != nil {
		return nil, err
	}
	return NewPermissions(state), nil
}

// Initialize a Permissions struct with a dsn
func NewWithDSN(connectionString string, database_name string) (*Permissions, error) {
	state, err := NewUserStateWithDSN(connectionString, database_name, true)
	if err != nil {
		return nil, err
	}
	return NewPermissions(state), nil
}

// Initialize a Permissions struct with the given UserState and
// a few default paths for admin/user/public path prefixes.
func NewPermissions(state *UserState) *Permissions {
	// default permissions
	return &Permissions{state,
		[]string{"/admin"},         // admin path prefixes
		[]string{"/repo", "/data"}, // user path prefixes
		[]string{"/", "/login", "/register", "/favicon.ico", "/style", "/img", "/js",
			"/favicon.ico", "/robots.txt", "/sitemap_index.xml"}, // public
		true,
		PermissionDenied}
}

// Specify the http.HandlerFunc for when the permissions are denied
func (perm *Permissions) SetDenyFunction(f http.HandlerFunc) {
	perm.denied = f
}

// Get the current http.HandlerFunc for when permissions are denied
func (perm *Permissions) DenyFunction() http.HandlerFunc {
	return perm.denied
}

// Retrieve the UserState struct
func (perm *Permissions) UserState() pinterface.IUserState {
	return perm.state
}

// Set everything to public
func (perm *Permissions) Clear() {
	perm.adminPathPrefixes = []string{}
	perm.userPathPrefixes = []string{}
}

// Add an url path prefix that is a page for the logged in administrators
func (perm *Permissions) AddAdminPath(prefix string) {
	perm.adminPathPrefixes = append(perm.adminPathPrefixes, prefix)
}

// Add an url path prefix that is a page for the logged in users
func (perm *Permissions) AddUserPath(prefix string) {
	perm.userPathPrefixes = append(perm.userPathPrefixes, prefix)
}

// Add an url path prefix that is a public page
func (perm *Permissions) AddPublicPath(prefix string) {
	perm.publicPathPrefixes = append(perm.publicPathPrefixes, prefix)
}

// Set all url path prefixes that are for the logged in administrator pages
func (perm *Permissions) SetAdminPath(pathPrefixes []string) {
	perm.adminPathPrefixes = pathPrefixes
}

// Set all url path prefixes that are for the logged in user pages
func (perm *Permissions) SetUserPath(pathPrefixes []string) {
	perm.userPathPrefixes = pathPrefixes
}

// Set all url path prefixes that are for the public pages
func (perm *Permissions) SetPublicPath(pathPrefixes []string) {
	perm.publicPathPrefixes = pathPrefixes
}

// The default "permission denied" http handler.
func PermissionDenied(w http.ResponseWriter, req *http.Request) {
	http.Error(w, "Permission denied.", http.StatusForbidden)
}

// Check if a given request should be rejected.
func (perm *Permissions) Rejected(w http.ResponseWriter, req *http.Request) bool {
	reject := false
	path := req.URL.Path // the path of the url that the user wish to visit

	// If it's not "/" and set to be public regardless of permissions
	if !(perm.rootIsPublic && path == "/") {

		// Reject if it is an admin page and user does not have admin permissions
		for _, prefix := range perm.adminPathPrefixes {
			if strings.HasPrefix(path, prefix) {
				if !perm.state.AdminRights(req) {
					reject = true
					break
				}
			}
		}

		if !reject {
			// Reject if it's a user page and the user does not have user rights
			for _, prefix := range perm.userPathPrefixes {
				if strings.HasPrefix(path, prefix) {
					if !perm.state.UserRights(req) {
						reject = true
						break
					}
				}
			}
		}

		if !reject {
			// Reject if it's not a public page
			found := false
			for _, prefix := range perm.publicPathPrefixes {
				if strings.HasPrefix(path, prefix) {
					found = true
					break
				}
			}
			if !found {
				reject = true
			}
		}

	}

	return reject
}

// Middleware handler (compatible with Negroni)
func (perm *Permissions) ServeHTTP(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
	// Check if the user has the right admin/user rights
	if perm.Rejected(w, req) {
		// Get and call the Permission Denied function
		perm.DenyFunction()(w, req)
		// Reject the request by not calling the next handler below
		return
	}

	// Call the next middleware handler
	next(w, req)
}
<commit_msg>Package description<commit_after>// Package permissionsql provides a way to keeping track of users, login states and permissions.
package permissionsql

import (
	"github.com/xyproto/pinterface"
	"net/http"
	"strings"
)

// The structure that keeps track of the permissions for various path prefixes
type Permissions struct {
	state              *UserState
	adminPathPrefixes  []string
	userPathPrefixes   []string
	publicPathPrefixes []string
	rootIsPublic       bool
	denied             http.HandlerFunc
}

const (
	// Version number. Stable API within major version numbers.
	Version = 2.0
)

// Initialize a Permissions struct with all the default settings.
// This will also connect to the database host at port 3306.
func New() (*Permissions, error) {
	state, err := NewUserStateSimple()
	if err != nil {
		return nil, err
	}
	return NewPermissions(state), nil
}

// Initialize a Permissions struct with a database connection string
func NewWithConf(connectionString string) (*Permissions, error) {
	state, err := NewUserState(connectionString, true)
	if err != nil {
		return nil, err
	}
	return NewPermissions(state), nil
}

// Initialize a Permissions struct with a dsn
func NewWithDSN(connectionString string, database_name string) (*Permissions, error) {
	state, err := NewUserStateWithDSN(connectionString, database_name, true)
	if err != nil {
		return nil, err
	}
	return NewPermissions(state), nil
}

// Initialize a Permissions struct with the given UserState and
// a few default paths for admin/user/public path prefixes.
func NewPermissions(state *UserState) *Permissions {
	// default permissions
	return &Permissions{state,
		[]string{"/admin"},         // admin path prefixes
		[]string{"/repo", "/data"}, // user path prefixes
		[]string{"/", "/login", "/register", "/favicon.ico", "/style", "/img", "/js",
			"/favicon.ico", "/robots.txt", "/sitemap_index.xml"}, // public
		true,
		PermissionDenied}
}

// Specify the http.HandlerFunc for when the permissions are denied
func (perm *Permissions) SetDenyFunction(f http.HandlerFunc) {
	perm.denied = f
}

// Get the current http.HandlerFunc for when permissions are denied
func (perm *Permissions) DenyFunction() http.HandlerFunc {
	return perm.denied
}

// Retrieve the UserState struct
func (perm *Permissions) UserState() pinterface.IUserState {
	return perm.state
}

// Set everything to public
func (perm *Permissions) Clear() {
	perm.adminPathPrefixes = []string{}
	perm.userPathPrefixes = []string{}
}

// Add an url path prefix that is a page for the logged in administrators
func (perm *Permissions) AddAdminPath(prefix string) {
	perm.adminPathPrefixes = append(perm.adminPathPrefixes, prefix)
}

// Add an url path prefix that is a page for the logged in users
func (perm *Permissions) AddUserPath(prefix string) {
	perm.userPathPrefixes = append(perm.userPathPrefixes, prefix)
}

// Add an url path prefix that is a public page
func (perm *Permissions) AddPublicPath(prefix string) {
	perm.publicPathPrefixes = append(perm.publicPathPrefixes, prefix)
}

// Set all url path prefixes that are for the logged in administrator pages
func (perm *Permissions) SetAdminPath(pathPrefixes []string) {
	perm.adminPathPrefixes = pathPrefixes
}

// Set all url path prefixes that are for the logged in user pages
func (perm *Permissions) SetUserPath(pathPrefixes []string) {
	perm.userPathPrefixes = pathPrefixes
}

// Set all url path prefixes that are for the public pages
func (perm *Permissions) SetPublicPath(pathPrefixes []string) {
	perm.publicPathPrefixes = pathPrefixes
}

// The default "permission denied" http handler.
func PermissionDenied(w http.ResponseWriter, req *http.Request) {
	http.Error(w, "Permission denied.", http.StatusForbidden)
}

// Check if a given request should be rejected.
func (perm *Permissions) Rejected(w http.ResponseWriter, req *http.Request) bool {
	reject := false
	path := req.URL.Path // the path of the url that the user wish to visit

	// If it's not "/" and set to be public regardless of permissions
	if !(perm.rootIsPublic && path == "/") {

		// Reject if it is an admin page and user does not have admin permissions
		for _, prefix := range perm.adminPathPrefixes {
			if strings.HasPrefix(path, prefix) {
				if !perm.state.AdminRights(req) {
					reject = true
					break
				}
			}
		}

		if !reject {
			// Reject if it's a user page and the user does not have user rights
			for _, prefix := range perm.userPathPrefixes {
				if strings.HasPrefix(path, prefix) {
					if !perm.state.UserRights(req) {
						reject = true
						break
					}
				}
			}
		}

		if !reject {
			// Reject if it's not a public page
			found := false
			for _, prefix := range perm.publicPathPrefixes {
				if strings.HasPrefix(path, prefix) {
					found = true
					break
				}
			}
			if !found {
				reject = true
			}
		}

	}

	return reject
}

// Middleware handler (compatible with Negroni)
func (perm *Permissions) ServeHTTP(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
	// Check if the user has the right admin/user rights
	if perm.Rejected(w, req) {
		// Get and call the Permission Denied function
		perm.DenyFunction()(w, req)
		// Reject the request by not calling the next handler below
		return
	}

	// Call the next middleware handler
	next(w, req)
}
<|endoftext|>


<commit_before>package winrm

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"time"

	"github.com/masterzen/winrm/soap"
	. "gopkg.in/check.v1"
)

func (s *WinRMSuite) TestExecuteCommand(c *C) {
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}
	count := 0
	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		switch count {
		case 0:
			{
				c.Assert(message.String(), Contains, "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Command")
				count = 1
				return executeCommandResponse, nil
			}
		case 1:
			{
				c.Assert(message.String(), Contains, "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive")
				count = 2
				return outputResponse, nil
			}
		default:
			{
				return doneCommandResponse, nil
			}
		}
	}

	command, _ := shell.Execute("ipconfig /all")
	var stdout, stderr bytes.Buffer
	var wg sync.WaitGroup
	f := func(b *bytes.Buffer, r *commandReader) {
		wg.Add(1)
		defer wg.Done()
		io.Copy(b, r)
	}
	go f(&stdout, command.Stdout)
	go f(&stderr, command.Stderr)
	command.Wait()
	wg.Wait()
	c.Assert(stdout.String(), Equals, "That's all folks!!!")
	c.Assert(stderr.String(), Equals, "This is stderr, I'm pretty sure!")
}

func (s *WinRMSuite) TestStdinCommand(c *C) {
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{
		client: client,
		ID:     "67A74734-DD32-4F10-89DE-49A060483810",
	}

	count := 0

	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		if strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Send") {
			c.Assert(message.String(), Contains, "c3RhbmRhcmQgaW5wdXQ=")
			return "", nil
		} else {
			if strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Command") {
				return executeCommandResponse, nil
			} else if count != 1 && strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive") {
				count = 1
				return outputResponse, nil
			} else {
				return doneCommandResponse, nil
			}
		}
	}

	command, _ := shell.Execute("ipconfig /all")
	command.Stdin.Write([]byte("standard input"))
	// slurp output from command
	var outWriter, errWriter bytes.Buffer
	go io.Copy(&outWriter, command.Stdout)
	go io.Copy(&errWriter, command.Stderr)
	command.Wait()
}

func (s *WinRMSuite) TestCommandExitCode(c *C) {
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{
		client: client,
		ID:     "67A74734-DD32-4F10-89DE-49A060483810",
	}

	count := 0

	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		defer func() { count++ }()
		switch count {
		case 0:
			return executeCommandResponse, nil
		case 1:
			return doneCommandResponse, nil
		default:
			c.Log("Mimicking some observed Windows behavior where only the first 'done' response has the actual exit code and 0 afterwards")
			return doneCommandExitCode0Response, nil
		}
	}

	command, _ := shell.Execute("ipconfig /all")

	command.Wait()
	<-time.After(time.Second) // to make the test fail if fetchOutput races to re-set the exit code

	c.Assert(command.ExitCode(), Equals, 123)
}

func (s *WinRMSuite) TestCloseCommandStopsFetch(c *C) {
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}

	http := make(chan string)
	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		switch {
		case strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive"):
			c.Log("Request for command output received by server")
			r := <-http
			c.Log("Returning command output")
			return r, nil
		case strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Command"):
			return executeCommandResponse, nil
		case strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Signal"):
			c.Log("Signal message received by server")
			return "", nil // response is not used
		default:
			c.Logf("Unexpected message: %s", message)
			return "", nil
		}
	}

	command, _ := shell.Execute("ipconfig /all")
	// need to be reading Stdout/Stderr, otherwise, the writes to these are blocking...
	go ioutil.ReadAll(command.Stdout)
	go ioutil.ReadAll(command.Stderr)

	http <- outputResponse // wait for command to enter fetch/slurp

	command.Close()

	select {
	case http <- outputResponse: // return to fetch from slurp
		c.Log("Fetch loop 'drained' one last reponse before realizing that the command is now closed")
	case <-time.After(1 * time.Second):
		c.Log("no poll within one second, fetch may have stopped")
	}

	select {
	case http <- outputResponse:
		c.Log("Fetch loop is still polling after command.Close()")
		c.FailNow()
	case <-time.After(1 * time.Second):
		c.Log("no poll within one second, assuming fetch has stopped")
	}
}

func (s *WinRMSuite) TestConnectionTimeout(c *C) {
	count := 0
	ts, host, port, err := StartTestServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(2 * time.Second)
		w.Header().Set("Content-Type", "application/soap+xml")
		switch count {
		case 0:
			{
				count = 1
				fmt.Fprintln(w, executeCommandResponse)
			}
		case 1:
			{
				count = 2
				fmt.Fprintln(w, outputResponse)
			}
		default:
			{
				fmt.Fprintln(w, doneCommandResponse)
			}
		}
	}))
	defer ts.Close()

	if err != nil {
		c.Error(err)
	}

	client, err := NewClient(&Endpoint{Host: host, Port: port, Timeout: 1 * time.Second}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}
	_, err = shell.Execute("ipconfig /all")
	c.Assert(err, ErrorMatches, ".*timeout.*")
}

func (s *WinRMSuite) TestOperationTimeoutSupport(c *C) {
	count := 0
	ts, host, port, err := StartTestServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/soap+xml")
		switch count {
		case 0:
			{
				count = 1
				fmt.Fprintln(w, executeCommandResponse)
			}
		case 1:
			{
				count = 2
				w.WriteHeader(500)
				fmt.Fprintln(w, operationTimeoutResponse)
			}
		case 2:
			{
				count = 3
				fmt.Fprintln(w, outputResponse)
			}
		default:
			{
				fmt.Fprintln(w, doneCommandResponse)
			}
		}
	}))
	defer ts.Close()

	if err != nil {
		c.Error(err)
	}

	client, err := NewClient(&Endpoint{Host: host, Port: port}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}
	command, _ := shell.Execute("ipconfig /all")
	var stdout, stderr bytes.Buffer
	var wg sync.WaitGroup
	f := func(b *bytes.Buffer, r *commandReader) {
		wg.Add(1)
		defer wg.Done()
		io.Copy(b, r)
	}
	go f(&stdout, command.Stdout)
	go f(&stderr, command.Stderr)
	command.Wait()
	wg.Wait()
	c.Assert(stdout.String(), Equals, "That's all folks!!!")
	c.Assert(stderr.String(), Equals, "This is stderr, I'm pretty sure!")
}

func (s *WinRMSuite) TestEOFError(c *C) {
	count := 0
	var ts *httptest.Server
	ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/soap+xml")
		switch count {
		case 0:
			{
				count = 1
				fmt.Fprintln(w, executeCommandResponse)
			}
		case 1:
			{
				ts.CloseClientConnections()
			}
		default:
			{
				fmt.Fprintln(w, doneCommandResponse)
			}
		}
	}))

	host, port, err := FindHostAndPortFromURL(ts.URL)
	if err != nil {
		c.Error(err)
	}

	client, err := NewClient(&Endpoint{Host: host, Port: port}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}
	command, _ := shell.Execute("ipconfig /all")
	ts.CloseClientConnections()
	command.Wait()
	c.Assert(command.exitCode, Equals, 16001)
	c.Assert(command.err.Error(), Contains, "EOF")
}
<commit_msg>Fix failing EOF test<commit_after>package winrm

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/masterzen/winrm/soap"
	. "gopkg.in/check.v1"
)

func (s *WinRMSuite) TestExecuteCommand(c *C) {
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}
	count := 0
	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		switch count {
		case 0:
			{
				c.Assert(message.String(), Contains, "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Command")
				count = 1
				return executeCommandResponse, nil
			}
		case 1:
			{
				c.Assert(message.String(), Contains, "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive")
				count = 2
				return outputResponse, nil
			}
		default:
			{
				return doneCommandResponse, nil
			}
		}
	}

	command, _ := shell.Execute("ipconfig /all")
	var stdout, stderr bytes.Buffer
	var wg sync.WaitGroup
	f := func(b *bytes.Buffer, r *commandReader) {
		wg.Add(1)
		defer wg.Done()
		io.Copy(b, r)
	}
	go f(&stdout, command.Stdout)
	go f(&stderr, command.Stderr)
	command.Wait()
	wg.Wait()
	c.Assert(stdout.String(), Equals, "That's all folks!!!")
	c.Assert(stderr.String(), Equals, "This is stderr, I'm pretty sure!")
}

func (s *WinRMSuite) TestStdinCommand(c *C) {
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{
		client: client,
		ID:     "67A74734-DD32-4F10-89DE-49A060483810",
	}

	count := 0

	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		if strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Send") {
			c.Assert(message.String(), Contains, "c3RhbmRhcmQgaW5wdXQ=")
			return "", nil
		} else {
			if strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Command") {
				return executeCommandResponse, nil
			} else if count != 1 && strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive") {
				count = 1
				return outputResponse, nil
			} else {
				return doneCommandResponse, nil
			}
		}
	}

	command, _ := shell.Execute("ipconfig /all")
	command.Stdin.Write([]byte("standard input"))
	// slurp output from command
	var outWriter, errWriter bytes.Buffer
	go io.Copy(&outWriter, command.Stdout)
	go io.Copy(&errWriter, command.Stderr)
	command.Wait()
}

func (s *WinRMSuite) TestCommandExitCode(c *C) {
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{
		client: client,
		ID:     "67A74734-DD32-4F10-89DE-49A060483810",
	}

	count := 0

	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		defer func() { count++ }()
		switch count {
		case 0:
			return executeCommandResponse, nil
		case 1:
			return doneCommandResponse, nil
		default:
			c.Log("Mimicking some observed Windows behavior where only the first 'done' response has the actual exit code and 0 afterwards")
			return doneCommandExitCode0Response, nil
		}
	}

	command, _ := shell.Execute("ipconfig /all")

	command.Wait()
	<-time.After(time.Second) // to make the test fail if fetchOutput races to re-set the exit code

	c.Assert(command.ExitCode(), Equals, 123)
}

func (s *WinRMSuite) TestCloseCommandStopsFetch(c *C) {
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}

	http := make(chan string)
	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		switch {
		case strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Receive"):
			c.Log("Request for command output received by server")
			r := <-http
			c.Log("Returning command output")
			return r, nil
		case strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Command"):
			return executeCommandResponse, nil
		case strings.Contains(message.String(), "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Signal"):
			c.Log("Signal message received by server")
			return "", nil // response is not used
		default:
			c.Logf("Unexpected message: %s", message)
			return "", nil
		}
	}

	command, _ := shell.Execute("ipconfig /all")
	// need to be reading Stdout/Stderr, otherwise, the writes to these are blocking...
	go ioutil.ReadAll(command.Stdout)
	go ioutil.ReadAll(command.Stderr)

	http <- outputResponse // wait for command to enter fetch/slurp

	command.Close()

	select {
	case http <- outputResponse: // return to fetch from slurp
		c.Log("Fetch loop 'drained' one last reponse before realizing that the command is now closed")
	case <-time.After(1 * time.Second):
		c.Log("no poll within one second, fetch may have stopped")
	}

	select {
	case http <- outputResponse:
		c.Log("Fetch loop is still polling after command.Close()")
		c.FailNow()
	case <-time.After(1 * time.Second):
		c.Log("no poll within one second, assuming fetch has stopped")
	}
}

func (s *WinRMSuite) TestConnectionTimeout(c *C) {
	count := 0
	ts, host, port, err := StartTestServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(2 * time.Second)
		w.Header().Set("Content-Type", "application/soap+xml")
		switch count {
		case 0:
			{
				count = 1
				fmt.Fprintln(w, executeCommandResponse)
			}
		case 1:
			{
				count = 2
				fmt.Fprintln(w, outputResponse)
			}
		default:
			{
				fmt.Fprintln(w, doneCommandResponse)
			}
		}
	}))
	defer ts.Close()

	if err != nil {
		c.Error(err)
	}

	client, err := NewClient(&Endpoint{Host: host, Port: port, Timeout: 1 * time.Second}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}
	_, err = shell.Execute("ipconfig /all")
	c.Assert(err, ErrorMatches, ".*timeout.*")
}

func (s *WinRMSuite) TestOperationTimeoutSupport(c *C) {
	count := 0
	ts, host, port, err := StartTestServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/soap+xml")
		switch count {
		case 0:
			{
				count = 1
				fmt.Fprintln(w, executeCommandResponse)
			}
		case 1:
			{
				count = 2
				w.WriteHeader(500)
				fmt.Fprintln(w, operationTimeoutResponse)
			}
		case 2:
			{
				count = 3
				fmt.Fprintln(w, outputResponse)
			}
		default:
			{
				fmt.Fprintln(w, doneCommandResponse)
			}
		}
	}))
	defer ts.Close()

	if err != nil {
		c.Error(err)
	}

	client, err := NewClient(&Endpoint{Host: host, Port: port}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}
	command, _ := shell.Execute("ipconfig /all")
	var stdout, stderr bytes.Buffer
	var wg sync.WaitGroup
	f := func(b *bytes.Buffer, r *commandReader) {
		wg.Add(1)
		defer wg.Done()
		io.Copy(b, r)
	}
	go f(&stdout, command.Stdout)
	go f(&stderr, command.Stderr)
	command.Wait()
	wg.Wait()
	c.Assert(stdout.String(), Equals, "That's all folks!!!")
	c.Assert(stderr.String(), Equals, "This is stderr, I'm pretty sure!")
}

func (s *WinRMSuite) TestEOFError(c *C) {
	count := 0
	client, err := NewClient(&Endpoint{Host: "localhost", Port: 5985}, "Administrator", "v3r1S3cre7")
	c.Assert(err, IsNil)

	// simulating a dropped client connection
	client.http = func(client *Client, message *soap.SoapMessage) (string, error) {
		defer func() { count++ }()
		switch count {
		case 0:
			return executeCommandResponse, nil
		case 1:
			return "", fmt.Errorf("http response error: 200 - /wsman EOF")
		default:
			return doneCommandExitCode0Response, nil
		}
	}

	shell := &Shell{client: client, ID: "67A74734-DD32-4F10-89DE-49A060483810"}
	command, _ := shell.Execute("ipconfig /all")

	command.Wait()
	c.Assert(command.exitCode, Equals, 16001)
	c.Assert(command.err.Error(), Contains, "EOF")
}
<|endoftext|>


<commit_before>package ipset

import (
	"log"
	"os/exec"
	"strings"

	"github.com/pkg/errors"
	"k8s.io/apimachinery/pkg/types"
)

type Name string

type Type string

const (
	ListSet = Type("list:set")
	HashIP  = Type("hash:ip")
	HashNet = Type("hash:net")
)

type Interface interface {
	Create(ipsetName Name, ipsetType Type) error
	AddEntry(user types.UID, ipsetName Name, entry string, comment string) error
	DelEntry(user types.UID, ipsetName Name, entry string) error
	Exist(user types.UID, ipsetName Name, entry string) bool
	Flush(ipsetName Name) error
	Destroy(ipsetName Name) error

	List(prefix string) ([]Name, error)

	FlushAll() error
	DestroyAll() error
}

type entryKey struct {
	ipsetName Name
	entry     string
}

type ipset struct {
	*log.Logger
	enableComments bool
	// List of users per ipset entry. User is either a namespace or a pod.
	// There might be multiple users for the same ipset & entry pair because
	// events from k8s API server might be out of order causing duplicate IPs:
	// https://github.com/weaveworks/weave/issues/2792.
	users map[entryKey]map[types.UID]struct{}
}

func New(logger *log.Logger) Interface {
	ips := &ipset{
		Logger:         logger,
		enableComments: true,
		users:          make(map[entryKey]map[types.UID]struct{}),
	}

	// Check for comment support
	testIpsetName := Name("weave-test-comment")
	// Clear it out if it already exists
	_ = ips.Destroy(testIpsetName)
	// Test for comment support
	if err := ips.Create(testIpsetName, HashIP); err != nil {
		ips.Logger.Printf("failed to create %s; disabling comment support", testIpsetName)
		ips.enableComments = false
	}
	// If it was created, destroy it
	_ = ips.Destroy(testIpsetName)

	return ips
}

func (i *ipset) Create(ipsetName Name, ipsetType Type) error {
	args := []string{"create", string(ipsetName), string(ipsetType)}
	if i.enableComments {
		args = append(args, "comment")
	}
	return doExec(args...)
}

func (i *ipset) AddEntry(user types.UID, ipsetName Name, entry string, comment string) error {
	i.Logger.Printf("adding entry %s to %s of %s", entry, ipsetName, user)

	if !i.addUser(user, ipsetName, entry) { // already in the set
		return nil
	}

	i.Logger.Printf("added entry %s to %s of %s", entry, ipsetName, user)

	args := []string{"add", string(ipsetName), entry}
	if i.enableComments {
		args = append(args, "comment", comment)
	}
	return doExec(args...)
}

func (i *ipset) DelEntry(user types.UID, ipsetName Name, entry string) error {
	i.Logger.Printf("deleting entry %s from %s of %s", entry, ipsetName, user)

	if !i.delUser(user, ipsetName, entry) { // still needed
		return nil
	}

	i.Logger.Printf("deleted entry %s from %s of %s", entry, ipsetName, user)

	return doExec("del", string(ipsetName), entry)
}

func (i *ipset) Exist(user types.UID, ipsetName Name, entry string) bool {
	return i.existUser(user, ipsetName, entry)
}

func (i *ipset) Flush(ipsetName Name) error {
	i.removeSetFromUsers(ipsetName)
	return doExec("flush", string(ipsetName))
}

func (i *ipset) FlushAll() error {
	i.users = make(map[entryKey]map[types.UID]struct{})
	return doExec("flush")
}

func (i *ipset) Destroy(ipsetName Name) error {
	i.removeSetFromUsers(ipsetName)
	return doExec("destroy", string(ipsetName))
}

func (i *ipset) DestroyAll() error {
	i.users = make(map[entryKey]map[types.UID]struct{})
	return doExec("destroy")
}

// Fetch a list of all existing sets with a given prefix
func (i *ipset) List(prefix string) ([]Name, error) {
	output, err := exec.Command("ipset", "list", "-name", "-output", "plain").Output()
	if err != nil {
		return nil, err
	}

	var selected []Name
	sets := strings.Split(string(output), "\n")
	for _, v := range sets {
		if strings.HasPrefix(v, prefix) {
			selected = append(selected, Name(v))
		}
	}

	return selected, err
}

// Returns true if entry does not exist in ipset (entry has to be inserted into ipset).
func (i *ipset) addUser(user types.UID, ipsetName Name, entry string) bool {
	k := entryKey{ipsetName, entry}
	add := false

	if i.users[k] == nil {
		i.users[k] = make(map[types.UID]struct{})
	}
	if len(i.users[k]) == 0 {
		add = true
	}
	i.users[k][user] = struct{}{}

	return add
}

// Returns true if user is the last owner of entry (entry has to be removed from ipset).
func (i *ipset) delUser(user types.UID, ipsetName Name, entry string) bool {
	k := entryKey{ipsetName, entry}

	oneLeft := len(i.users[k]) == 1
	delete(i.users[k], user)

	if len(i.users[k]) == 0 {
		delete(i.users, k)
	}

	return oneLeft && (len(i.users[k]) == 0)
}

func (i *ipset) existUser(user types.UID, ipsetName Name, entry string) bool {
	_, ok := i.users[entryKey{ipsetName, entry}][user]
	return ok
}

func (i *ipset) removeSetFromUsers(ipsetName Name) {
	for k := range i.users {
		if k.ipsetName == ipsetName {
			delete(i.users, k)
		}
	}
}

func doExec(args ...string) error {
	if output, err := exec.Command("ipset", args...).CombinedOutput(); err != nil {
		return errors.Wrapf(err, "ipset %v failed: %s", args, output)
	}
	return nil
}
<commit_msg>Add random nonce to test ipset name to prevent race<commit_after>package ipset

import (
	"crypto/rand"
	"encoding/hex"
	"log"
	"os/exec"
	"strings"

	"github.com/pkg/errors"
	"k8s.io/apimachinery/pkg/types"
)

type Name string

type Type string

const (
	ListSet = Type("list:set")
	HashIP  = Type("hash:ip")
	HashNet = Type("hash:net")
)

type Interface interface {
	Create(ipsetName Name, ipsetType Type) error
	AddEntry(user types.UID, ipsetName Name, entry string, comment string) error
	DelEntry(user types.UID, ipsetName Name, entry string) error
	Exist(user types.UID, ipsetName Name, entry string) bool
	Flush(ipsetName Name) error
	Destroy(ipsetName Name) error

	List(prefix string) ([]Name, error)

	FlushAll() error
	DestroyAll() error
}

type entryKey struct {
	ipsetName Name
	entry     string
}

type ipset struct {
	*log.Logger
	enableComments bool
	// List of users per ipset entry. User is either a namespace or a pod.
	// There might be multiple users for the same ipset & entry pair because
	// events from k8s API server might be out of order causing duplicate IPs:
	// https://github.com/weaveworks/weave/issues/2792.
	users map[entryKey]map[types.UID]struct{}
}

func New(logger *log.Logger) Interface {
	ips := &ipset{
		Logger:         logger,
		enableComments: true,
		users:          make(map[entryKey]map[types.UID]struct{}),
	}

	// Check for comment support

	// To prevent from a race when more than one process check for the support
	// we append a random nonce to the test ipset name. The final name is
	// shorter than 31 chars (max ipset name).
	nonce := make([]byte, 4)
	rand.Read(nonce)
	testIpsetName := Name("weave-test-comment" + hex.EncodeToString(nonce))

	// Clear it out if it already exists
	_ = ips.Destroy(testIpsetName)
	// Test for comment support
	if err := ips.Create(testIpsetName, HashIP); err != nil {
		ips.Logger.Printf("failed to create %s; disabling comment support", testIpsetName)
		ips.enableComments = false
	}
	// If it was created, destroy it
	_ = ips.Destroy(testIpsetName)

	return ips
}

func (i *ipset) Create(ipsetName Name, ipsetType Type) error {
	args := []string{"create", string(ipsetName), string(ipsetType)}
	if i.enableComments {
		args = append(args, "comment")
	}
	return doExec(args...)
}

func (i *ipset) AddEntry(user types.UID, ipsetName Name, entry string, comment string) error {
	i.Logger.Printf("adding entry %s to %s of %s", entry, ipsetName, user)

	if !i.addUser(user, ipsetName, entry) { // already in the set
		return nil
	}

	i.Logger.Printf("added entry %s to %s of %s", entry, ipsetName, user)

	args := []string{"add", string(ipsetName), entry}
	if i.enableComments {
		args = append(args, "comment", comment)
	}
	return doExec(args...)
}

func (i *ipset) DelEntry(user types.UID, ipsetName Name, entry string) error {
	i.Logger.Printf("deleting entry %s from %s of %s", entry, ipsetName, user)

	if !i.delUser(user, ipsetName, entry) { // still needed
		return nil
	}

	i.Logger.Printf("deleted entry %s from %s of %s", entry, ipsetName, user)

	return doExec("del", string(ipsetName), entry)
}

func (i *ipset) Exist(user types.UID, ipsetName Name, entry string) bool {
	return i.existUser(user, ipsetName, entry)
}

func (i *ipset) Flush(ipsetName Name) error {
	i.removeSetFromUsers(ipsetName)
	return doExec("flush", string(ipsetName))
}

func (i *ipset) FlushAll() error {
	i.users = make(map[entryKey]map[types.UID]struct{})
	return doExec("flush")
}

func (i *ipset) Destroy(ipsetName Name) error {
	i.removeSetFromUsers(ipsetName)
	return doExec("destroy", string(ipsetName))
}

func (i *ipset) DestroyAll() error {
	i.users = make(map[entryKey]map[types.UID]struct{})
	return doExec("destroy")
}

// Fetch a list of all existing sets with a given prefix
func (i *ipset) List(prefix string) ([]Name, error) {
	output, err := exec.Command("ipset", "list", "-name", "-output", "plain").Output()
	if err != nil {
		return nil, err
	}

	var selected []Name
	sets := strings.Split(string(output), "\n")
	for _, v := range sets {
		if strings.HasPrefix(v, prefix) {
			selected = append(selected, Name(v))
		}
	}

	return selected, err
}

// Returns true if entry does not exist in ipset (entry has to be inserted into ipset).
func (i *ipset) addUser(user types.UID, ipsetName Name, entry string) bool {
	k := entryKey{ipsetName, entry}
	add := false

	if i.users[k] == nil {
		i.users[k] = make(map[types.UID]struct{})
	}
	if len(i.users[k]) == 0 {
		add = true
	}
	i.users[k][user] = struct{}{}

	return add
}

// Returns true if user is the last owner of entry (entry has to be removed from ipset).
func (i *ipset) delUser(user types.UID, ipsetName Name, entry string) bool {
	k := entryKey{ipsetName, entry}

	oneLeft := len(i.users[k]) == 1
	delete(i.users[k], user)

	if len(i.users[k]) == 0 {
		delete(i.users, k)
	}

	return oneLeft && (len(i.users[k]) == 0)
}

func (i *ipset) existUser(user types.UID, ipsetName Name, entry string) bool {
	_, ok := i.users[entryKey{ipsetName, entry}][user]
	return ok
}

func (i *ipset) removeSetFromUsers(ipsetName Name) {
	for k := range i.users {
		if k.ipsetName == ipsetName {
			delete(i.users, k)
		}
	}
}

func doExec(args ...string) error {
	if output, err := exec.Command("ipset", args...).CombinedOutput(); err != nil {
		return errors.Wrapf(err, "ipset %v failed: %s", args, output)
	}
	return nil
}
<|endoftext|>


<commit_before>package opts

import (
	"fmt"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/docker/docker/pkg/parsers"
)

// ListOpts type
type ListOpts struct {
	values    []string
	validator ValidatorFctType
}

func NewListOpts(validator ValidatorFctType) ListOpts {
	return ListOpts{
		validator: validator,
	}
}

func (opts *ListOpts) String() string {
	return fmt.Sprintf("%v", []string(opts.values))
}

// Set validates if needed the input value and add it to the
// internal slice.
func (opts *ListOpts) Set(value string) error {
	if opts.validator != nil {
		v, err := opts.validator(value)
		if err != nil {
			return err
		}
		value = v
	}
	opts.values = append(opts.values, value)
	return nil
}

// Delete remove the given element from the slice.
func (opts *ListOpts) Delete(key string) {
	for i, k := range opts.values {
		if k == key {
			opts.values = append(opts.values[:i], opts.values[i+1:]...)
			return
		}
	}
}

// GetMap returns the content of values in a map in order to avoid
// duplicates.
// FIXME: can we remove this?
func (opts *ListOpts) GetMap() map[string]struct{} {
	ret := make(map[string]struct{})
	for _, k := range opts.values {
		ret[k] = struct{}{}
	}
	return ret
}

// GetAll returns the values' slice.
// FIXME: Can we remove this?
func (opts *ListOpts) GetAll() []string {
	return opts.values
}

// Get checks the existence of the given key.
func (opts *ListOpts) Get(key string) bool {
	for _, k := range opts.values {
		if k == key {
			return true
		}
	}
	return false
}

// Len returns the amount of element in the slice.
func (opts *ListOpts) Len() int {
	return len(opts.values)
}

// Validators
type ValidatorFctType func(val string) (string, error)

func ValidateAttach(val string) (string, error) {
	if val != "stdin" && val != "stdout" && val != "stderr" {
		return val, fmt.Errorf("Unsupported stream name: %s", val)
	}
	return val, nil
}

func ValidateLink(val string) (string, error) {
	if _, err := parsers.PartParser("name:alias", val); err != nil {
		return val, err
	}
	return val, nil
}

func ValidatePath(val string) (string, error) {
	var containerPath string

	if strings.Count(val, ":") > 2 {
		return val, fmt.Errorf("bad format for volumes: %s", val)
	}

	splited := strings.SplitN(val, ":", 2)
	if len(splited) == 1 {
		containerPath = splited[0]
		val = filepath.Clean(splited[0])
	} else {
		containerPath = splited[1]
		val = fmt.Sprintf("%s:%s", splited[0], filepath.Clean(splited[1]))
	}

	if !filepath.IsAbs(containerPath) {
		return val, fmt.Errorf("%s is not an absolute path", containerPath)
	}
	return val, nil
}

func ValidateEnv(val string) (string, error) {
	arr := strings.Split(val, "=")
	if len(arr) > 1 {
		return val, nil
	}
	return fmt.Sprintf("%s=%s", val, os.Getenv(val)), nil
}

func ValidateIPAddress(val string) (string, error) {
	var ip = net.ParseIP(strings.TrimSpace(val))
	if ip != nil {
		return ip.String(), nil
	}
	return "", fmt.Errorf("%s is not an ip address", val)
}

// Validates domain for resolvconf search configuration.
// A zero length domain is represented by .
func ValidateDnsSearch(val string) (string, error) {
	if val = strings.Trim(val, " "); val == "." {
		return val, nil
	}
	return validateDomain(val)
}

func validateDomain(val string) (string, error) {
	alpha := regexp.MustCompile(`[a-zA-Z]`)
	if alpha.FindString(val) == "" {
		return "", fmt.Errorf("%s is not a valid domain", val)
	}
	re := regexp.MustCompile(`^(:?(:?[a-zA-Z0-9]|(:?[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9]))(:?\.(:?[a-zA-Z0-9]|(:?[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])))*)\.?\s*$`)
	ns := re.FindSubmatch([]byte(val))
	if len(ns) > 0 {
		return string(ns[1]), nil
	}
	return "", fmt.Errorf("%s is not a valid domain", val)
}
<commit_msg>Make --attach case-insensitive<commit_after>package opts

import (
	"fmt"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/docker/docker/pkg/parsers"
)

// ListOpts type
type ListOpts struct {
	values    []string
	validator ValidatorFctType
}

func NewListOpts(validator ValidatorFctType) ListOpts {
	return ListOpts{
		validator: validator,
	}
}

func (opts *ListOpts) String() string {
	return fmt.Sprintf("%v", []string(opts.values))
}

// Set validates if needed the input value and add it to the
// internal slice.
func (opts *ListOpts) Set(value string) error {
	if opts.validator != nil {
		v, err := opts.validator(value)
		if err != nil {
			return err
		}
		value = v
	}
	opts.values = append(opts.values, value)
	return nil
}

// Delete remove the given element from the slice.
func (opts *ListOpts) Delete(key string) {
	for i, k := range opts.values {
		if k == key {
			opts.values = append(opts.values[:i], opts.values[i+1:]...)
			return
		}
	}
}

// GetMap returns the content of values in a map in order to avoid
// duplicates.
// FIXME: can we remove this?
func (opts *ListOpts) GetMap() map[string]struct{} {
	ret := make(map[string]struct{})
	for _, k := range opts.values {
		ret[k] = struct{}{}
	}
	return ret
}

// GetAll returns the values' slice.
// FIXME: Can we remove this?
func (opts *ListOpts) GetAll() []string {
	return opts.values
}

// Get checks the existence of the given key.
func (opts *ListOpts) Get(key string) bool {
	for _, k := range opts.values {
		if k == key {
			return true
		}
	}
	return false
}

// Len returns the amount of element in the slice.
func (opts *ListOpts) Len() int {
	return len(opts.values)
}

// Validators
type ValidatorFctType func(val string) (string, error)

func ValidateAttach(val string) (string, error) {
	s := strings.ToLower(val)
	for _, str := range []string{"stdin", "stdout", "stderr"} {
		if s == str {
			return s, nil
		}
	}
	return val, fmt.Errorf("valid streams are STDIN, STDOUT and STDERR.")
}

func ValidateLink(val string) (string, error) {
	if _, err := parsers.PartParser("name:alias", val); err != nil {
		return val, err
	}
	return val, nil
}

func ValidatePath(val string) (string, error) {
	var containerPath string

	if strings.Count(val, ":") > 2 {
		return val, fmt.Errorf("bad format for volumes: %s", val)
	}

	splited := strings.SplitN(val, ":", 2)
	if len(splited) == 1 {
		containerPath = splited[0]
		val = filepath.Clean(splited[0])
	} else {
		containerPath = splited[1]
		val = fmt.Sprintf("%s:%s", splited[0], filepath.Clean(splited[1]))
	}

	if !filepath.IsAbs(containerPath) {
		return val, fmt.Errorf("%s is not an absolute path", containerPath)
	}
	return val, nil
}

func ValidateEnv(val string) (string, error) {
	arr := strings.Split(val, "=")
	if len(arr) > 1 {
		return val, nil
	}
	return fmt.Sprintf("%s=%s", val, os.Getenv(val)), nil
}

func ValidateIPAddress(val string) (string, error) {
	var ip = net.ParseIP(strings.TrimSpace(val))
	if ip != nil {
		return ip.String(), nil
	}
	return "", fmt.Errorf("%s is not an ip address", val)
}

// Validates domain for resolvconf search configuration.
// A zero length domain is represented by .
func ValidateDnsSearch(val string) (string, error) {
	if val = strings.Trim(val, " "); val == "." {
		return val, nil
	}
	return validateDomain(val)
}

func validateDomain(val string) (string, error) {
	alpha := regexp.MustCompile(`[a-zA-Z]`)
	if alpha.FindString(val) == "" {
		return "", fmt.Errorf("%s is not a valid domain", val)
	}
	re := regexp.MustCompile(`^(:?(:?[a-zA-Z0-9]|(:?[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9]))(:?\.(:?[a-zA-Z0-9]|(:?[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])))*)\.?\s*$`)
	ns := re.FindSubmatch([]byte(val))
	if len(ns) > 0 {
		return string(ns[1]), nil
	}
	return "", fmt.Errorf("%s is not a valid domain", val)
}
<|endoftext|>


<commit_before>// Copyright 2013, Cong Ding. 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.
//
// author: Cong Ding <dinggnu@gmail.com>

// Package logging implements log library for other applications. It provides
// functions Debug, Info, Warning, Error, Critical, and formatting version
// Logf.
//
// Example:
//
//	logger := logging.SimpleLogger("main")
//	logger.SetLevel(logging.WARNING)
//	logger.Error("test for error")
//	logger.Warning("test for warning", "second parameter")
//	logger.Debug("test for debug")
//
package logging

import (
	"io"
	"os"
	"sync"
	"sync/atomic"
	"time"
)

// Pre-defined formats
const (
	DefaultFileName       = "logging.log"                   // default logging filename
	DefaultConfigFileName = "logging.conf"                  // default configuration filename
	DefaultTimeFormat     = "2006-01-02 15:04:05.999999999" // defaulttime format
	bufSize               = 1000                            // buffer size for writer
	queueSize             = 10000                           // chan queue size in async logging
	reqSize               = 10000                           // chan queue size in async logging
	flushSize             = 1                               // chen queue used to flush the writer
)

// Logger is the logging struct.
type Logger struct {

	// Be careful of the alignment issue of the variable seqid because it
	// uses the sync/atomic.AddUint64() operation. If the alignment is
	// wrong, it will cause a panic. To solve the alignment issue in an
	// easy way, we put seqid to the beginning of the structure.
	// seqid is only visiable internally.
	seqid uint64 // last used sequence number in record

	// These variables can be configured by users.
	name       string    // logger name
	level      Level     // record level higher than this will be printed
	rfmt       string    // format of the record
	rargs      []string  // arguments to be used in the rfmt
	out        io.Writer // writer
	sync       bool      // use sync or async way to record logs
	timeFormat string    // format for time

	// These variables are visible to users.
	startTime time.Time // start time of the logger

	// Internally used variables, which don't have get and set functions.
	wlock   sync.Mutex   // writer lock
	queue   chan string  // queue used in async logging
	request chan request // queue used in non-runtime logging
	flush   chan bool    // flush signal for the watcher to write
	quit    chan bool    // quit signal for the watcher to quit
	fd      *os.File     // file handler, used to close the file on destroy
	runtime bool         // with runtime operation or not
}

// request struct stores the logger request
type request struct {
	level  Level
	format string
	v      []interface{}
}

// SimpleLogger creates a new logger with simple configuration.
func SimpleLogger(name string) (*Logger, error) {
	return createLogger(name, WARNING, BasicFormat, os.Stdout, false)
}

// BasicLogger creates a new logger with basic configuration.
func BasicLogger(name string) (*Logger, error) {
	return FileLogger(name, WARNING, BasicFormat, DefaultFileName, false)
}

// RichLogger creates a new logger with simple configuration.
func RichLogger(name string) (*Logger, error) {
	return FileLogger(name, NOTSET, RichFormat, DefaultFileName, false)
}

// FileLogger creates a new logger with file output.
func FileLogger(name string, level Level, format string, file string, sync bool) (*Logger, error) {
	out, err := os.Create(file)
	if err != nil {
		return new(Logger), err
	}
	logger, err := createLogger(name, level, format, out, sync)
	if err == nil {
		logger.fd = out
	}
	return logger, err
}

func WriterLogger(name string, level Level, format string, out io.Writer, sync bool) (*Logger, error) {
	return createLogger(name, level, format, out, sync)
}

// createLogger create a new logger
func createLogger(name string, level Level, format string, out io.Writer, sync bool) (*Logger, error) {
	logger := new(Logger)

	err := logger.parseFormat(format)
	if err != nil {
		return logger, err
	}

	// asign values to logger
	logger.name = name
	logger.level = level
	logger.out = out
	logger.seqid = 0
	logger.sync = sync
	logger.queue = make(chan string, queueSize)
	logger.request = make(chan request, reqSize)
	logger.flush = make(chan bool, flushSize)
	logger.quit = make(chan bool)
	logger.startTime = time.Now()
	logger.fd = nil
	logger.timeFormat = DefaultTimeFormat

	// start watcher to write logs if it is async or no runtime field
	if !logger.sync {
		go logger.watcher()
	}

	return logger, nil
}

// Destroy sends quit signal to watcher and releases all the resources.
func (logger *Logger) Destroy() {
	logger.quitWatcher()
	// clean up
	if logger.fd != nil {
		logger.fd.Close()
	}
}

func (logger *Logger) quitWatcher() {
	if !logger.sync {
		// quit watcher
		logger.quit <- true
		// wait for watcher quit
		<-logger.quit
	}
}

// Flush the writer
func (logger *Logger) Flush() {
	if !logger.sync {
		logger.flush <- true
	}
}

// Get and set the configuration of the logger

func (logger *Logger) Name() string {
	return logger.name
}

func (logger *Logger) SetName(name string) {
	logger.name = name
}

func (logger *Logger) TimeFormat() string {
	return logger.timeFormat
}

func (logger *Logger) Level() Level {
	return Level(atomic.LoadInt32((*int32)(&logger.level)))
}

func (logger *Logger) SetLevel(level Level) {
	atomic.StoreInt32((*int32)(&logger.level), int32(level))
}

func (logger *Logger) Format() string {
	return logger.rfmt
}

func (logger *Logger) Fargs() []string {
	return logger.rargs
}

func (logger *Logger) Writer() io.Writer {
	return logger.out
}

func (logger *Logger) AddWriter(out io.Writer) {
	logger.out = io.MultiWriter(logger.out, out)
}

func (logger *Logger) AddWriters(out ...io.Writer) {
	logger.out = io.MultiWriter(append(out, logger.out)...)
}

func (logger *Logger) SetWriter(out io.Writer) {
	logger.out = out
}

func (logger *Logger) SetWriters(out ...io.Writer) {
	logger.out = io.MultiWriter(out...)
}

func (logger *Logger) Sync() bool {
	return logger.sync
}
<commit_msg>clean up<commit_after>// Copyright 2013, Cong Ding. 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.
//
// author: Cong Ding <dinggnu@gmail.com>

// Package logging implements log library for other applications. It provides
// functions Debug, Info, Warning, Error, Critical, and formatting version
// Logf.
//
// Example:
//
//	logger := logging.SimpleLogger("main")
//	logger.SetLevel(logging.WARNING)
//	logger.Error("test for error")
//	logger.Warning("test for warning", "second parameter")
//	logger.Debug("test for debug")
//
package logging

import (
	"io"
	"os"
	"sync"
	"sync/atomic"
	"time"
)

// Pre-defined formats
const (
	DefaultFileName       = "logging.log"                   // default logging filename
	DefaultConfigFileName = "logging.conf"                  // default configuration filename
	DefaultTimeFormat     = "2006-01-02 15:04:05.999999999" // defaulttime format
	bufSize               = 1000                            // buffer size for writer
	queueSize             = 10000                           // chan queue size in async logging
	reqSize               = 10000                           // chan queue size in async logging
	flushSize             = 1                               // chen queue used to flush the writer
)

// Logger is the logging struct.
type Logger struct {

	// Be careful of the alignment issue of the variable seqid because it
	// uses the sync/atomic.AddUint64() operation. If the alignment is
	// wrong, it will cause a panic. To solve the alignment issue in an
	// easy way, we put seqid to the beginning of the structure.
	// seqid is only visiable internally.
	seqid uint64 // last used sequence number in record

	// These variables can be configured by users.
	name       string    // logger name
	level      Level     // record level higher than this will be printed
	rfmt       string    // format of the record
	rargs      []string  // arguments to be used in the rfmt
	out        io.Writer // writer
	sync       bool      // use sync or async way to record logs
	timeFormat string    // format for time

	// These variables are visible to users.
	startTime time.Time // start time of the logger

	// Internally used variables, which don't have get and set functions.
	wlock   sync.Mutex   // writer lock
	queue   chan string  // queue used in async logging
	request chan request // queue used in non-runtime logging
	flush   chan bool    // flush signal for the watcher to write
	quit    chan bool    // quit signal for the watcher to quit
	fd      *os.File     // file handler, used to close the file on destroy
	runtime bool         // with runtime operation or not
}

// request struct stores the logger request
type request struct {
	level  Level
	format string
	v      []interface{}
}

// SimpleLogger creates a new logger with simple configuration.
func SimpleLogger(name string) (*Logger, error) {
	return createLogger(name, WARNING, BasicFormat, os.Stdout, false)
}

// BasicLogger creates a new logger with basic configuration.
func BasicLogger(name string) (*Logger, error) {
	return FileLogger(name, WARNING, BasicFormat, DefaultFileName, false)
}

// RichLogger creates a new logger with simple configuration.
func RichLogger(name string) (*Logger, error) {
	return FileLogger(name, NOTSET, RichFormat, DefaultFileName, false)
}

// FileLogger creates a new logger with file output.
func FileLogger(name string, level Level, format string, file string, sync bool) (*Logger, error) {
	out, err := os.Create(file)
	if err != nil {
		return new(Logger), err
	}
	logger, err := createLogger(name, level, format, out, sync)
	if err == nil {
		logger.fd = out
	}
	return logger, err
}

func WriterLogger(name string, level Level, format string, out io.Writer, sync bool) (*Logger, error) {
	return createLogger(name, level, format, out, sync)
}

// createLogger create a new logger
func createLogger(name string, level Level, format string, out io.Writer, sync bool) (*Logger, error) {
	logger := new(Logger)

	err := logger.parseFormat(format)
	if err != nil {
		return logger, err
	}

	// asign values to logger
	logger.name = name
	logger.level = level
	logger.out = out
	logger.seqid = 0
	logger.sync = sync
	logger.queue = make(chan string, queueSize)
	logger.request = make(chan request, reqSize)
	logger.flush = make(chan bool, flushSize)
	logger.quit = make(chan bool)
	logger.startTime = time.Now()
	logger.fd = nil
	logger.timeFormat = DefaultTimeFormat

	// start watcher to write logs if it is async or no runtime field
	if !logger.sync {
		go logger.watcher()
	}

	return logger, nil
}

// Destroy sends quit signal to watcher and releases all the resources.
func (logger *Logger) Destroy() {
	if !logger.sync {
		// quit watcher
		logger.quit <- true
		// wait for watcher quit
		<-logger.quit
	}
	// clean up
	if logger.fd != nil {
		logger.fd.Close()
	}
}

// Flush the writer
func (logger *Logger) Flush() {
	if !logger.sync {
		logger.flush <- true
	}
}

// Get and set the configuration of the logger

func (logger *Logger) Name() string {
	return logger.name
}

func (logger *Logger) SetName(name string) {
	logger.name = name
}

func (logger *Logger) TimeFormat() string {
	return logger.timeFormat
}

func (logger *Logger) Level() Level {
	return Level(atomic.LoadInt32((*int32)(&logger.level)))
}

func (logger *Logger) SetLevel(level Level) {
	atomic.StoreInt32((*int32)(&logger.level), int32(level))
}

func (logger *Logger) Format() string {
	return logger.rfmt
}

func (logger *Logger) Fargs() []string {
	return logger.rargs
}

func (logger *Logger) Writer() io.Writer {
	return logger.out
}

func (logger *Logger) AddWriter(out io.Writer) {
	logger.out = io.MultiWriter(logger.out, out)
}

func (logger *Logger) AddWriters(out ...io.Writer) {
	logger.out = io.MultiWriter(append(out, logger.out)...)
}

func (logger *Logger) SetWriter(out io.Writer) {
	logger.out = out
}

func (logger *Logger) SetWriters(out ...io.Writer) {
	logger.out = io.MultiWriter(out...)
}

func (logger *Logger) Sync() bool {
	return logger.sync
}
<|endoftext|>


<commit_before>package jsonrpc

// Router is a JSON-RPC request routing middleware.
type Router struct {
	jsonrpc        *App
	reqRoutes      map[string]func(ctx *ReqContext)
	notRoutes      map[string]func(ctx *NotContext)
	pendinRequests map[string]chan *Response // requests sent from the router that are pending responses from clients
}

// NewRouter creates a JSON-RPC router instance and registers it with the Neptulon JSON-RPC app.
func NewRouter(app *App) (*Router, error) {
	r := Router{
		jsonrpc:        app,
		reqRoutes:      make(map[string]func(ctx *ReqContext)),
		notRoutes:      make(map[string]func(ctx *NotContext)),
		pendinRequests: make(map[string]chan *Response),
	}

	app.ReqMiddleware(r.reqMiddleware)
	app.NotMiddleware(r.notMiddleware)
	app.ResMiddleware(r.resMiddleware)
	return &r, nil
}

// Request adds a new request route registry.
// Optionally, you can pass in a data structure that the returned JSON-RPC response result data will be serialized into. Otherwise json.Unmarshal defaults apply.
func (r *Router) Request(route string, resultData interface{}, handler func(ctx *ReqContext)) {
	r.reqRoutes[route] = handler
}

// Notification adds a new notification route registry.
func (r *Router) Notification(route string, handler func(ctx *NotContext)) {
	r.notRoutes[route] = handler
}

// SendRequest sends a JSON-RPC request throught the connection denoted by the connection ID.
func (r *Router) SendRequest(connID string, req *Request) chan<- *Response {
	r.jsonrpc.Send(connID, req)
	ch := make(chan *Response)
	r.pendinRequests[req.ID] = ch
	return ch
}

// SendNotification sends a JSON-RPC notification through the connection denoted by the connection ID.
func (r *Router) SendNotification(connID string, not *Notification) {
	r.jsonrpc.Send(connID, not)
}

func (r *Router) reqMiddleware(ctx *ReqContext) {
	if handler, ok := r.reqRoutes[ctx.Req.Method]; ok {
		handler(ctx)
	}
}

func (r *Router) notMiddleware(ctx *NotContext) {
	if handler, ok := r.notRoutes[ctx.Not.Method]; ok {
		handler(ctx)
	}
}

func (r *Router) resMiddleware(ctx *ResContext) {
	if ch, ok := r.pendinRequests[ctx.Res.ID]; ok {
		ch <- ctx.Res
		delete(r.pendinRequests, ctx.Res.ID)
	}
}
<commit_msg>add type to Router.SendRequest<commit_after>package jsonrpc

// Router is a JSON-RPC message routing middleware.
type Router struct {
	jsonrpc        *App
	reqRoutes      map[string]func(ctx *ReqContext)
	notRoutes      map[string]func(ctx *NotContext)
	pendinRequests map[string]chan *Response // requests sent from the router that are pending responses from clients
}

// NewRouter creates a JSON-RPC router instance and registers it with the Neptulon JSON-RPC app.
func NewRouter(app *App) (*Router, error) {
	r := Router{
		jsonrpc:        app,
		reqRoutes:      make(map[string]func(ctx *ReqContext)),
		notRoutes:      make(map[string]func(ctx *NotContext)),
		pendinRequests: make(map[string]chan *Response),
	}

	app.ReqMiddleware(r.reqMiddleware)
	app.NotMiddleware(r.notMiddleware)
	app.ResMiddleware(r.resMiddleware)
	return &r, nil
}

// Request adds a new incoming request route registry.
// Optionally, you can pass in a data structure that the returned JSON-RPC response result data will be serialized into. Otherwise json.Unmarshal defaults apply.
func (r *Router) Request(route string, resultData interface{}, handler func(ctx *ReqContext)) {
	r.reqRoutes[route] = handler
}

// Notification adds a new incoming notification route registry.
func (r *Router) Notification(route string, handler func(ctx *NotContext)) {
	r.notRoutes[route] = handler
}

// SendRequest sends a JSON-RPC request throught the connection denoted by the connection ID.
// Optionally, you can pass in a data structure that the returned JSON-RPC response result data will be serialized into. Otherwise json.Unmarshal defaults apply.
func (r *Router) SendRequest(connID string, resultData interface{}, req *Request) chan<- *Response {
	r.jsonrpc.Send(connID, req)
	ch := make(chan *Response)
	r.pendinRequests[req.ID] = ch
	return ch
}

// SendNotification sends a JSON-RPC notification through the connection denoted by the connection ID.
func (r *Router) SendNotification(connID string, not *Notification) {
	r.jsonrpc.Send(connID, not)
}

func (r *Router) reqMiddleware(ctx *ReqContext) {
	if handler, ok := r.reqRoutes[ctx.Req.Method]; ok {
		handler(ctx)
	}
}

func (r *Router) notMiddleware(ctx *NotContext) {
	if handler, ok := r.notRoutes[ctx.Not.Method]; ok {
		handler(ctx)
	}
}

func (r *Router) resMiddleware(ctx *ResContext) {
	if ch, ok := r.pendinRequests[ctx.Res.ID]; ok {
		ch <- ctx.Res
		delete(r.pendinRequests, ctx.Res.ID)
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 Netflix Inc
// Author: Colin McIntosh (colin@netflix.com)
//
// 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 prometheus

import (
	"errors"
	"fmt"
	"github.com/openconfig/gnmi-gateway/gateway/configuration"
	"github.com/openconfig/gnmi-gateway/gateway/exporters"
	"github.com/openconfig/gnmi-gateway/gateway/openconfig"
	"github.com/openconfig/gnmi/cache"
	"github.com/openconfig/gnmi/ctree"
	gnmipb "github.com/openconfig/gnmi/proto/gnmi"
	prom "github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"net/http"
	"strings"
)

const Name = "prometheus"

func init() {
	exporters.Register(Name, NewPrometheusExporter)
}

func NewPrometheusExporter(config *configuration.GatewayConfig) exporters.Exporter {
	return &PrometheusExporter{
		config:     config,
		deltaCalc:  NewDeltaCalculator(),
		metrics:    make(map[Hash]prom.Metric),
		typeLookup: new(openconfig.TypeLookup),
	}
}

type PrometheusExporter struct {
	config     *configuration.GatewayConfig
	cache      *cache.Cache
	deltaCalc  *DeltaCalculator
	metrics    map[Hash]prom.Metric
	typeLookup *openconfig.TypeLookup
}

func (e *PrometheusExporter) Name() string {
	return Name
}

func (e *PrometheusExporter) Export(leaf *ctree.Leaf) {
	notification := leaf.Value().(*gnmipb.Notification)
	for _, update := range notification.Update {
		value, isNumber := GetNumberValues(update.Val)
		if !isNumber {
			continue
		}
		metricName, labels := UpdateToMetricNameAndLabels(notification.GetPrefix(), update)
		metricHash := NewStringMapHash(metricName, labels)

		metric, exists := e.metrics[metricHash]
		if !exists {
			var path []string
			for _, elem := range update.Path.Elem {
				path = append(path, elem.Name)
			}
			metricType := e.typeLookup.GetTypeByPath(path)

			switch metricType {
			case "counter64":
				metric = promauto.NewCounter(prom.CounterOpts{
					Name:        metricName,
					ConstLabels: labels,
				})
			case "gauge32":
			default:
				metric = promauto.NewGauge(prom.GaugeOpts{
					Name:        metricName,
					ConstLabels: labels,
				})
			}
			e.metrics[metricHash] = metric
		}

		switch m := metric.(type) {
		case prom.Counter:
			delta, exists := e.deltaCalc.Calc(metricHash, value)
			if exists {
				m.Add(delta)
			}
		case prom.Gauge:
			m.Set(value)
		}
	}
}

func (e *PrometheusExporter) Start(cache *cache.Cache) error {
	e.config.Log.Info().Msg("Starting Prometheus exporter.")
	if e.config.OpenConfigDirectory == "" {
		return errors.New("value is not set for OpenConfigDirectory configuration")
	}
	e.cache = cache
	err := e.typeLookup.LoadAllModules(e.config.OpenConfigDirectory)
	if err != nil {
		e.config.Log.Error().Err(err).Msgf("Unable to load OpenConfig modules in %s: %v", e.config.OpenConfigDirectory, err)
		return err
	}

	e.config.Log.Info().Msg("Starting Prometheus exporter.")
	e.cache.SetClient(e.Export)
	go e.runHttpServer()
	return nil
}

func (e *PrometheusExporter) runHttpServer() {
	var errCount = 0
	var lastError error
	for {
		e.config.Log.Info().Msg("Starting Prometheus HTTP server.")
		http.Handle("/metrics", promhttp.Handler())
		err := http.ListenAndServe(":59100", nil)
		if err != nil {
			e.config.Log.Error().Err(err).Msgf("Prometheus HTTP server stopped with an error: %v", err)
			if err.Error() == lastError.Error() {
				errCount = errCount + 1
				if errCount >= 3 {
					panic(fmt.Errorf("too many errors returned by Prometheus HTTP server: %s", err.Error()))
				}
			} else {
				errCount = 0
				lastError = err
			}
		}
	}
}

func GetNumberValues(tv *gnmipb.TypedValue) (float64, bool) {
	if tv != nil && tv.Value != nil {
		switch tv.Value.(type) {
		case *gnmipb.TypedValue_StringVal:
			return 0, false
		case *gnmipb.TypedValue_IntVal:
			return float64(tv.GetIntVal()), true
		case *gnmipb.TypedValue_UintVal:
			return float64(tv.GetUintVal()), true
		case *gnmipb.TypedValue_BoolVal:
			if tv.GetBoolVal() {
				return 1, true
			} else {
				return 0, false
			}
		case *gnmipb.TypedValue_FloatVal:
			return float64(tv.GetFloatVal()), true
		case *gnmipb.TypedValue_LeaflistVal:
			return 0, false
		case *gnmipb.TypedValue_BytesVal:
			return 0, false
		default:
			return 0, false
		}
	}
	return 0, false
}

func UpdateToMetricNameAndLabels(prefix *gnmipb.Path, update *gnmipb.Update) (string, map[string]string) {
	metricName := ""
	labels := make(map[string]string)

	if prefix != nil {
		target := prefix.GetTarget()
		if target != "" {
			labels["target"] = target
		}
	}

	for _, elem := range update.Path.Elem {
		elemName := strings.ReplaceAll(elem.Name, "-", "_")
		if metricName == "" {
			metricName = elemName
		} else {
			metricName = metricName + "_" + elemName
		}

		for key, value := range elem.Key {
			labelKey := metricName + "_" + strings.ReplaceAll(key, "-", "_")
			labels[labelKey] = value
		}
	}
	return metricName, labels
}
<commit_msg>Remove code that overwrites cache client. This is old code that should have been previously removed.<commit_after>// Copyright 2020 Netflix Inc
// Author: Colin McIntosh (colin@netflix.com)
//
// 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 prometheus

import (
	"errors"
	"fmt"
	"github.com/openconfig/gnmi-gateway/gateway/configuration"
	"github.com/openconfig/gnmi-gateway/gateway/exporters"
	"github.com/openconfig/gnmi-gateway/gateway/openconfig"
	"github.com/openconfig/gnmi/cache"
	"github.com/openconfig/gnmi/ctree"
	gnmipb "github.com/openconfig/gnmi/proto/gnmi"
	prom "github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"net/http"
	"strings"
)

const Name = "prometheus"

func init() {
	exporters.Register(Name, NewPrometheusExporter)
}

func NewPrometheusExporter(config *configuration.GatewayConfig) exporters.Exporter {
	return &PrometheusExporter{
		config:     config,
		deltaCalc:  NewDeltaCalculator(),
		metrics:    make(map[Hash]prom.Metric),
		typeLookup: new(openconfig.TypeLookup),
	}
}

type PrometheusExporter struct {
	config     *configuration.GatewayConfig
	cache      *cache.Cache
	deltaCalc  *DeltaCalculator
	metrics    map[Hash]prom.Metric
	typeLookup *openconfig.TypeLookup
}

func (e *PrometheusExporter) Name() string {
	return Name
}

func (e *PrometheusExporter) Export(leaf *ctree.Leaf) {
	notification := leaf.Value().(*gnmipb.Notification)
	for _, update := range notification.Update {
		value, isNumber := GetNumberValues(update.Val)
		if !isNumber {
			continue
		}
		metricName, labels := UpdateToMetricNameAndLabels(notification.GetPrefix(), update)
		metricHash := NewStringMapHash(metricName, labels)

		metric, exists := e.metrics[metricHash]
		if !exists {
			var path []string
			for _, elem := range update.Path.Elem {
				path = append(path, elem.Name)
			}
			metricType := e.typeLookup.GetTypeByPath(path)

			switch metricType {
			case "counter64":
				metric = promauto.NewCounter(prom.CounterOpts{
					Name:        metricName,
					ConstLabels: labels,
				})
			case "gauge32":
			default:
				metric = promauto.NewGauge(prom.GaugeOpts{
					Name:        metricName,
					ConstLabels: labels,
				})
			}
			e.metrics[metricHash] = metric
		}

		switch m := metric.(type) {
		case prom.Counter:
			delta, exists := e.deltaCalc.Calc(metricHash, value)
			if exists {
				m.Add(delta)
			}
		case prom.Gauge:
			m.Set(value)
		}
	}
}

func (e *PrometheusExporter) Start(cache *cache.Cache) error {
	e.config.Log.Info().Msg("Starting Prometheus exporter.")
	if e.config.OpenConfigDirectory == "" {
		return errors.New("value is not set for OpenConfigDirectory configuration")
	}
	e.cache = cache
	err := e.typeLookup.LoadAllModules(e.config.OpenConfigDirectory)
	if err != nil {
		e.config.Log.Error().Err(err).Msgf("Unable to load OpenConfig modules in %s: %v", e.config.OpenConfigDirectory, err)
		return err
	}
	go e.runHttpServer()
	return nil
}

func (e *PrometheusExporter) runHttpServer() {
	var errCount = 0
	var lastError error
	for {
		e.config.Log.Info().Msg("Starting Prometheus HTTP server.")
		http.Handle("/metrics", promhttp.Handler())
		err := http.ListenAndServe(":59100", nil)
		if err != nil {
			e.config.Log.Error().Err(err).Msgf("Prometheus HTTP server stopped with an error: %v", err)
			if err.Error() == lastError.Error() {
				errCount = errCount + 1
				if errCount >= 3 {
					panic(fmt.Errorf("too many errors returned by Prometheus HTTP server: %s", err.Error()))
				}
			} else {
				errCount = 0
				lastError = err
			}
		}
	}
}

func GetNumberValues(tv *gnmipb.TypedValue) (float64, bool) {
	if tv != nil && tv.Value != nil {
		switch tv.Value.(type) {
		case *gnmipb.TypedValue_StringVal:
			return 0, false
		case *gnmipb.TypedValue_IntVal:
			return float64(tv.GetIntVal()), true
		case *gnmipb.TypedValue_UintVal:
			return float64(tv.GetUintVal()), true
		case *gnmipb.TypedValue_BoolVal:
			if tv.GetBoolVal() {
				return 1, true
			} else {
				return 0, false
			}
		case *gnmipb.TypedValue_FloatVal:
			return float64(tv.GetFloatVal()), true
		case *gnmipb.TypedValue_LeaflistVal:
			return 0, false
		case *gnmipb.TypedValue_BytesVal:
			return 0, false
		default:
			return 0, false
		}
	}
	return 0, false
}

func UpdateToMetricNameAndLabels(prefix *gnmipb.Path, update *gnmipb.Update) (string, map[string]string) {
	metricName := ""
	labels := make(map[string]string)

	if prefix != nil {
		target := prefix.GetTarget()
		if target != "" {
			labels["target"] = target
		}
	}

	for _, elem := range update.Path.Elem {
		elemName := strings.ReplaceAll(elem.Name, "-", "_")
		if metricName == "" {
			metricName = elemName
		} else {
			metricName = metricName + "_" + elemName
		}

		for key, value := range elem.Key {
			labelKey := metricName + "_" + strings.ReplaceAll(key, "-", "_")
			labels[labelKey] = value
		}
	}
	return metricName, labels
}
<|endoftext|>


<commit_before>// Copyright 2015 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.

// +build ignore
//
// This build tag means that "go install golang.org/x/exp/shiny/..." doesn't
// install this example program. Use "go run main.go board.go xy.go" to run it.

// Goban is a simple example of a graphics program using shiny.
// It implements a Go board that two people can use to play the game.
// TODO: Improve the main function.
// TODO: Provide more functionality.
package main

import (
	"flag"
	"image"
	"image/color"
	stdDraw "image/draw"
	"log"
	"math/rand"
	"time"

	"golang.org/x/exp/shiny/driver"
	"golang.org/x/exp/shiny/screen"
	"golang.org/x/mobile/event/key"
	"golang.org/x/mobile/event/mouse"
	"golang.org/x/mobile/event/paint"
	"golang.org/x/mobile/event/size"
)

var dirty bool
var uploading bool

var scale = flag.Int("scale", 35, "`percent` to scale images (TODO: a poor design)")

func main() {
	flag.Parse()

	rand.Seed(int64(time.Now().Nanosecond()))
	board := NewBoard(9, *scale)

	driver.Main(func(s screen.Screen) {
		w, err := s.NewWindow(nil)
		if err != nil {
			log.Fatal(err)
		}
		defer w.Release()

		r := board.image.Bounds()
		winSize := r.Size()
		var b screen.Buffer
		defer func() {
			if b != nil {
				b.Release()
			}
		}()

		var sz size.Event

		for e := range w.Events() {
			switch e := e.(type) {
			default:

			case mouse.Event:
				if e.Direction == mouse.DirRelease && e.Button != 0 {
					// Invert y. TODO: for Darwin gldriver bug that will be fixed by https://go-review.googlesource.com/#/c/13917/
					y := b.RGBA().Bounds().Dy() - int(e.Y)
					board.click(b.RGBA(), int(e.X), y, int(e.Button))
					dirty = true
				}

			case key.Event:
				if e.Code == key.CodeEscape {
					return
				}

			case paint.Event:
				// TODO: we shouldn't rely on the library sending us a constant
				// stream of paint events. If we're dirty, we should just draw
				// on the window and then call w.Publish.
				//
				// TODO: This check should save CPU time but causes flicker on Darwin.
				//if dirty && !uploading {
				w.Fill(sz.Bounds(), color.RGBA{0x00, 0x00, 0x3f, 0xff}, stdDraw.Src)
				w.Upload(image.Point{0, 0}, b, b.Bounds(), w) // TODO: On Darwin always writes to 0,0, ignoring first arg.
				dirty = false
				uploading = true
				//}
				w.Publish()

			case screen.UploadedEvent:
				// No-op.
				uploading = false

			case size.Event:
				// TODO: Set board size.
				sz = e
				if b != nil {
					b.Release()
				}
				winSize = image.Point{sz.WidthPx, sz.HeightPx}
				b, err = s.NewBuffer(winSize)
				if err != nil {
					log.Fatal(err)
				}
				render(b.RGBA(), board)

			case error:
				log.Print(e)
			}
		}
	})
}

func render(m *image.RGBA, board *Board) {
	board.Draw(m)
	dirty = true
}
<commit_msg>shiny/example/goban: remove y-inversion workaround.<commit_after>// Copyright 2015 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.

// +build ignore
//
// This build tag means that "go install golang.org/x/exp/shiny/..." doesn't
// install this example program. Use "go run main.go board.go xy.go" to run it.

// Goban is a simple example of a graphics program using shiny.
// It implements a Go board that two people can use to play the game.
// TODO: Improve the main function.
// TODO: Provide more functionality.
package main

import (
	"flag"
	"image"
	"log"
	"math/rand"
	"time"

	"golang.org/x/exp/shiny/driver"
	"golang.org/x/exp/shiny/screen"
	"golang.org/x/mobile/event/key"
	"golang.org/x/mobile/event/mouse"
	"golang.org/x/mobile/event/paint"
	"golang.org/x/mobile/event/size"
)

var dirty bool
var uploading bool

var scale = flag.Int("scale", 35, "`percent` to scale images (TODO: a poor design)")

func main() {
	flag.Parse()

	rand.Seed(int64(time.Now().Nanosecond()))
	board := NewBoard(9, *scale)

	driver.Main(func(s screen.Screen) {
		w, err := s.NewWindow(nil)
		if err != nil {
			log.Fatal(err)
		}
		defer w.Release()

		var b screen.Buffer
		defer func() {
			if b != nil {
				b.Release()
			}
		}()

		for e := range w.Events() {
			switch e := e.(type) {
			case mouse.Event:
				if e.Direction == mouse.DirRelease && e.Button != 0 {
					board.click(b.RGBA(), int(e.X), int(e.Y), int(e.Button))
					dirty = true
				}

			case key.Event:
				if e.Code == key.CodeEscape {
					return
				}

			case paint.Event:
				// TODO: we shouldn't rely on the library sending us a constant
				// stream of paint events. If we're dirty, we should just draw
				// on the window and then call w.Publish.
				//
				// TODO: This check should save CPU time but causes flicker on Darwin.
				//if dirty && !uploading {
				w.Upload(image.Point{}, b, b.Bounds(), w) // TODO: On Darwin always writes to 0,0, ignoring first arg.
				dirty = false
				uploading = true
				//}
				w.Publish()

			case screen.UploadedEvent:
				// No-op.
				uploading = false

			case size.Event:
				// TODO: Set board size.
				if b != nil {
					b.Release()
				}
				b, err = s.NewBuffer(e.Size())
				if err != nil {
					log.Fatal(err)
				}
				render(b.RGBA(), board)

			case error:
				log.Print(e)
			}
		}
	})
}

func render(m *image.RGBA, board *Board) {
	board.Draw(m)
	dirty = true
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package scheduling

import (
	"fmt"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	podName = "pfpod"
)

var _ = SIGDescribe("LimitRange", func() {
	f := framework.NewDefaultFramework("limitrange")

	It("should create a LimitRange with defaults and ensure pod has those defaults applied.", func() {
		By("Creating a LimitRange")

		min := getResourceList("50m", "100Mi", "100Gi")
		max := getResourceList("500m", "500Mi", "500Gi")
		defaultLimit := getResourceList("500m", "500Mi", "500Gi")
		defaultRequest := getResourceList("100m", "200Mi", "200Gi")
		maxLimitRequestRatio := v1.ResourceList{}
		limitRange := newLimitRange("limit-range", v1.LimitTypeContainer,
			min, max,
			defaultLimit, defaultRequest,
			maxLimitRequestRatio)

		By("Setting up watch")
		selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": limitRange.Name}))
		options := metav1.ListOptions{LabelSelector: selector.String()}
		limitRanges, err := f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).List(options)
		Expect(err).NotTo(HaveOccurred(), "failed to query for limitRanges")
		Expect(len(limitRanges.Items)).To(Equal(0))
		options = metav1.ListOptions{
			LabelSelector:   selector.String(),
			ResourceVersion: limitRanges.ListMeta.ResourceVersion,
		}
		w, err := f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Watch(metav1.ListOptions{})
		Expect(err).NotTo(HaveOccurred(), "failed to set up watch")

		By("Submitting a LimitRange")
		limitRange, err = f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Create(limitRange)
		Expect(err).NotTo(HaveOccurred())

		By("Verifying LimitRange creation was observed")
		select {
		case event, _ := <-w.ResultChan():
			if event.Type != watch.Added {
				framework.Failf("Failed to observe pod creation: %v", event)
			}
		case <-time.After(framework.ServiceRespondingTimeout):
			framework.Failf("Timeout while waiting for LimitRange creation")
		}

		By("Fetching the LimitRange to ensure it has proper values")
		limitRange, err = f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Get(limitRange.Name, metav1.GetOptions{})
		expected := v1.ResourceRequirements{Requests: defaultRequest, Limits: defaultLimit}
		actual := v1.ResourceRequirements{Requests: limitRange.Spec.Limits[0].DefaultRequest, Limits: limitRange.Spec.Limits[0].Default}
		err = equalResourceRequirement(expected, actual)
		Expect(err).NotTo(HaveOccurred())

		By("Creating a Pod with no resource requirements")
		pod := f.NewTestPod("pod-no-resources", v1.ResourceList{}, v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).NotTo(HaveOccurred())

		By("Ensuring Pod has resource requirements applied from LimitRange")
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Get(pod.Name, metav1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())
		for i := range pod.Spec.Containers {
			err = equalResourceRequirement(expected, pod.Spec.Containers[i].Resources)
			if err != nil {
				// Print the pod to help in debugging.
				framework.Logf("Pod %+v does not have the expected requirements", pod)
				Expect(err).NotTo(HaveOccurred())
			}
		}

		By("Creating a Pod with partial resource requirements")
		pod = f.NewTestPod("pod-partial-resources", getResourceList("", "150Mi", "150Gi"), getResourceList("300m", "", ""))
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).NotTo(HaveOccurred())

		By("Ensuring Pod has merged resource requirements applied from LimitRange")
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Get(pod.Name, metav1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())
		// This is an interesting case, so it's worth a comment
		// If you specify a Limit, and no Request, the Limit will default to the Request
		// This means that the LimitRange.DefaultRequest will ONLY take affect if a container.resources.limit is not supplied
		expected = v1.ResourceRequirements{Requests: getResourceList("300m", "150Mi", "150Gi"), Limits: getResourceList("300m", "500Mi", "500Gi")}
		for i := range pod.Spec.Containers {
			err = equalResourceRequirement(expected, pod.Spec.Containers[i].Resources)
			if err != nil {
				// Print the pod to help in debugging.
				framework.Logf("Pod %+v does not have the expected requirements", pod)
				Expect(err).NotTo(HaveOccurred())
			}
		}

		By("Failing to create a Pod with less than min resources")
		pod = f.NewTestPod(podName, getResourceList("10m", "50Mi", "50Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).To(HaveOccurred())

		By("Failing to create a Pod with more than max resources")
		pod = f.NewTestPod(podName, getResourceList("600m", "600Mi", "600Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).To(HaveOccurred())

		By("Updating a LimitRange")
		newMin := getResourceList("9m", "49Mi", "49Gi")
		limitRange.Spec.Limits[0].Min = newMin
		limitRange, err = f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Update(limitRange)
		Expect(err).NotTo(HaveOccurred())

		By("Creating a Pod with less than former min resources")
		pod = f.NewTestPod(podName, getResourceList("10m", "50Mi", "50Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).NotTo(HaveOccurred())

		By("Failing to create a Pod with more than max resources")
		pod = f.NewTestPod(podName, getResourceList("600m", "600Mi", "600Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).To(HaveOccurred())

		By("Deleting a LimitRange")
		err = f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Delete(limitRange.Name, metav1.NewDeleteOptions(30))
		Expect(err).NotTo(HaveOccurred())

		By("Verifying the LimitRange was deleted")
		Expect(wait.Poll(time.Second*5, time.Second*30, func() (bool, error) {
			selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": limitRange.Name}))
			options := metav1.ListOptions{LabelSelector: selector.String()}
			limitRanges, err := f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).List(options)

			if err != nil {
				framework.Logf("Unable to retrieve LimitRanges: %v", err)
				return false, nil
			}

			if len(limitRanges.Items) == 0 {
				framework.Logf("limitRange is already deleted")
				return true, nil
			}

			if len(limitRanges.Items) > 0 {
				if limitRanges.Items[0].ObjectMeta.DeletionTimestamp == nil {
					framework.Logf("deletion has not yet been observed")
					return false, nil
				}
				return true, nil
			}

			return false, nil

		})).NotTo(HaveOccurred(), "kubelet never observed the termination notice")

		By("Creating a Pod with more than former max resources")
		pod = f.NewTestPod(podName+"2", getResourceList("600m", "600Mi", "600Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).NotTo(HaveOccurred())
	})

})

func equalResourceRequirement(expected v1.ResourceRequirements, actual v1.ResourceRequirements) error {
	framework.Logf("Verifying requests: expected %v with actual %v", expected.Requests, actual.Requests)
	err := equalResourceList(expected.Requests, actual.Requests)
	if err != nil {
		return err
	}
	framework.Logf("Verifying limits: expected %v with actual %v", expected.Limits, actual.Limits)
	err = equalResourceList(expected.Limits, actual.Limits)
	return err
}

func equalResourceList(expected v1.ResourceList, actual v1.ResourceList) error {
	for k, v := range expected {
		if actualValue, found := actual[k]; !found || (v.Cmp(actualValue) != 0) {
			return fmt.Errorf("resource %v expected %v actual %v", k, v.String(), actualValue.String())
		}
	}
	for k, v := range actual {
		if expectedValue, found := expected[k]; !found || (v.Cmp(expectedValue) != 0) {
			return fmt.Errorf("resource %v expected %v actual %v", k, expectedValue.String(), v.String())
		}
	}
	return nil
}

func getResourceList(cpu, memory string, ephemeralStorage string) v1.ResourceList {
	res := v1.ResourceList{}
	if cpu != "" {
		res[v1.ResourceCPU] = resource.MustParse(cpu)
	}
	if memory != "" {
		res[v1.ResourceMemory] = resource.MustParse(memory)
	}
	if ephemeralStorage != "" {
		res[v1.ResourceEphemeralStorage] = resource.MustParse(ephemeralStorage)
	}
	return res
}

// newLimitRange returns a limit range with specified data
func newLimitRange(name string, limitType v1.LimitType,
	min, max,
	defaultLimit, defaultRequest,
	maxLimitRequestRatio v1.ResourceList) *v1.LimitRange {
	return &v1.LimitRange{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
		},
		Spec: v1.LimitRangeSpec{
			Limits: []v1.LimitRangeItem{
				{
					Type:                 limitType,
					Min:                  min,
					Max:                  max,
					Default:              defaultLimit,
					DefaultRequest:       defaultRequest,
					MaxLimitRequestRatio: maxLimitRequestRatio,
				},
			},
		},
	}
}
<commit_msg>UPSTREAM: <carry>: Disable ephemeral storage in tests<commit_after>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package scheduling

import (
	"fmt"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	podName = "pfpod"
)

var _ = SIGDescribe("LimitRange", func() {
	f := framework.NewDefaultFramework("limitrange")

	It("should create a LimitRange with defaults and ensure pod has those defaults applied.", func() {
		By("Creating a LimitRange")

		min := getResourceList("50m", "100Mi", "")
		max := getResourceList("500m", "500Mi", "")
		defaultLimit := getResourceList("500m", "500Mi", "")
		defaultRequest := getResourceList("100m", "200Mi", "")
		maxLimitRequestRatio := v1.ResourceList{}
		limitRange := newLimitRange("limit-range", v1.LimitTypeContainer,
			min, max,
			defaultLimit, defaultRequest,
			maxLimitRequestRatio)

		By("Setting up watch")
		selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": limitRange.Name}))
		options := metav1.ListOptions{LabelSelector: selector.String()}
		limitRanges, err := f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).List(options)
		Expect(err).NotTo(HaveOccurred(), "failed to query for limitRanges")
		Expect(len(limitRanges.Items)).To(Equal(0))
		options = metav1.ListOptions{
			LabelSelector:   selector.String(),
			ResourceVersion: limitRanges.ListMeta.ResourceVersion,
		}
		w, err := f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Watch(metav1.ListOptions{})
		Expect(err).NotTo(HaveOccurred(), "failed to set up watch")

		By("Submitting a LimitRange")
		limitRange, err = f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Create(limitRange)
		Expect(err).NotTo(HaveOccurred())

		By("Verifying LimitRange creation was observed")
		select {
		case event, _ := <-w.ResultChan():
			if event.Type != watch.Added {
				framework.Failf("Failed to observe pod creation: %v", event)
			}
		case <-time.After(framework.ServiceRespondingTimeout):
			framework.Failf("Timeout while waiting for LimitRange creation")
		}

		By("Fetching the LimitRange to ensure it has proper values")
		limitRange, err = f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Get(limitRange.Name, metav1.GetOptions{})
		expected := v1.ResourceRequirements{Requests: defaultRequest, Limits: defaultLimit}
		actual := v1.ResourceRequirements{Requests: limitRange.Spec.Limits[0].DefaultRequest, Limits: limitRange.Spec.Limits[0].Default}
		err = equalResourceRequirement(expected, actual)
		Expect(err).NotTo(HaveOccurred())

		By("Creating a Pod with no resource requirements")
		pod := f.NewTestPod("pod-no-resources", v1.ResourceList{}, v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).NotTo(HaveOccurred())

		By("Ensuring Pod has resource requirements applied from LimitRange")
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Get(pod.Name, metav1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())
		for i := range pod.Spec.Containers {
			err = equalResourceRequirement(expected, pod.Spec.Containers[i].Resources)
			if err != nil {
				// Print the pod to help in debugging.
				framework.Logf("Pod %+v does not have the expected requirements", pod)
				Expect(err).NotTo(HaveOccurred())
			}
		}

		By("Creating a Pod with partial resource requirements")
		pod = f.NewTestPod("pod-partial-resources", getResourceList("", "150Mi", "150Gi"), getResourceList("300m", "", ""))
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).NotTo(HaveOccurred())

		By("Ensuring Pod has merged resource requirements applied from LimitRange")
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Get(pod.Name, metav1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())
		// This is an interesting case, so it's worth a comment
		// If you specify a Limit, and no Request, the Limit will default to the Request
		// This means that the LimitRange.DefaultRequest will ONLY take affect if a container.resources.limit is not supplied
		expected = v1.ResourceRequirements{Requests: getResourceList("300m", "150Mi", "150Gi"), Limits: getResourceList("300m", "500Mi", "500Gi")}
		for i := range pod.Spec.Containers {
			err = equalResourceRequirement(expected, pod.Spec.Containers[i].Resources)
			if err != nil {
				// Print the pod to help in debugging.
				framework.Logf("Pod %+v does not have the expected requirements", pod)
				Expect(err).NotTo(HaveOccurred())
			}
		}

		By("Failing to create a Pod with less than min resources")
		pod = f.NewTestPod(podName, getResourceList("10m", "50Mi", "50Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).To(HaveOccurred())

		By("Failing to create a Pod with more than max resources")
		pod = f.NewTestPod(podName, getResourceList("600m", "600Mi", "600Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).To(HaveOccurred())

		By("Updating a LimitRange")
		newMin := getResourceList("9m", "49Mi", "49Gi")
		limitRange.Spec.Limits[0].Min = newMin
		limitRange, err = f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Update(limitRange)
		Expect(err).NotTo(HaveOccurred())

		By("Creating a Pod with less than former min resources")
		pod = f.NewTestPod(podName, getResourceList("10m", "50Mi", "50Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).NotTo(HaveOccurred())

		By("Failing to create a Pod with more than max resources")
		pod = f.NewTestPod(podName, getResourceList("600m", "600Mi", "600Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).To(HaveOccurred())

		By("Deleting a LimitRange")
		err = f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).Delete(limitRange.Name, metav1.NewDeleteOptions(30))
		Expect(err).NotTo(HaveOccurred())

		By("Verifying the LimitRange was deleted")
		Expect(wait.Poll(time.Second*5, time.Second*30, func() (bool, error) {
			selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": limitRange.Name}))
			options := metav1.ListOptions{LabelSelector: selector.String()}
			limitRanges, err := f.ClientSet.CoreV1().LimitRanges(f.Namespace.Name).List(options)

			if err != nil {
				framework.Logf("Unable to retrieve LimitRanges: %v", err)
				return false, nil
			}

			if len(limitRanges.Items) == 0 {
				framework.Logf("limitRange is already deleted")
				return true, nil
			}

			if len(limitRanges.Items) > 0 {
				if limitRanges.Items[0].ObjectMeta.DeletionTimestamp == nil {
					framework.Logf("deletion has not yet been observed")
					return false, nil
				}
				return true, nil
			}

			return false, nil

		})).NotTo(HaveOccurred(), "kubelet never observed the termination notice")

		By("Creating a Pod with more than former max resources")
		pod = f.NewTestPod(podName+"2", getResourceList("600m", "600Mi", "600Gi"), v1.ResourceList{})
		pod, err = f.ClientSet.CoreV1().Pods(f.Namespace.Name).Create(pod)
		Expect(err).NotTo(HaveOccurred())
	})

})

func equalResourceRequirement(expected v1.ResourceRequirements, actual v1.ResourceRequirements) error {
	framework.Logf("Verifying requests: expected %v with actual %v", expected.Requests, actual.Requests)
	err := equalResourceList(expected.Requests, actual.Requests)
	if err != nil {
		return err
	}
	framework.Logf("Verifying limits: expected %v with actual %v", expected.Limits, actual.Limits)
	err = equalResourceList(expected.Limits, actual.Limits)
	return err
}

func equalResourceList(expected v1.ResourceList, actual v1.ResourceList) error {
	for k, v := range expected {
		if actualValue, found := actual[k]; !found || (v.Cmp(actualValue) != 0) {
			return fmt.Errorf("resource %v expected %v actual %v", k, v.String(), actualValue.String())
		}
	}
	for k, v := range actual {
		if expectedValue, found := expected[k]; !found || (v.Cmp(expectedValue) != 0) {
			return fmt.Errorf("resource %v expected %v actual %v", k, expectedValue.String(), v.String())
		}
	}
	return nil
}

func getResourceList(cpu, memory string, ephemeralStorage string) v1.ResourceList {
	res := v1.ResourceList{}
	if cpu != "" {
		res[v1.ResourceCPU] = resource.MustParse(cpu)
	}
	if memory != "" {
		res[v1.ResourceMemory] = resource.MustParse(memory)
	}
	if ephemeralStorage != "" {
		res[v1.ResourceEphemeralStorage] = resource.MustParse(ephemeralStorage)
	}
	return res
}

// newLimitRange returns a limit range with specified data
func newLimitRange(name string, limitType v1.LimitType,
	min, max,
	defaultLimit, defaultRequest,
	maxLimitRequestRatio v1.ResourceList) *v1.LimitRange {
	return &v1.LimitRange{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
		},
		Spec: v1.LimitRangeSpec{
			Limits: []v1.LimitRangeItem{
				{
					Type:                 limitType,
					Min:                  min,
					Max:                  max,
					Default:              defaultLimit,
					DefaultRequest:       defaultRequest,
					MaxLimitRequestRatio: maxLimitRequestRatio,
				},
			},
		},
	}
}
<|endoftext|>