<commit_before>// Package protocol defines the communication between the components
// of the Kite infrastructure. It defines some constants and structures
// designed to be sent between those components.
package protocol

import (
	"encoding/json"
	"errors"
	"strings"

	"github.com/koding/kite/dnode"
)

// Kite is the base struct containing the public fields. It is usually embeded
// in other structs, including the db model. The access model is in the form:
// username.environment.name.version.region.hostname.id
type Kite struct {
	// Short name identifying the type of the kite. Example: fs, terminal...
	Name string `json:"name"`

	// Owner of the Kite
	Username string `json:"username"`

	// Every Kite instance has different identifier.
	// If a kite is restarted, it's id will change.
	// This is generated on the Kite.
	ID string `json:"id"`

	// Environment is defines as something like "production", "testing",
	// "staging" or whatever.  This allows you to differentiate between a
	// cluster of kites.
	Environment string `json:"environment"`

	// Region of the kite it is running. Like "Europe", "Asia" or some other
	// locations.
	Region string `json:"region"`

	// 3-digit semantic version.
	Version string `json:"version"`

	// os.Hostname() of the Kite.
	Hostname string `json:"hostname"`
}

func (k Kite) String() string {
	return "/" + k.Username + "/" + k.Environment + "/" + k.Name + "/" + k.Version + "/" + k.Region + "/" + k.Hostname + "/" + k.ID
}

func (k *Kite) Validate() error {
	s := k.String()
	if strings.Contains(s, "//") {
		return errors.New("empty field")
	}
	if strings.Count(s, "/") != 7 {
		return errors.New(`fields cannot contain "/"`)
	}
	return nil
}

// RegisterArgs is used as the function argument to the Kontrol's register
// method.
type RegisterArgs struct {
	URL string `json:"url"`
}

// RegisterResult is a response to Register request from Kite to Kontrol.
type RegisterResult struct {
	URL string `json:"url"`
}

type GetKitesArgs struct {
	Query         KontrolQuery    `json:"query"`
	WatchCallback dnode.Callback  `json:"watchCallback,omitempty"`
	Who           json.RawMessage `json:"who,omitempty"`
}

type GetKitesResult struct {
	Kites     []*KiteWithToken `json:"kites"`
	WatcherID string           `json:"watcherID,omitempty"`
}

type KiteWithToken struct {
	Kite  Kite   `json:"kite"`
	URL   string `json:"url"`
	Token string `json:"token"`
}

// KiteEvent is the struct that is sent as an argument in watchCallback of
// getKites method of Kontrol.
type KiteEvent struct {
	Action KiteAction `json:"action"`
	Kite   Kite       `json:"kite"`

	// Required to connect when Action == Register
	URL   string `json:"url,omitempty"`
	Token string `json:"token,omitempty"`
}

type KiteAction string

const (
	Register   KiteAction = "REGISTER"
	Deregister KiteAction = "DEREGISTER"
)

// KontrolQuery is a structure of message sent to Kontrol. It is used for
// querying kites based on the incoming field parameters. Missing fields are
// not counted during the query (for example if the "version" field is empty,
// any kite with different version is going to be matched).
// Order of the fields is from general to specific.
type KontrolQuery struct {
	Username    string `json:"username"`
	Environment string `json:"environment"`
	Name        string `json:"name"`
	Version     string `json:"version"`
	Region      string `json:"region"`
	Hostname    string `json:"hostname"`
	ID          string `json:"id"`
}
<commit_msg>add WhoResult struct<commit_after>// Package protocol defines the communication between the components
// of the Kite infrastructure. It defines some constants and structures
// designed to be sent between those components.
package protocol

import (
	"encoding/json"
	"errors"
	"strings"

	"github.com/koding/kite/dnode"
)

// Kite is the base struct containing the public fields. It is usually embeded
// in other structs, including the db model. The access model is in the form:
// username.environment.name.version.region.hostname.id
type Kite struct {
	// Short name identifying the type of the kite. Example: fs, terminal...
	Name string `json:"name"`

	// Owner of the Kite
	Username string `json:"username"`

	// Every Kite instance has different identifier.
	// If a kite is restarted, it's id will change.
	// This is generated on the Kite.
	ID string `json:"id"`

	// Environment is defines as something like "production", "testing",
	// "staging" or whatever.  This allows you to differentiate between a
	// cluster of kites.
	Environment string `json:"environment"`

	// Region of the kite it is running. Like "Europe", "Asia" or some other
	// locations.
	Region string `json:"region"`

	// 3-digit semantic version.
	Version string `json:"version"`

	// os.Hostname() of the Kite.
	Hostname string `json:"hostname"`
}

func (k Kite) String() string {
	return "/" + k.Username + "/" + k.Environment + "/" + k.Name + "/" + k.Version + "/" + k.Region + "/" + k.Hostname + "/" + k.ID
}

func (k *Kite) Validate() error {
	s := k.String()
	if strings.Contains(s, "//") {
		return errors.New("empty field")
	}
	if strings.Count(s, "/") != 7 {
		return errors.New(`fields cannot contain "/"`)
	}
	return nil
}

// RegisterArgs is used as the function argument to the Kontrol's register
// method.
type RegisterArgs struct {
	URL string `json:"url"`
}

// RegisterResult is a response to Register request from Kite to Kontrol.
type RegisterResult struct {
	URL string `json:"url"`
}

type GetKitesArgs struct {
	Query         KontrolQuery    `json:"query"`
	WatchCallback dnode.Callback  `json:"watchCallback,omitempty"`
	Who           json.RawMessage `json:"who,omitempty"`
}

type WhoResult struct {
	Query KontrolQuery `json:"query"`
}

type GetKitesResult struct {
	Kites     []*KiteWithToken `json:"kites"`
	WatcherID string           `json:"watcherID,omitempty"`
}

type KiteWithToken struct {
	Kite  Kite   `json:"kite"`
	URL   string `json:"url"`
	Token string `json:"token"`
}

// KiteEvent is the struct that is sent as an argument in watchCallback of
// getKites method of Kontrol.
type KiteEvent struct {
	Action KiteAction `json:"action"`
	Kite   Kite       `json:"kite"`

	// Required to connect when Action == Register
	URL   string `json:"url,omitempty"`
	Token string `json:"token,omitempty"`
}

type KiteAction string

const (
	Register   KiteAction = "REGISTER"
	Deregister KiteAction = "DEREGISTER"
)

// KontrolQuery is a structure of message sent to Kontrol. It is used for
// querying kites based on the incoming field parameters. Missing fields are
// not counted during the query (for example if the "version" field is empty,
// any kite with different version is going to be matched).
// Order of the fields is from general to specific.
type KontrolQuery struct {
	Username    string `json:"username"`
	Environment string `json:"environment"`
	Name        string `json:"name"`
	Version     string `json:"version"`
	Region      string `json:"region"`
	Hostname    string `json:"hostname"`
	ID          string `json:"id"`
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"os"

	"github.com/lib/pq"
	"github.com/mozilla/tls-observatory/certificate"
	"github.com/mozilla/tls-observatory/database"
)

func main() {
	db, err := database.RegisterConnection(
		os.Getenv("TLSOBS_POSTGRESDB"),
		os.Getenv("TLSOBS_POSTGRESUSER"),
		os.Getenv("TLSOBS_POSTGRESPASS"),
		os.Getenv("TLSOBS_POSTGRES"),
		"require")
	defer db.Close()
	if err != nil {
		panic(err)
	}
	batch := 0
	for {
		fmt.Printf("\nProcessing batch %d to %d: ", batch, batch+100)
		rows, err := db.Query(`SELECT id, raw_cert
					FROM certificates
					ORDER BY id ASC
					LIMIT 100`)
		if rows != nil {
			defer rows.Close()
		}
		if err != nil {
			panic(fmt.Errorf("Error while retrieving certs: '%v'", err))
		}
		i := 0
		for rows.Next() {
			i++
			var raw string
			var id int64
			err = rows.Scan(&id, &raw)
			if err != nil {
				fmt.Println("error while parsing cert", id, ":", err)
				continue
			}
			certdata, err := base64.StdEncoding.DecodeString(raw)
			if err != nil {
				fmt.Println("error decoding base64 of cert", id, ":", err)
				continue
			}
			c, err := x509.ParseCertificate(certdata)
			if err != nil {
				fmt.Println("error while x509 parsing cert", id, ":", err)
				continue
			}
			var valInfo certificate.ValidationInfo
			cert := certificate.CertToStored(c, "", "", "", "", &valInfo)
			policies, err := json.Marshal(cert.X509v3Extensions.PolicyIdentifiers)
			if err != nil {
				log.Printf("error while marshalling policies for cert %d: %v", id, err)
				continue
			}
			if err != nil {
				log.Printf("error while marshalling permitted names for cert %d: %v", id, err)
				continue
			}
			if cert.X509v3Extensions.PermittedDNSDomains == nil {
				cert.X509v3Extensions.PermittedDNSDomains = make([]string, 0)
			}
			if cert.X509v3Extensions.ExcludedDNSDomains == nil {
				cert.X509v3Extensions.ExcludedDNSDomains = make([]string, 0)
			}
			log.Printf("id=%d, subject=%s, policies=%s, is_technically_constrained=%t", id, cert.Subject.String(),
				policies, cert.X509v3Extensions.IsTechnicallyConstrained)
			_, err = db.Exec(`UPDATE certificates
						SET x509_certificatepolicies=$1,
						    permitted_dns_domains=$2,
						    permitted_ip_addresses=$3,
						    excluded_dns_domains=$4,
						    excluded_ip_addresses=$5,
						    is_technically_constrained=$6
						WHERE id=$7`,
				policies,
				pq.Array(&cert.X509v3Extensions.PermittedDNSDomains),
				pq.Array(&cert.X509v3Extensions.PermittedIPAddresses),
				pq.Array(&cert.X509v3Extensions.ExcludedDNSDomains),
				pq.Array(&cert.X509v3Extensions.ExcludedIPAddresses),
				cert.X509v3Extensions.IsTechnicallyConstrained,
				id)
			if err != nil {
				fmt.Println("error while updating cert", id, "in database:", err)
			}
			fmt.Printf(".")
		}
		if i == 0 {
			fmt.Println("done!")
			break
		}
		//offset += limit
		batch += 100
	}
}
<commit_msg>Improve fixExtensions tool<commit_after>package main

import (
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/lib/pq"
	"github.com/mozilla/tls-observatory/certificate"
	"github.com/mozilla/tls-observatory/database"
)

func main() {
	db, err := database.RegisterConnection(
		os.Getenv("TLSOBS_POSTGRESDB"),
		os.Getenv("TLSOBS_POSTGRESUSER"),
		os.Getenv("TLSOBS_POSTGRESPASS"),
		os.Getenv("TLSOBS_POSTGRES"),
		"require")
	defer db.Close()
	if err != nil {
		panic(err)
	}
	batch := 0
	var lastid int64
	for {
		log.Printf("Processing batch %d to %d", batch, batch+100)
		rows, err := db.Query(`SELECT id, raw_cert
					FROM certificates
					ORDER BY id ASC
					OFFSET $1 LIMIT 100`, batch)
		if rows != nil {
			defer rows.Close()
		}
		if err != nil {
			panic(fmt.Errorf("Error while retrieving certs: '%v'", err))
		}
		i := 0
		for rows.Next() {
			var raw string
			var id int64
			err = rows.Scan(&id, &raw)
			if err != nil {
				fmt.Println("error while parsing cert", id, ":", err)
				continue
			}

			if id == lastid {
				// We're processing an ID that was already processed earlier, which
				// means we're looping over the end rows. it's time to exit
				goto done
			}
			lastid = id
			i++

			certdata, err := base64.StdEncoding.DecodeString(raw)
			if err != nil {
				fmt.Println("error decoding base64 of cert", id, ":", err)
				continue
			}
			c, err := x509.ParseCertificate(certdata)
			if err != nil {
				fmt.Println("error while x509 parsing cert", id, ":", err)
				continue
			}
			var valInfo certificate.ValidationInfo
			cert := certificate.CertToStored(c, "", "", "", "", &valInfo)
			if err != nil {
				log.Printf("error while marshalling permitted names for cert %d: %v", id, err)
				continue
			}
			if cert.X509v3Extensions.PermittedDNSDomains == nil {
				cert.X509v3Extensions.PermittedDNSDomains = make([]string, 0)
			}
			if cert.X509v3Extensions.ExcludedDNSDomains == nil {
				cert.X509v3Extensions.ExcludedDNSDomains = make([]string, 0)
			}
			if cert.X509v3Extensions.IsTechnicallyConstrained {
				log.Printf("id=%d, permitted_dns_domains=%v, permitted_ip_addresses=%v, excluded_dns_domains=%v, excluded_ip_addresses=%v, is_technically_constrained=%t",
					id, cert.X509v3Extensions.PermittedDNSDomains, cert.X509v3Extensions.PermittedIPAddresses, cert.X509v3Extensions.ExcludedDNSDomains, cert.X509v3Extensions.ExcludedIPAddresses, cert.X509v3Extensions.IsTechnicallyConstrained)
				_, err = db.Exec(`UPDATE certificates
						SET permitted_dns_domains=$1,
						    permitted_ip_addresses=$2,
						    excluded_dns_domains=$3,
						    excluded_ip_addresses=$4,
						    is_technically_constrained=$5
						WHERE id=$6`,
					pq.Array(&cert.X509v3Extensions.PermittedDNSDomains),
					pq.Array(&cert.X509v3Extensions.PermittedIPAddresses),
					pq.Array(&cert.X509v3Extensions.ExcludedDNSDomains),
					pq.Array(&cert.X509v3Extensions.ExcludedIPAddresses),
					cert.X509v3Extensions.IsTechnicallyConstrained,
					id)
				if err != nil {
					fmt.Println("error while updating cert", id, "in database:", err)
				}
			}
			if strings.HasPrefix(cert.Serial, "-") {
				log.Printf("id=%d, serial=%s", id, cert.Serial)
				_, err = db.Exec(`UPDATE certificates
						SET serial_number=$1
						WHERE id=$2`,
					cert.Serial, id)
				if err != nil {
					fmt.Println("error while updating cert", id, "in database:", err)
				}
			}
		}
		if i == 0 {
			goto done
		}
		batch += 100
	}
done:
	fmt.Println("Processing done. Goodbye!")
}
<|endoftext|>


<commit_before>package proxy

import (
	"bytes"
	"strings"

	"github.com/juju/errors"
	log "github.com/ngaut/logging"
	"github.com/wandoulabs/cm/hack"
	"github.com/wandoulabs/cm/mysql"
	"github.com/wandoulabs/cm/sqlparser"
)

func (c *Conn) handleSimpleSelect(sql string, stmt *sqlparser.SimpleSelect) error {
	//todo: handle multi set statement like: set x=1, y=2
	/*
		if len(stmt.SelectExprs) != 1 {
			return errors.Errorf("support select one informaction function, %s", sql)
		}
	*/
	log.Debug(sql)
	expr, ok := stmt.SelectExprs[0].(*sqlparser.NonStarExpr)
	if !ok {
		return errors.Errorf("support select informaction function, %s", sql)
	}

	var funcExpr *sqlparser.FuncExpr
	var specialColumn *sqlparser.ColName

	switch v := expr.Expr.(type) {
	case *sqlparser.FuncExpr:
		funcExpr = v
	case *sqlparser.ColName:
		specialColumn = v
		log.Debug(string(specialColumn.Name))
	case sqlparser.NumVal: //select 1
		return errors.Trace(c.handleShow(stmt, sql, nil))
	case *sqlparser.BinaryExpr: //select 2 * 3
		return errors.Trace(c.handleShow(stmt, sql, nil))
	default:
		return errors.Errorf("support select informaction function, %s, %T", sql, v)
	}

	var r *mysql.Resultset
	var err error

	if funcExpr != nil {
		switch strings.ToLower(string(funcExpr.Name)) {
		case "last_insert_id":
			r, err = c.buildSimpleSelectResult(c.lastInsertId, funcExpr.Name, expr.As)
		case "row_count":
			r, err = c.buildSimpleSelectResult(c.affectedRows, funcExpr.Name, expr.As)
		case "version":
			r, err = c.buildSimpleSelectResult(mysql.ServerVersion, funcExpr.Name, expr.As)
		case "connection_id":
			r, err = c.buildSimpleSelectResult(c.connectionId, funcExpr.Name, expr.As)
		case "database":
			if len(c.db) > 0 {
				r, err = c.buildSimpleSelectResult(c.db, funcExpr.Name, expr.As)
			} else {
				r, err = c.buildSimpleSelectResult("NULL", funcExpr.Name, expr.As)
			}
		case "user": // select USER()
			r, err = c.buildSimpleSelectResult(c.user, funcExpr.Name, expr.As)
		default:
			log.Warning(c.connectionId, sql)
			//todo: more strict
			return errors.Trace(c.handleShow(stmt, sql, nil))
			//return errors.Errorf("function %s not support, %+v", funcExpr.Name, funcExpr)
		}

		if err != nil {
			return errors.Trace(err)
		}
	} else {
		return errors.Trace(c.handleShow(stmt, sql, nil))
	}

	return errors.Trace(c.writeResultset(c.status, r))
}

func (c *Conn) buildSimpleSelectResult(value interface{}, name []byte, asName []byte) (*mysql.Resultset, error) {
	field := &mysql.Field{Name: name, OrgName: name}
	if asName != nil {
		field.Name = asName
	}

	formatField(field, value)

	r := &mysql.Resultset{Fields: []*mysql.Field{field}}
	row := mysql.Raw(byte(field.Type), value, false)
	r.RowDatas = append(r.RowDatas, mysql.PutLengthEncodedString(row, c.alloc))

	return r, nil
}

func (c *Conn) getShardIds(table string) ([]string, error) {
	r := c.schema().r
	rule := r.GetRule(table)
	var shardIds []string
	if rule == nil {
		//using default rules
		if len(r.Default) != 0 {
			shardIds = r.Default
		} else {
			return nil, errors.Errorf("no rule for table %s, %+v, please check config file", table, c.schema)
		}
	} else {
		shardIds = rule.MapToShards
	}

	return shardIds, nil
}

func (c *Conn) handleFieldList(data []byte) error {
	index := bytes.IndexByte(data, 0x00)
	table := hack.String(data[0:index])
	wildcard := hack.String(data[index+1:])
	shardIds, err := c.getShardIds(table)
	if err != nil {
		return err
	}

	//todo: pass through
	if len(shardIds) == 0 {
		return errors.Errorf("no rule for table %s, %+v, please check config file", table, c.schema)
	}

	//hard code, assume all of the shard has the same schema
	n := c.server.GetShard(shardIds[0])
	if n == nil {
		return errors.Errorf("shard %s not found, %+v", shardIds, c.schema)
	}

	co, err := n.getMasterConn()
	if err != nil {
		return errors.Trace(err)
	}
	defer co.Close()

	if err = co.UseDB(c.db); err != nil {
		return errors.Trace(err)
	}

	if fs, err := co.FieldList(table, wildcard); err != nil {
		return errors.Trace(err)
	} else {
		return errors.Trace(c.writeFieldList(c.status, fs))
	}
}

func (c *Conn) writeFieldList(status uint16, fs []*mysql.Field) error {
	c.affectedRows = int64(-1)

	data := make([]byte, 4, 1024)

	for _, v := range fs {
		data = data[0:4]
		data = append(data, v.Dump(c.alloc)...)
		if err := c.writePacket(data); err != nil {
			return errors.Trace(err)
		}
	}

	err := c.writeEOF(status)
	if err != nil {
		return errors.Trace(err)
	}

	return errors.Trace(c.flush())
}
<commit_msg>fix getShardIds<commit_after>package proxy

import (
	"bytes"
	"strings"

	"github.com/juju/errors"
	log "github.com/ngaut/logging"
	"github.com/wandoulabs/cm/hack"
	"github.com/wandoulabs/cm/mysql"
	"github.com/wandoulabs/cm/sqlparser"
)

func (c *Conn) handleSimpleSelect(sql string, stmt *sqlparser.SimpleSelect) error {
	//todo: handle multi set statement like: set x=1, y=2
	/*
		if len(stmt.SelectExprs) != 1 {
			return errors.Errorf("support select one informaction function, %s", sql)
		}
	*/
	log.Debug(sql)
	expr, ok := stmt.SelectExprs[0].(*sqlparser.NonStarExpr)
	if !ok {
		return errors.Errorf("support select informaction function, %s", sql)
	}

	var funcExpr *sqlparser.FuncExpr
	var specialColumn *sqlparser.ColName

	switch v := expr.Expr.(type) {
	case *sqlparser.FuncExpr:
		funcExpr = v
	case *sqlparser.ColName:
		specialColumn = v
		log.Debug(string(specialColumn.Name))
	case sqlparser.NumVal: //select 1
		return errors.Trace(c.handleShow(stmt, sql, nil))
	case *sqlparser.BinaryExpr: //select 2 * 3
		return errors.Trace(c.handleShow(stmt, sql, nil))
	default:
		return errors.Errorf("support select informaction function, %s, %T", sql, v)
	}

	var r *mysql.Resultset
	var err error

	if funcExpr != nil {
		switch strings.ToLower(string(funcExpr.Name)) {
		case "last_insert_id":
			r, err = c.buildSimpleSelectResult(c.lastInsertId, funcExpr.Name, expr.As)
		case "row_count":
			r, err = c.buildSimpleSelectResult(c.affectedRows, funcExpr.Name, expr.As)
		case "version":
			r, err = c.buildSimpleSelectResult(mysql.ServerVersion, funcExpr.Name, expr.As)
		case "connection_id":
			r, err = c.buildSimpleSelectResult(c.connectionId, funcExpr.Name, expr.As)
		case "database":
			if len(c.db) > 0 {
				r, err = c.buildSimpleSelectResult(c.db, funcExpr.Name, expr.As)
			} else {
				r, err = c.buildSimpleSelectResult("NULL", funcExpr.Name, expr.As)
			}
		case "user": // select USER()
			r, err = c.buildSimpleSelectResult(c.user, funcExpr.Name, expr.As)
		default:
			log.Warning(c.connectionId, sql)
			//todo: more strict
			return errors.Trace(c.handleShow(stmt, sql, nil))
			//return errors.Errorf("function %s not support, %+v", funcExpr.Name, funcExpr)
		}

		if err != nil {
			return errors.Trace(err)
		}
	} else {
		return errors.Trace(c.handleShow(stmt, sql, nil))
	}

	return errors.Trace(c.writeResultset(c.status, r))
}

func (c *Conn) buildSimpleSelectResult(value interface{}, name []byte, asName []byte) (*mysql.Resultset, error) {
	field := &mysql.Field{Name: name, OrgName: name}
	if asName != nil {
		field.Name = asName
	}

	formatField(field, value)

	r := &mysql.Resultset{Fields: []*mysql.Field{field}}
	row := mysql.Raw(byte(field.Type), value, false)
	r.RowDatas = append(r.RowDatas, mysql.PutLengthEncodedString(row, c.alloc))

	return r, nil
}

func (c *Conn) getShardIds(table string) ([]string, error) {
	r := c.schema().r
	rule := r.GetRule(table)
	var shardIds []string
	if rule == nil {
		return nil, errors.Errorf("no rule for table %s, %+v, please check config file", table, c.schema)
	}

	shardIds = rule.MapToShards
	if len(shardIds) == 0 {
		//using default rules
		shardIds = r.Default
	}

	return shardIds, nil
}

func (c *Conn) handleFieldList(data []byte) error {
	index := bytes.IndexByte(data, 0x00)
	table := hack.String(data[0:index])
	wildcard := hack.String(data[index+1:])
	shardIds, err := c.getShardIds(table)
	if err != nil {
		return err
	}

	//todo: pass through
	if len(shardIds) == 0 {
		return errors.Errorf("no rule for table %s, %+v, please check config file", table, c.schema)
	}

	//hard code, assume all of the shard has the same schema
	n := c.server.GetShard(shardIds[0])
	if n == nil {
		return errors.Errorf("shard %s not found, %+v", shardIds, c.schema)
	}

	co, err := n.getMasterConn()
	if err != nil {
		return errors.Trace(err)
	}
	defer co.Close()

	if err = co.UseDB(c.db); err != nil {
		return errors.Trace(err)
	}

	if fs, err := co.FieldList(table, wildcard); err != nil {
		return errors.Trace(err)
	} else {
		return errors.Trace(c.writeFieldList(c.status, fs))
	}
}

func (c *Conn) writeFieldList(status uint16, fs []*mysql.Field) error {
	c.affectedRows = int64(-1)

	data := make([]byte, 4, 1024)

	for _, v := range fs {
		data = data[0:4]
		data = append(data, v.Dump(c.alloc)...)
		if err := c.writePacket(data); err != nil {
			return errors.Trace(err)
		}
	}

	err := c.writeEOF(status)
	if err != nil {
		return errors.Trace(err)
	}

	return errors.Trace(c.flush())
}
<|endoftext|>


<commit_before>/* -------------------------------------------------------------------------- */
/* Copyright 2002-2019, OpenNebula Project, OpenNebula Systems                */
/*                                                                            */
/* 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 goca

import (
	"testing"

	. "gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { TestingT(t) }

type VMSuite struct {
	templateID uint
	vmID       uint
	hostID     uint
}

var _ = Suite(&VMSuite{})

func (s *VMSuite) SetUpSuite(c *C) {
	// Create template
	tpl := NewTemplateBuilder()

	tpl.AddValue("NAME", GenName("VMSuite-template"))
	tpl.AddValue("CPU", 1)
	tpl.AddValue("MEMORY", "64")

	templateID, err := CreateTemplate(tpl.String())
	c.Assert(err, IsNil)

	s.templateID = templateID

	s.hostID, _ = CreateHost("dummy-test", "dummy", "dummy", 0)

	tmpl := "TM_MAD=dummy\nDS_MAD=dummy"

	ds_img := NewDatastore(1)
	ds_img.Update(tmpl, 1)

	ds_sys := NewDatastore(0)
	ds_sys.Update(tmpl, 1)

}

func (s *VMSuite) SetUpTest(c *C) {
	template := NewTemplate(s.templateID)

	vmID, err := template.Instantiate("", true, "", false)
	c.Assert(err, IsNil)
	s.vmID = vmID
}

func (s *VMSuite) TearDownTest(c *C) {
	vm := NewVM(s.vmID)

	err := vm.TerminateHard()
	if err != nil {
		err = vm.RecoverDelete()
	}
	c.Assert(err, IsNil)
}

func (s *VMSuite) TearDownSuite(c *C) {
	template := NewTemplate(s.templateID)
	template.Delete()

	host := NewHost(s.hostID)
	host.Delete()
}

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

func VMExpectState(c *C, vm *VM, state, lcmState string) func() bool {
	return func() bool {
		vm.Info()

		s, l, err := vm.StateString()
		if err != nil {
			return false
		}

		if lcmState != "" && l == lcmState {
			return true
		}

		if state != "" && s == state && lcmState == ""{
			return true
		}

		c.Logf("VM: %d. Expecting: %s/%s, Got: %s/%s", vm.ID, state, lcmState, s, l)

		return false
	}
}

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

func (s *VMSuite) TestVMDeploy(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)
}

func (s *VMSuite) TestVMHoldRelease(c *C) {
	vm := NewVM(s.vmID)
	c.Assert(WaitResource(VMExpectState(c, vm, "HOLD", "")), Equals, true)

	err := vm.Release()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "PENDING", "")), Equals, true)
}

func (s *VMSuite) TestVMUpdate(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Update("A=B", 1)
	c.Assert(err, IsNil)

	err = vm.Info()
	c.Assert(err, IsNil)

	val := vm.UserTemplate.Dynamic.GetContentByName("A")
	c.Assert(val, Equals, "B")
}

// TODO: Hosts
// func (s *VMSuite) TestVMMigrate(c *C) {
// }
// func (s *VMSuite) TestVMLiveMigrate(c *C) {
// }

func (s *VMSuite) TestVMTerminate(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Terminate()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "DONE", "")), Equals, true)
}

func (s *VMSuite) TestVMTerminateHard(c *C) {
	vm := NewVM(s.vmID)
	err :=vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	vm.TerminateHard()
	c.Assert(WaitResource(VMExpectState(c, vm, "DONE", "")), Equals, true)
}

func (s *VMSuite) TestVMStop(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Stop()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "STOPPED", "")), Equals, true)
}

func (s *VMSuite) TestVMSuspend(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Suspend()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "SUSPENDED", "")), Equals, true)
}

func (s *VMSuite) TestVMResume(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Suspend()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "SUSPENDED", "")), Equals, true)

	err = vm.Resume()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)
}

func (s *VMSuite) TestVMResize(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Poweroff()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "POWEROFF", "")), Equals, true)

	err = vm.Resize("CPU=2.5\nMEMORY=512", false)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "POWEROFF", "")), Equals, true)

	err = vm.Resume()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Info()
	c.Assert(err, IsNil)

	c.Assert(vm.Template.CPU, Equals, 2.5)

	c.Assert(vm.Template.Memory, Equals, 512)
}
<commit_msg>M #-: Fix minor bug in GOCA tests (#3486)<commit_after>/* -------------------------------------------------------------------------- */
/* Copyright 2002-2019, OpenNebula Project, OpenNebula Systems                */
/*                                                                            */
/* 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 goca

import (
	"testing"

	. "gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { TestingT(t) }

type VMSuite struct {
	templateID uint
	vmID       uint
	hostID     uint
}

var _ = Suite(&VMSuite{})

func (s *VMSuite) SetUpSuite(c *C) {
	// Create template
	tpl := NewTemplateBuilder()

	tpl.AddValue("NAME", GenName("VMSuite-template"))
	tpl.AddValue("CPU", 1)
	tpl.AddValue("MEMORY", "64")

	templateID, err := CreateTemplate(tpl.String())
	c.Assert(err, IsNil)

	s.templateID = templateID

	s.hostID, _ = CreateHost("dummy-test", "dummy", "dummy", 0)

	tmpl := "TM_MAD=dummy\nDS_MAD=dummy"

	ds_img := NewDatastore(1)
	ds_img.Update(tmpl, 1)

	ds_sys := NewDatastore(0)
	ds_sys.Update(tmpl, 1)

}

func (s *VMSuite) SetUpTest(c *C) {
	template := NewTemplate(s.templateID)

	vmID, err := template.Instantiate("", true, "", false)
	c.Assert(err, IsNil)
	s.vmID = vmID
}

func (s *VMSuite) TearDownTest(c *C) {
	vm := NewVM(s.vmID)

	err := vm.TerminateHard()
	if err != nil {
		err = vm.RecoverDelete()
	}
	c.Assert(err, IsNil)
}

func (s *VMSuite) TearDownSuite(c *C) {
	template := NewTemplate(s.templateID)
	template.Delete()

	host := NewHost(s.hostID)
	host.Delete()
}

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

func VMExpectState(c *C, vm *VM, state, lcmState string) func() bool {
	return func() bool {
		vm.Info()

		s, l, err := vm.StateString()
		if err != nil {
			return false
		}

		if lcmState != "" && l == lcmState {
			return true
		}

		if state != "" && s == state && lcmState == ""{
			return true
		}

		c.Logf("VM: %d. Expecting: %s/%s, Got: %s/%s", vm.ID, state, lcmState, s, l)

		return false
	}
}

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

func (s *VMSuite) TestVMDeploy(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)
}

func (s *VMSuite) TestVMHoldRelease(c *C) {
	vm := NewVM(s.vmID)
	c.Assert(WaitResource(VMExpectState(c, vm, "HOLD", "")), Equals, true)

	err := vm.Release()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "PENDING", "")), Equals, true)
}

func (s *VMSuite) TestVMUpdate(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Update("A=B", 1)
	c.Assert(err, IsNil)

	err = vm.Info()
	c.Assert(err, IsNil)

	val := vm.UserTemplate.Dynamic.GetContentByName("A")
	c.Assert(val, Equals, "B")
}

// TODO: Hosts
// func (s *VMSuite) TestVMMigrate(c *C) {
// }
// func (s *VMSuite) TestVMLiveMigrate(c *C) {
// }

func (s *VMSuite) TestVMTerminate(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Terminate()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "DONE", "")), Equals, true)
}

func (s *VMSuite) TestVMTerminateHard(c *C) {
	vm := NewVM(s.vmID)
	err :=vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	vm.TerminateHard()
	c.Assert(WaitResource(VMExpectState(c, vm, "DONE", "")), Equals, true)
}

func (s *VMSuite) TestVMStop(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Stop()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "STOPPED", "")), Equals, true)
}

func (s *VMSuite) TestVMSuspend(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Suspend()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "SUSPENDED", "")), Equals, true)
}

func (s *VMSuite) TestVMResume(c *C) {
	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Suspend()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "SUSPENDED", "")), Equals, true)

	err = vm.Resume()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)
}

func (s *VMSuite) TestVMResize(c *C) {
	expectHostState := func(host *Host) func() bool {
		return func() bool {
			host.Info()
			state, _ := host.StateString()

			return state == "MONITORED"
		}
	}

	vm := NewVM(s.vmID)
	err := vm.Deploy(s.hostID, false, -1)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Poweroff()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "POWEROFF", "")), Equals, true)

	c.Assert(WaitResource(expectHostState(NewHost(s.hostID))), Equals, true)

	err = vm.Resize("CPU=2.5\nMEMORY=512", false)
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "POWEROFF", "")), Equals, true)

	err = vm.Resume()
	c.Assert(err, IsNil)
	c.Assert(WaitResource(VMExpectState(c, vm, "ACTIVE", "RUNNING")), Equals, true)

	err = vm.Info()
	c.Assert(err, IsNil)

	c.Assert(vm.Template.CPU, Equals, 2.5)

	c.Assert(vm.Template.Memory, Equals, 512)
}
<|endoftext|>


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

// Tests for transport.go

package http_test

import (
	"fmt"
	. "http"
	"http/httptest"
	"io/ioutil"
	"os"
	"testing"
	"time"
)

// TODO: test 5 pipelined requests with responses: 1) OK, 2) OK, Connection: Close
//       and then verify that the final 2 responses get errors back.

// hostPortHandler writes back the client's "host:port".
var hostPortHandler = HandlerFunc(func(w ResponseWriter, r *Request) {
	if r.FormValue("close") == "true" {
		w.Header().Set("Connection", "close")
	}
	fmt.Fprintf(w, "%s", r.RemoteAddr)
})

// Two subsequent requests and verify their response is the same.
// The response from the server is our own IP:port
func TestTransportKeepAlives(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	for _, disableKeepAlive := range []bool{false, true} {
		tr := &Transport{DisableKeepAlives: disableKeepAlive}
		c := &Client{Transport: tr}

		fetch := func(n int) string {
			res, _, err := c.Get(ts.URL)
			if err != nil {
				t.Fatalf("error in disableKeepAlive=%v, req #%d, GET: %v", disableKeepAlive, n, err)
			}
			body, err := ioutil.ReadAll(res.Body)
			if err != nil {
				t.Fatalf("error in disableKeepAlive=%v, req #%d, ReadAll: %v", disableKeepAlive, n, err)
			}
			return string(body)
		}

		body1 := fetch(1)
		body2 := fetch(2)

		bodiesDiffer := body1 != body2
		if bodiesDiffer != disableKeepAlive {
			t.Errorf("error in disableKeepAlive=%v. unexpected bodiesDiffer=%v; body1=%q; body2=%q",
				disableKeepAlive, bodiesDiffer, body1, body2)
		}
	}
}

func TestTransportConnectionCloseOnResponse(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	for _, connectionClose := range []bool{false, true} {
		tr := &Transport{}
		c := &Client{Transport: tr}

		fetch := func(n int) string {
			req := new(Request)
			var err os.Error
			req.URL, err = ParseURL(ts.URL + fmt.Sprintf("?close=%v", connectionClose))
			if err != nil {
				t.Fatalf("URL parse error: %v", err)
			}
			req.Method = "GET"
			req.Proto = "HTTP/1.1"
			req.ProtoMajor = 1
			req.ProtoMinor = 1

			res, err := c.Do(req)
			if err != nil {
				t.Fatalf("error in connectionClose=%v, req #%d, Do: %v", connectionClose, n, err)
			}
			body, err := ioutil.ReadAll(res.Body)
			defer res.Body.Close()
			if err != nil {
				t.Fatalf("error in connectionClose=%v, req #%d, ReadAll: %v", connectionClose, n, err)
			}
			return string(body)
		}

		body1 := fetch(1)
		body2 := fetch(2)
		bodiesDiffer := body1 != body2
		if bodiesDiffer != connectionClose {
			t.Errorf("error in connectionClose=%v. unexpected bodiesDiffer=%v; body1=%q; body2=%q",
				connectionClose, bodiesDiffer, body1, body2)
		}
	}
}

func TestTransportConnectionCloseOnRequest(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	for _, connectionClose := range []bool{false, true} {
		tr := &Transport{}
		c := &Client{Transport: tr}

		fetch := func(n int) string {
			req := new(Request)
			var err os.Error
			req.URL, err = ParseURL(ts.URL)
			if err != nil {
				t.Fatalf("URL parse error: %v", err)
			}
			req.Method = "GET"
			req.Proto = "HTTP/1.1"
			req.ProtoMajor = 1
			req.ProtoMinor = 1
			req.Close = connectionClose

			res, err := c.Do(req)
			if err != nil {
				t.Fatalf("error in connectionClose=%v, req #%d, Do: %v", connectionClose, n, err)
			}
			body, err := ioutil.ReadAll(res.Body)
			if err != nil {
				t.Fatalf("error in connectionClose=%v, req #%d, ReadAll: %v", connectionClose, n, err)
			}
			return string(body)
		}

		body1 := fetch(1)
		body2 := fetch(2)
		bodiesDiffer := body1 != body2
		if bodiesDiffer != connectionClose {
			t.Errorf("error in connectionClose=%v. unexpected bodiesDiffer=%v; body1=%q; body2=%q",
				connectionClose, bodiesDiffer, body1, body2)
		}
	}
}

func TestTransportIdleCacheKeys(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	tr := &Transport{DisableKeepAlives: false}
	c := &Client{Transport: tr}

	if e, g := 0, len(tr.IdleConnKeysForTesting()); e != g {
		t.Errorf("After CloseIdleConnections expected %d idle conn cache keys; got %d", e, g)
	}

	resp, _, err := c.Get(ts.URL)
	if err != nil {
		t.Error(err)
	}
	ioutil.ReadAll(resp.Body)

	keys := tr.IdleConnKeysForTesting()
	if e, g := 1, len(keys); e != g {
		t.Fatalf("After Get expected %d idle conn cache keys; got %d", e, g)
	}

	if e := "|http|" + ts.Listener.Addr().String(); keys[0] != e {
		t.Errorf("Expected idle cache key %q; got %q", e, keys[0])
	}

	tr.CloseIdleConnections()
	if e, g := 0, len(tr.IdleConnKeysForTesting()); e != g {
		t.Errorf("After CloseIdleConnections expected %d idle conn cache keys; got %d", e, g)
	}
}

func TestTransportMaxPerHostIdleConns(t *testing.T) {
	ch := make(chan string)
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		fmt.Fprintf(w, "%s", <-ch)
	}))
	defer ts.Close()
	maxIdleConns := 2
	tr := &Transport{DisableKeepAlives: false, MaxIdleConnsPerHost: maxIdleConns}
	c := &Client{Transport: tr}

	// Start 3 outstanding requests (will hang until we write to
	// ch)
	donech := make(chan bool)
	doReq := func() {
		resp, _, err := c.Get(ts.URL)
		if err != nil {
			t.Error(err)
		}
		ioutil.ReadAll(resp.Body)
		donech <- true
	}
	go doReq()
	go doReq()
	go doReq()

	if e, g := 0, len(tr.IdleConnKeysForTesting()); e != g {
		t.Fatalf("Before writes, expected %d idle conn cache keys; got %d", e, g)
	}

	ch <- "res1"
	<-donech
	keys := tr.IdleConnKeysForTesting()
	if e, g := 1, len(keys); e != g {
		t.Fatalf("after first response, expected %d idle conn cache keys; got %d", e, g)
	}
	cacheKey := "|http|" + ts.Listener.Addr().String()
	if keys[0] != cacheKey {
		t.Fatalf("Expected idle cache key %q; got %q", cacheKey, keys[0])
	}
	if e, g := 1, tr.IdleConnCountForTesting(cacheKey); e != g {
		t.Errorf("after first response, expected %d idle conns; got %d", e, g)
	}

	ch <- "res2"
	<-donech
	if e, g := 2, tr.IdleConnCountForTesting(cacheKey); e != g {
		t.Errorf("after second response, expected %d idle conns; got %d", e, g)
	}

	ch <- "res3"
	<-donech
	if e, g := maxIdleConns, tr.IdleConnCountForTesting(cacheKey); e != g {
		t.Errorf("after third response, still expected %d idle conns; got %d", e, g)
	}
}

func TestTransportServerClosingUnexpectedly(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	tr := &Transport{}
	c := &Client{Transport: tr}

	fetch := func(n int) string {
		res, _, err := c.Get(ts.URL)
		if err != nil {
			t.Fatalf("error in req #%d, GET: %v", n, err)
		}
		body, err := ioutil.ReadAll(res.Body)
		if err != nil {
			t.Fatalf("error in req #%d, ReadAll: %v", n, err)
		}
		res.Body.Close()
		return string(body)
	}

	body1 := fetch(1)
	body2 := fetch(2)

	ts.CloseClientConnections() // surprise!
	time.Sleep(25e6)            // idle for a bit (test is inherently racey, but expectedly)

	body3 := fetch(3)

	if body1 != body2 {
		t.Errorf("expected body1 and body2 to be equal")
	}
	if body2 == body3 {
		t.Errorf("expected body2 and body3 to be different")
	}
}

// TestTransportHeadResponses verifies that we deal with Content-Lengths
// with no bodies properly
func TestTransportHeadResponses(t *testing.T) {
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		if r.Method != "HEAD" {
			panic("expected HEAD; got " + r.Method)
		}
		w.Header().Set("Content-Length", "123")
		w.WriteHeader(200)
	}))
	defer ts.Close()

	tr := &Transport{DisableKeepAlives: false}
	c := &Client{Transport: tr}
	for i := 0; i < 2; i++ {
		res, err := c.Head(ts.URL)
		if err != nil {
			t.Errorf("error on loop %d: %v", i, err)
		}
		if e, g := "123", res.Header.Get("Content-Length"); e != g {
			t.Errorf("loop %d: expected Content-Length header of %q, got %q", e, g)
		}
		if e, g := int64(0), res.ContentLength; e != g {
			t.Errorf("loop %d: expected res.ContentLength of %v, got %v", e, g)
		}
	}
}

// TestTransportHeadChunkedResponse verifies that we ignore chunked transfer-encoding
// on responses to HEAD requests.
func TestTransportHeadChunkedResponse(t *testing.T) {
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		if r.Method != "HEAD" {
			panic("expected HEAD; got " + r.Method)
		}
		w.Header().Set("Transfer-Encoding", "chunked") // client should ignore
		w.Header().Set("x-client-ipport", r.RemoteAddr)
		w.WriteHeader(200)
	}))
	defer ts.Close()

	tr := &Transport{DisableKeepAlives: false}
	c := &Client{Transport: tr}

	res1, err := c.Head(ts.URL)
	if err != nil {
		t.Fatalf("request 1 error: %v", err)
	}
	res2, err := c.Head(ts.URL)
	if err != nil {
		t.Fatalf("request 2 error: %v", err)
	}
	if v1, v2 := res1.Header.Get("x-client-ipport"), res2.Header.Get("x-client-ipport"); v1 != v2 {
		t.Errorf("ip/ports differed between head requests: %q vs %q", v1, v2)
	}
}

func TestTransportNilURL(t *testing.T) {
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		fmt.Fprintf(w, "Hi")
	}))
	defer ts.Close()

	req := new(Request)
	req.URL = nil // what we're actually testing
	req.Method = "GET"
	req.RawURL = ts.URL
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1

	tr := &Transport{}
	res, err := tr.RoundTrip(req)
	if err != nil {
		t.Fatalf("unexpected RoundTrip error: %v", err)
	}
	body, err := ioutil.ReadAll(res.Body)
	if g, e := string(body), "Hi"; g != e {
		t.Fatalf("Expected response body of %q; got %q", e, g)
	}
}
<commit_msg>http: fix incorrect prints found by govet<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Tests for transport.go

package http_test

import (
	"fmt"
	. "http"
	"http/httptest"
	"io/ioutil"
	"os"
	"testing"
	"time"
)

// TODO: test 5 pipelined requests with responses: 1) OK, 2) OK, Connection: Close
//       and then verify that the final 2 responses get errors back.

// hostPortHandler writes back the client's "host:port".
var hostPortHandler = HandlerFunc(func(w ResponseWriter, r *Request) {
	if r.FormValue("close") == "true" {
		w.Header().Set("Connection", "close")
	}
	fmt.Fprintf(w, "%s", r.RemoteAddr)
})

// Two subsequent requests and verify their response is the same.
// The response from the server is our own IP:port
func TestTransportKeepAlives(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	for _, disableKeepAlive := range []bool{false, true} {
		tr := &Transport{DisableKeepAlives: disableKeepAlive}
		c := &Client{Transport: tr}

		fetch := func(n int) string {
			res, _, err := c.Get(ts.URL)
			if err != nil {
				t.Fatalf("error in disableKeepAlive=%v, req #%d, GET: %v", disableKeepAlive, n, err)
			}
			body, err := ioutil.ReadAll(res.Body)
			if err != nil {
				t.Fatalf("error in disableKeepAlive=%v, req #%d, ReadAll: %v", disableKeepAlive, n, err)
			}
			return string(body)
		}

		body1 := fetch(1)
		body2 := fetch(2)

		bodiesDiffer := body1 != body2
		if bodiesDiffer != disableKeepAlive {
			t.Errorf("error in disableKeepAlive=%v. unexpected bodiesDiffer=%v; body1=%q; body2=%q",
				disableKeepAlive, bodiesDiffer, body1, body2)
		}
	}
}

func TestTransportConnectionCloseOnResponse(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	for _, connectionClose := range []bool{false, true} {
		tr := &Transport{}
		c := &Client{Transport: tr}

		fetch := func(n int) string {
			req := new(Request)
			var err os.Error
			req.URL, err = ParseURL(ts.URL + fmt.Sprintf("?close=%v", connectionClose))
			if err != nil {
				t.Fatalf("URL parse error: %v", err)
			}
			req.Method = "GET"
			req.Proto = "HTTP/1.1"
			req.ProtoMajor = 1
			req.ProtoMinor = 1

			res, err := c.Do(req)
			if err != nil {
				t.Fatalf("error in connectionClose=%v, req #%d, Do: %v", connectionClose, n, err)
			}
			body, err := ioutil.ReadAll(res.Body)
			defer res.Body.Close()
			if err != nil {
				t.Fatalf("error in connectionClose=%v, req #%d, ReadAll: %v", connectionClose, n, err)
			}
			return string(body)
		}

		body1 := fetch(1)
		body2 := fetch(2)
		bodiesDiffer := body1 != body2
		if bodiesDiffer != connectionClose {
			t.Errorf("error in connectionClose=%v. unexpected bodiesDiffer=%v; body1=%q; body2=%q",
				connectionClose, bodiesDiffer, body1, body2)
		}
	}
}

func TestTransportConnectionCloseOnRequest(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	for _, connectionClose := range []bool{false, true} {
		tr := &Transport{}
		c := &Client{Transport: tr}

		fetch := func(n int) string {
			req := new(Request)
			var err os.Error
			req.URL, err = ParseURL(ts.URL)
			if err != nil {
				t.Fatalf("URL parse error: %v", err)
			}
			req.Method = "GET"
			req.Proto = "HTTP/1.1"
			req.ProtoMajor = 1
			req.ProtoMinor = 1
			req.Close = connectionClose

			res, err := c.Do(req)
			if err != nil {
				t.Fatalf("error in connectionClose=%v, req #%d, Do: %v", connectionClose, n, err)
			}
			body, err := ioutil.ReadAll(res.Body)
			if err != nil {
				t.Fatalf("error in connectionClose=%v, req #%d, ReadAll: %v", connectionClose, n, err)
			}
			return string(body)
		}

		body1 := fetch(1)
		body2 := fetch(2)
		bodiesDiffer := body1 != body2
		if bodiesDiffer != connectionClose {
			t.Errorf("error in connectionClose=%v. unexpected bodiesDiffer=%v; body1=%q; body2=%q",
				connectionClose, bodiesDiffer, body1, body2)
		}
	}
}

func TestTransportIdleCacheKeys(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	tr := &Transport{DisableKeepAlives: false}
	c := &Client{Transport: tr}

	if e, g := 0, len(tr.IdleConnKeysForTesting()); e != g {
		t.Errorf("After CloseIdleConnections expected %d idle conn cache keys; got %d", e, g)
	}

	resp, _, err := c.Get(ts.URL)
	if err != nil {
		t.Error(err)
	}
	ioutil.ReadAll(resp.Body)

	keys := tr.IdleConnKeysForTesting()
	if e, g := 1, len(keys); e != g {
		t.Fatalf("After Get expected %d idle conn cache keys; got %d", e, g)
	}

	if e := "|http|" + ts.Listener.Addr().String(); keys[0] != e {
		t.Errorf("Expected idle cache key %q; got %q", e, keys[0])
	}

	tr.CloseIdleConnections()
	if e, g := 0, len(tr.IdleConnKeysForTesting()); e != g {
		t.Errorf("After CloseIdleConnections expected %d idle conn cache keys; got %d", e, g)
	}
}

func TestTransportMaxPerHostIdleConns(t *testing.T) {
	ch := make(chan string)
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		fmt.Fprintf(w, "%s", <-ch)
	}))
	defer ts.Close()
	maxIdleConns := 2
	tr := &Transport{DisableKeepAlives: false, MaxIdleConnsPerHost: maxIdleConns}
	c := &Client{Transport: tr}

	// Start 3 outstanding requests (will hang until we write to
	// ch)
	donech := make(chan bool)
	doReq := func() {
		resp, _, err := c.Get(ts.URL)
		if err != nil {
			t.Error(err)
		}
		ioutil.ReadAll(resp.Body)
		donech <- true
	}
	go doReq()
	go doReq()
	go doReq()

	if e, g := 0, len(tr.IdleConnKeysForTesting()); e != g {
		t.Fatalf("Before writes, expected %d idle conn cache keys; got %d", e, g)
	}

	ch <- "res1"
	<-donech
	keys := tr.IdleConnKeysForTesting()
	if e, g := 1, len(keys); e != g {
		t.Fatalf("after first response, expected %d idle conn cache keys; got %d", e, g)
	}
	cacheKey := "|http|" + ts.Listener.Addr().String()
	if keys[0] != cacheKey {
		t.Fatalf("Expected idle cache key %q; got %q", cacheKey, keys[0])
	}
	if e, g := 1, tr.IdleConnCountForTesting(cacheKey); e != g {
		t.Errorf("after first response, expected %d idle conns; got %d", e, g)
	}

	ch <- "res2"
	<-donech
	if e, g := 2, tr.IdleConnCountForTesting(cacheKey); e != g {
		t.Errorf("after second response, expected %d idle conns; got %d", e, g)
	}

	ch <- "res3"
	<-donech
	if e, g := maxIdleConns, tr.IdleConnCountForTesting(cacheKey); e != g {
		t.Errorf("after third response, still expected %d idle conns; got %d", e, g)
	}
}

func TestTransportServerClosingUnexpectedly(t *testing.T) {
	ts := httptest.NewServer(hostPortHandler)
	defer ts.Close()

	tr := &Transport{}
	c := &Client{Transport: tr}

	fetch := func(n int) string {
		res, _, err := c.Get(ts.URL)
		if err != nil {
			t.Fatalf("error in req #%d, GET: %v", n, err)
		}
		body, err := ioutil.ReadAll(res.Body)
		if err != nil {
			t.Fatalf("error in req #%d, ReadAll: %v", n, err)
		}
		res.Body.Close()
		return string(body)
	}

	body1 := fetch(1)
	body2 := fetch(2)

	ts.CloseClientConnections() // surprise!
	time.Sleep(25e6)            // idle for a bit (test is inherently racey, but expectedly)

	body3 := fetch(3)

	if body1 != body2 {
		t.Errorf("expected body1 and body2 to be equal")
	}
	if body2 == body3 {
		t.Errorf("expected body2 and body3 to be different")
	}
}

// TestTransportHeadResponses verifies that we deal with Content-Lengths
// with no bodies properly
func TestTransportHeadResponses(t *testing.T) {
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		if r.Method != "HEAD" {
			panic("expected HEAD; got " + r.Method)
		}
		w.Header().Set("Content-Length", "123")
		w.WriteHeader(200)
	}))
	defer ts.Close()

	tr := &Transport{DisableKeepAlives: false}
	c := &Client{Transport: tr}
	for i := 0; i < 2; i++ {
		res, err := c.Head(ts.URL)
		if err != nil {
			t.Errorf("error on loop %d: %v", i, err)
		}
		if e, g := "123", res.Header.Get("Content-Length"); e != g {
			t.Errorf("loop %d: expected Content-Length header of %q, got %q", i, e, g)
		}
		if e, g := int64(0), res.ContentLength; e != g {
			t.Errorf("loop %d: expected res.ContentLength of %v, got %v", i, e, g)
		}
	}
}

// TestTransportHeadChunkedResponse verifies that we ignore chunked transfer-encoding
// on responses to HEAD requests.
func TestTransportHeadChunkedResponse(t *testing.T) {
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		if r.Method != "HEAD" {
			panic("expected HEAD; got " + r.Method)
		}
		w.Header().Set("Transfer-Encoding", "chunked") // client should ignore
		w.Header().Set("x-client-ipport", r.RemoteAddr)
		w.WriteHeader(200)
	}))
	defer ts.Close()

	tr := &Transport{DisableKeepAlives: false}
	c := &Client{Transport: tr}

	res1, err := c.Head(ts.URL)
	if err != nil {
		t.Fatalf("request 1 error: %v", err)
	}
	res2, err := c.Head(ts.URL)
	if err != nil {
		t.Fatalf("request 2 error: %v", err)
	}
	if v1, v2 := res1.Header.Get("x-client-ipport"), res2.Header.Get("x-client-ipport"); v1 != v2 {
		t.Errorf("ip/ports differed between head requests: %q vs %q", v1, v2)
	}
}

func TestTransportNilURL(t *testing.T) {
	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
		fmt.Fprintf(w, "Hi")
	}))
	defer ts.Close()

	req := new(Request)
	req.URL = nil // what we're actually testing
	req.Method = "GET"
	req.RawURL = ts.URL
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1

	tr := &Transport{}
	res, err := tr.RoundTrip(req)
	if err != nil {
		t.Fatalf("unexpected RoundTrip error: %v", err)
	}
	body, err := ioutil.ReadAll(res.Body)
	if g, e := string(body), "Hi"; g != e {
		t.Fatalf("Expected response body of %q; got %q", e, g)
	}
}
<|endoftext|>


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

// Network interface identification for Linux

package net

import (
	"os"
	"syscall"
	"unsafe"
)

// If the ifindex is zero, interfaceTable returns mappings of all
// network interfaces.  Otherwise it returns a mapping of a specific
// interface.
func interfaceTable(ifindex int) ([]Interface, error) {
	tab, err := syscall.NetlinkRIB(syscall.RTM_GETLINK, syscall.AF_UNSPEC)
	if err != nil {
		return nil, os.NewSyscallError("netlink rib", err)
	}

	msgs, err := syscall.ParseNetlinkMessage(tab)
	if err != nil {
		return nil, os.NewSyscallError("netlink message", err)
	}

	var ift []Interface
	for _, m := range msgs {
		switch m.Header.Type {
		case syscall.NLMSG_DONE:
			goto done
		case syscall.RTM_NEWLINK:
			ifim := (*syscall.IfInfomsg)(unsafe.Pointer(&m.Data[0]))
			if ifindex == 0 || ifindex == int(ifim.Index) {
				attrs, err := syscall.ParseNetlinkRouteAttr(&m)
				if err != nil {
					return nil, os.NewSyscallError("netlink routeattr", err)
				}
				ifi := newLink(ifim, attrs)
				ift = append(ift, ifi)
			}
		}
	}
done:
	return ift, nil
}

func newLink(ifim *syscall.IfInfomsg, attrs []syscall.NetlinkRouteAttr) Interface {
	ifi := Interface{Index: int(ifim.Index), Flags: linkFlags(ifim.Flags)}
	for _, a := range attrs {
		switch a.Attr.Type {
		case syscall.IFLA_ADDRESS:
			var nonzero bool
			for _, b := range a.Value {
				if b != 0 {
					nonzero = true
				}
			}
			if nonzero {
				ifi.HardwareAddr = a.Value[:]
			}
		case syscall.IFLA_IFNAME:
			ifi.Name = string(a.Value[:len(a.Value)-1])
		case syscall.IFLA_MTU:
			ifi.MTU = int(*(*uint32)(unsafe.Pointer(&a.Value[:4][0])))
		}
	}
	return ifi
}

func linkFlags(rawFlags uint32) Flags {
	var f Flags
	if rawFlags&syscall.IFF_UP != 0 {
		f |= FlagUp
	}
	if rawFlags&syscall.IFF_BROADCAST != 0 {
		f |= FlagBroadcast
	}
	if rawFlags&syscall.IFF_LOOPBACK != 0 {
		f |= FlagLoopback
	}
	if rawFlags&syscall.IFF_POINTOPOINT != 0 {
		f |= FlagPointToPoint
	}
	if rawFlags&syscall.IFF_MULTICAST != 0 {
		f |= FlagMulticast
	}
	return f
}

// If the ifindex is zero, interfaceAddrTable returns addresses
// for all network interfaces.  Otherwise it returns addresses
// for a specific interface.
func interfaceAddrTable(ifindex int) ([]Addr, error) {
	tab, err := syscall.NetlinkRIB(syscall.RTM_GETADDR, syscall.AF_UNSPEC)
	if err != nil {
		return nil, os.NewSyscallError("netlink rib", err)
	}

	msgs, err := syscall.ParseNetlinkMessage(tab)
	if err != nil {
		return nil, os.NewSyscallError("netlink message", err)
	}

	ifat, err := addrTable(msgs, ifindex)
	if err != nil {
		return nil, err
	}
	return ifat, nil
}

func addrTable(msgs []syscall.NetlinkMessage, ifindex int) ([]Addr, error) {
	var ifat []Addr
	for _, m := range msgs {
		switch m.Header.Type {
		case syscall.NLMSG_DONE:
			goto done
		case syscall.RTM_NEWADDR:
			ifam := (*syscall.IfAddrmsg)(unsafe.Pointer(&m.Data[0]))
			if ifindex == 0 || ifindex == int(ifam.Index) {
				attrs, err := syscall.ParseNetlinkRouteAttr(&m)
				if err != nil {
					return nil, os.NewSyscallError("netlink routeattr", err)
				}
				ifat = append(ifat, newAddr(attrs, int(ifam.Family), int(ifam.Prefixlen)))
			}
		}
	}
done:
	return ifat, nil
}

func newAddr(attrs []syscall.NetlinkRouteAttr, family, pfxlen int) Addr {
	ifa := &IPNet{}
	for _, a := range attrs {
		switch a.Attr.Type {
		case syscall.IFA_ADDRESS:
			switch family {
			case syscall.AF_INET:
				ifa.IP = IPv4(a.Value[0], a.Value[1], a.Value[2], a.Value[3])
				ifa.Mask = CIDRMask(pfxlen, 8*IPv4len)
			case syscall.AF_INET6:
				ifa.IP = make(IP, IPv6len)
				copy(ifa.IP, a.Value[:])
				ifa.Mask = CIDRMask(pfxlen, 8*IPv6len)
			}
		}
	}
	return ifa
}

// If the ifindex is zero, interfaceMulticastAddrTable returns
// addresses for all network interfaces.  Otherwise it returns
// addresses for a specific interface.
func interfaceMulticastAddrTable(ifindex int) ([]Addr, error) {
	var (
		err error
		ifi *Interface
	)
	if ifindex > 0 {
		ifi, err = InterfaceByIndex(ifindex)
		if err != nil {
			return nil, err
		}
	}
	ifmat4 := parseProcNetIGMP("/proc/net/igmp", ifi)
	ifmat6 := parseProcNetIGMP6("/proc/net/igmp6", ifi)
	return append(ifmat4, ifmat6...), nil
}

func parseProcNetIGMP(path string, ifi *Interface) []Addr {
	fd, err := open(path)
	if err != nil {
		return nil
	}
	defer fd.close()

	var (
		ifmat []Addr
		name  string
	)
	fd.readLine() // skip first line
	b := make([]byte, IPv4len)
	for l, ok := fd.readLine(); ok; l, ok = fd.readLine() {
		f := splitAtBytes(l, " :\r\t\n")
		if len(f) < 4 {
			continue
		}
		switch {
		case l[0] != ' ' && l[0] != '\t': // new interface line
			name = f[1]
		case len(f[0]) == 8:
			if ifi == nil || name == ifi.Name {
				// The Linux kernel puts the IP
				// address in /proc/net/igmp in native
				// endianness.
				for i := 0; i+1 < len(f[0]); i += 2 {
					b[i/2], _ = xtoi2(f[0][i:i+2], 0)
				}
				i := *(*uint32)(unsafe.Pointer(&b[:4][0]))
				ifma := IPAddr{IP: IPv4(byte(i>>24), byte(i>>16), byte(i>>8), byte(i))}
				ifmat = append(ifmat, ifma.toAddr())
			}
		}
	}
	return ifmat
}

func parseProcNetIGMP6(path string, ifi *Interface) []Addr {
	fd, err := open(path)
	if err != nil {
		return nil
	}
	defer fd.close()

	var ifmat []Addr
	b := make([]byte, IPv6len)
	for l, ok := fd.readLine(); ok; l, ok = fd.readLine() {
		f := splitAtBytes(l, " \r\t\n")
		if len(f) < 6 {
			continue
		}
		if ifi == nil || f[1] == ifi.Name {
			for i := 0; i+1 < len(f[2]); i += 2 {
				b[i/2], _ = xtoi2(f[2][i:i+2], 0)
			}
			ifma := IPAddr{IP: IP{b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7], b[8], b[9], b[10], b[11], b[12], b[13], b[14], b[15]}}
			ifmat = append(ifmat, ifma.toAddr())
		}
	}
	return ifmat
}
<commit_msg>net: return correct point-to-point interface address on linux<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Network interface identification for Linux

package net

import (
	"os"
	"syscall"
	"unsafe"
)

// If the ifindex is zero, interfaceTable returns mappings of all
// network interfaces.  Otherwise it returns a mapping of a specific
// interface.
func interfaceTable(ifindex int) ([]Interface, error) {
	tab, err := syscall.NetlinkRIB(syscall.RTM_GETLINK, syscall.AF_UNSPEC)
	if err != nil {
		return nil, os.NewSyscallError("netlink rib", err)
	}

	msgs, err := syscall.ParseNetlinkMessage(tab)
	if err != nil {
		return nil, os.NewSyscallError("netlink message", err)
	}

	var ift []Interface
	for _, m := range msgs {
		switch m.Header.Type {
		case syscall.NLMSG_DONE:
			goto done
		case syscall.RTM_NEWLINK:
			ifim := (*syscall.IfInfomsg)(unsafe.Pointer(&m.Data[0]))
			if ifindex == 0 || ifindex == int(ifim.Index) {
				attrs, err := syscall.ParseNetlinkRouteAttr(&m)
				if err != nil {
					return nil, os.NewSyscallError("netlink routeattr", err)
				}
				ifi := newLink(ifim, attrs)
				ift = append(ift, ifi)
			}
		}
	}
done:
	return ift, nil
}

func newLink(ifim *syscall.IfInfomsg, attrs []syscall.NetlinkRouteAttr) Interface {
	ifi := Interface{Index: int(ifim.Index), Flags: linkFlags(ifim.Flags)}
	for _, a := range attrs {
		switch a.Attr.Type {
		case syscall.IFLA_ADDRESS:
			var nonzero bool
			for _, b := range a.Value {
				if b != 0 {
					nonzero = true
				}
			}
			if nonzero {
				ifi.HardwareAddr = a.Value[:]
			}
		case syscall.IFLA_IFNAME:
			ifi.Name = string(a.Value[:len(a.Value)-1])
		case syscall.IFLA_MTU:
			ifi.MTU = int(*(*uint32)(unsafe.Pointer(&a.Value[:4][0])))
		}
	}
	return ifi
}

func linkFlags(rawFlags uint32) Flags {
	var f Flags
	if rawFlags&syscall.IFF_UP != 0 {
		f |= FlagUp
	}
	if rawFlags&syscall.IFF_BROADCAST != 0 {
		f |= FlagBroadcast
	}
	if rawFlags&syscall.IFF_LOOPBACK != 0 {
		f |= FlagLoopback
	}
	if rawFlags&syscall.IFF_POINTOPOINT != 0 {
		f |= FlagPointToPoint
	}
	if rawFlags&syscall.IFF_MULTICAST != 0 {
		f |= FlagMulticast
	}
	return f
}

// If the ifindex is zero, interfaceAddrTable returns addresses
// for all network interfaces.  Otherwise it returns addresses
// for a specific interface.
func interfaceAddrTable(ifindex int) ([]Addr, error) {
	tab, err := syscall.NetlinkRIB(syscall.RTM_GETADDR, syscall.AF_UNSPEC)
	if err != nil {
		return nil, os.NewSyscallError("netlink rib", err)
	}

	msgs, err := syscall.ParseNetlinkMessage(tab)
	if err != nil {
		return nil, os.NewSyscallError("netlink message", err)
	}

	ifat, err := addrTable(msgs, ifindex)
	if err != nil {
		return nil, err
	}
	return ifat, nil
}

func addrTable(msgs []syscall.NetlinkMessage, ifindex int) ([]Addr, error) {
	var ifat []Addr
	for _, m := range msgs {
		switch m.Header.Type {
		case syscall.NLMSG_DONE:
			goto done
		case syscall.RTM_NEWADDR:
			ifam := (*syscall.IfAddrmsg)(unsafe.Pointer(&m.Data[0]))
			ifi, err := InterfaceByIndex(int(ifam.Index))
			if err != nil {
				return nil, err
			}
			if ifindex == 0 || ifindex == int(ifam.Index) {
				attrs, err := syscall.ParseNetlinkRouteAttr(&m)
				if err != nil {
					return nil, os.NewSyscallError("netlink routeattr", err)
				}
				ifat = append(ifat, newAddr(attrs, ifi, ifam))
			}
		}
	}
done:
	return ifat, nil
}

func newAddr(attrs []syscall.NetlinkRouteAttr, ifi *Interface, ifam *syscall.IfAddrmsg) Addr {
	ifa := &IPNet{}
	for _, a := range attrs {
		if ifi.Flags&FlagPointToPoint != 0 && a.Attr.Type == syscall.IFA_LOCAL ||
			ifi.Flags&FlagPointToPoint == 0 && a.Attr.Type == syscall.IFA_ADDRESS {
			switch ifam.Family {
			case syscall.AF_INET:
				ifa.IP = IPv4(a.Value[0], a.Value[1], a.Value[2], a.Value[3])
				ifa.Mask = CIDRMask(int(ifam.Prefixlen), 8*IPv4len)
			case syscall.AF_INET6:
				ifa.IP = make(IP, IPv6len)
				copy(ifa.IP, a.Value[:])
				ifa.Mask = CIDRMask(int(ifam.Prefixlen), 8*IPv6len)
			}
		}
	}
	return ifa
}

// If the ifindex is zero, interfaceMulticastAddrTable returns
// addresses for all network interfaces.  Otherwise it returns
// addresses for a specific interface.
func interfaceMulticastAddrTable(ifindex int) ([]Addr, error) {
	var (
		err error
		ifi *Interface
	)
	if ifindex > 0 {
		ifi, err = InterfaceByIndex(ifindex)
		if err != nil {
			return nil, err
		}
	}
	ifmat4 := parseProcNetIGMP("/proc/net/igmp", ifi)
	ifmat6 := parseProcNetIGMP6("/proc/net/igmp6", ifi)
	return append(ifmat4, ifmat6...), nil
}

func parseProcNetIGMP(path string, ifi *Interface) []Addr {
	fd, err := open(path)
	if err != nil {
		return nil
	}
	defer fd.close()

	var (
		ifmat []Addr
		name  string
	)
	fd.readLine() // skip first line
	b := make([]byte, IPv4len)
	for l, ok := fd.readLine(); ok; l, ok = fd.readLine() {
		f := splitAtBytes(l, " :\r\t\n")
		if len(f) < 4 {
			continue
		}
		switch {
		case l[0] != ' ' && l[0] != '\t': // new interface line
			name = f[1]
		case len(f[0]) == 8:
			if ifi == nil || name == ifi.Name {
				// The Linux kernel puts the IP
				// address in /proc/net/igmp in native
				// endianness.
				for i := 0; i+1 < len(f[0]); i += 2 {
					b[i/2], _ = xtoi2(f[0][i:i+2], 0)
				}
				i := *(*uint32)(unsafe.Pointer(&b[:4][0]))
				ifma := IPAddr{IP: IPv4(byte(i>>24), byte(i>>16), byte(i>>8), byte(i))}
				ifmat = append(ifmat, ifma.toAddr())
			}
		}
	}
	return ifmat
}

func parseProcNetIGMP6(path string, ifi *Interface) []Addr {
	fd, err := open(path)
	if err != nil {
		return nil
	}
	defer fd.close()

	var ifmat []Addr
	b := make([]byte, IPv6len)
	for l, ok := fd.readLine(); ok; l, ok = fd.readLine() {
		f := splitAtBytes(l, " \r\t\n")
		if len(f) < 6 {
			continue
		}
		if ifi == nil || f[1] == ifi.Name {
			for i := 0; i+1 < len(f[2]); i += 2 {
				b[i/2], _ = xtoi2(f[2][i:i+2], 0)
			}
			ifma := IPAddr{IP: IP{b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7], b[8], b[9], b[10], b[11], b[12], b[13], b[14], b[15]}}
			ifmat = append(ifmat, ifma.toAddr())
		}
	}
	return ifmat
}
<|endoftext|>


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

// Package pprof writes runtime profiling data in the format expected
// by the pprof visualization tool.
// For more information about pprof, see
// http://code.google.com/p/google-perftools/.
package pprof

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"runtime"
	"sync"
)

// WriteHeapProfile writes a pprof-formatted heap profile to w.
// If a write to w returns an error, WriteHeapProfile returns that error.
// Otherwise, WriteHeapProfile returns nil.
func WriteHeapProfile(w io.Writer) os.Error {
	// Find out how many records there are (MemProfile(nil, false)),
	// allocate that many records, and get the data.
	// There's a race—more records might be added between
	// the two calls—so allocate a few extra records for safety
	// and also try again if we're very unlucky.
	// The loop should only execute one iteration in the common case.
	var p []runtime.MemProfileRecord
	n, ok := runtime.MemProfile(nil, false)
	for {
		// Allocate room for a slightly bigger profile,
		// in case a few more entries have been added
		// since the call to MemProfile.
		p = make([]runtime.MemProfileRecord, n+50)
		n, ok = runtime.MemProfile(p, false)
		if ok {
			p = p[0:n]
			break
		}
		// Profile grew; try again.
	}

	var total runtime.MemProfileRecord
	for i := range p {
		r := &p[i]
		total.AllocBytes += r.AllocBytes
		total.AllocObjects += r.AllocObjects
		total.FreeBytes += r.FreeBytes
		total.FreeObjects += r.FreeObjects
	}

	// Technically the rate is MemProfileRate not 2*MemProfileRate,
	// but early versions of the C++ heap profiler reported 2*MemProfileRate,
	// so that's what pprof has come to expect.
	b := bufio.NewWriter(w)
	fmt.Fprintf(b, "heap profile: %d: %d [%d: %d] @ heap/%d\n",
		total.InUseObjects(), total.InUseBytes(),
		total.AllocObjects, total.AllocBytes,
		2*runtime.MemProfileRate)

	for i := range p {
		r := &p[i]
		fmt.Fprintf(b, "%d: %d [%d: %d] @",
			r.InUseObjects(), r.InUseBytes(),
			r.AllocObjects, r.AllocBytes)
		for _, pc := range r.Stack() {
			fmt.Fprintf(b, " %#x", pc)
		}
		fmt.Fprintf(b, "\n")
	}

	// Print memstats information too.
	// Pprof will ignore, but useful for people.
	s := &runtime.MemStats
	fmt.Fprintf(b, "\n# runtime.MemStats\n")
	fmt.Fprintf(b, "# Alloc = %d\n", s.Alloc)
	fmt.Fprintf(b, "# TotalAlloc = %d\n", s.TotalAlloc)
	fmt.Fprintf(b, "# Sys = %d\n", s.Sys)
	fmt.Fprintf(b, "# Lookups = %d\n", s.Lookups)
	fmt.Fprintf(b, "# Mallocs = %d\n", s.Mallocs)

	fmt.Fprintf(b, "# HeapAlloc = %d\n", s.HeapAlloc)
	fmt.Fprintf(b, "# HeapSys = %d\n", s.HeapSys)
	fmt.Fprintf(b, "# HeapIdle = %d\n", s.HeapIdle)
	fmt.Fprintf(b, "# HeapInuse = %d\n", s.HeapInuse)

	fmt.Fprintf(b, "# Stack = %d / %d\n", s.StackInuse, s.StackSys)
	fmt.Fprintf(b, "# MSpan = %d / %d\n", s.MSpanInuse, s.MSpanSys)
	fmt.Fprintf(b, "# MCache = %d / %d\n", s.MCacheInuse, s.MCacheSys)
	fmt.Fprintf(b, "# BuckHashSys = %d\n", s.BuckHashSys)

	fmt.Fprintf(b, "# NextGC = %d\n", s.NextGC)
	fmt.Fprintf(b, "# PauseNs = %d\n", s.PauseNs)
	fmt.Fprintf(b, "# NumGC = %d\n", s.NumGC)
	fmt.Fprintf(b, "# EnableGC = %v\n", s.EnableGC)
	fmt.Fprintf(b, "# DebugGC = %v\n", s.DebugGC)

	fmt.Fprintf(b, "# BySize = Size * (Active = Mallocs - Frees)\n")
	fmt.Fprintf(b, "# (Excluding large blocks.)\n")
	for _, t := range s.BySize {
		if t.Mallocs > 0 {
			fmt.Fprintf(b, "#   %d * (%d = %d - %d)\n", t.Size, t.Mallocs-t.Frees, t.Mallocs, t.Frees)
		}
	}
	return b.Flush()
}

var cpu struct {
	sync.Mutex
	profiling bool
	done      chan bool
}

// StartCPUProfile enables CPU profiling for the current process.
// While profiling, the profile will be buffered and written to w.
// StartCPUProfile returns an error if profiling is already enabled.
func StartCPUProfile(w io.Writer) os.Error {
	// The runtime routines allow a variable profiling rate,
	// but in practice operating systems cannot trigger signals
	// at more than about 500 Hz, and our processing of the
	// signal is not cheap (mostly getting the stack trace).
	// 100 Hz is a reasonable choice: it is frequent enough to
	// produce useful data, rare enough not to bog down the
	// system, and a nice round number to make it easy to
	// convert sample counts to seconds.  Instead of requiring
	// each client to specify the frequency, we hard code it.
	const hz = 100

	// Avoid queueing behind StopCPUProfile.
	// Could use TryLock instead if we had it.
	if cpu.profiling {
		return fmt.Errorf("cpu profiling already in use")
	}

	cpu.Lock()
	defer cpu.Unlock()
	if cpu.done == nil {
		cpu.done = make(chan bool)
	}
	// Double-check.
	if cpu.profiling {
		return fmt.Errorf("cpu profiling already in use")
	}
	cpu.profiling = true
	runtime.SetCPUProfileRate(hz)
	go profileWriter(w)
	return nil
}

func profileWriter(w io.Writer) {
	for {
		data := runtime.CPUProfile()
		if data == nil {
			break
		}
		w.Write(data)
	}
	cpu.done <- true
}

// StopCPUProfile stops the current CPU profile, if any.
// StopCPUProfile only returns after all the writes for the
// profile have completed.
func StopCPUProfile() {
	cpu.Lock()
	defer cpu.Unlock()

	if !cpu.profiling {
		return
	}
	cpu.profiling = false
	runtime.SetCPUProfileRate(0)
	<-cpu.done
}
<commit_msg>runtime/pprof: document OS X being broken<commit_after>// Copyright 2010 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package pprof writes runtime profiling data in the format expected
// by the pprof visualization tool.
// For more information about pprof, see
// http://code.google.com/p/google-perftools/.
package pprof

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"runtime"
	"sync"
)

// BUG(rsc): CPU profiling is broken on OS X, due to an Apple kernel bug.
// For details, see http://code.google.com/p/go/source/detail?r=35b716c94225.

// WriteHeapProfile writes a pprof-formatted heap profile to w.
// If a write to w returns an error, WriteHeapProfile returns that error.
// Otherwise, WriteHeapProfile returns nil.
func WriteHeapProfile(w io.Writer) os.Error {
	// Find out how many records there are (MemProfile(nil, false)),
	// allocate that many records, and get the data.
	// There's a race—more records might be added between
	// the two calls—so allocate a few extra records for safety
	// and also try again if we're very unlucky.
	// The loop should only execute one iteration in the common case.
	var p []runtime.MemProfileRecord
	n, ok := runtime.MemProfile(nil, false)
	for {
		// Allocate room for a slightly bigger profile,
		// in case a few more entries have been added
		// since the call to MemProfile.
		p = make([]runtime.MemProfileRecord, n+50)
		n, ok = runtime.MemProfile(p, false)
		if ok {
			p = p[0:n]
			break
		}
		// Profile grew; try again.
	}

	var total runtime.MemProfileRecord
	for i := range p {
		r := &p[i]
		total.AllocBytes += r.AllocBytes
		total.AllocObjects += r.AllocObjects
		total.FreeBytes += r.FreeBytes
		total.FreeObjects += r.FreeObjects
	}

	// Technically the rate is MemProfileRate not 2*MemProfileRate,
	// but early versions of the C++ heap profiler reported 2*MemProfileRate,
	// so that's what pprof has come to expect.
	b := bufio.NewWriter(w)
	fmt.Fprintf(b, "heap profile: %d: %d [%d: %d] @ heap/%d\n",
		total.InUseObjects(), total.InUseBytes(),
		total.AllocObjects, total.AllocBytes,
		2*runtime.MemProfileRate)

	for i := range p {
		r := &p[i]
		fmt.Fprintf(b, "%d: %d [%d: %d] @",
			r.InUseObjects(), r.InUseBytes(),
			r.AllocObjects, r.AllocBytes)
		for _, pc := range r.Stack() {
			fmt.Fprintf(b, " %#x", pc)
		}
		fmt.Fprintf(b, "\n")
	}

	// Print memstats information too.
	// Pprof will ignore, but useful for people.
	s := &runtime.MemStats
	fmt.Fprintf(b, "\n# runtime.MemStats\n")
	fmt.Fprintf(b, "# Alloc = %d\n", s.Alloc)
	fmt.Fprintf(b, "# TotalAlloc = %d\n", s.TotalAlloc)
	fmt.Fprintf(b, "# Sys = %d\n", s.Sys)
	fmt.Fprintf(b, "# Lookups = %d\n", s.Lookups)
	fmt.Fprintf(b, "# Mallocs = %d\n", s.Mallocs)

	fmt.Fprintf(b, "# HeapAlloc = %d\n", s.HeapAlloc)
	fmt.Fprintf(b, "# HeapSys = %d\n", s.HeapSys)
	fmt.Fprintf(b, "# HeapIdle = %d\n", s.HeapIdle)
	fmt.Fprintf(b, "# HeapInuse = %d\n", s.HeapInuse)

	fmt.Fprintf(b, "# Stack = %d / %d\n", s.StackInuse, s.StackSys)
	fmt.Fprintf(b, "# MSpan = %d / %d\n", s.MSpanInuse, s.MSpanSys)
	fmt.Fprintf(b, "# MCache = %d / %d\n", s.MCacheInuse, s.MCacheSys)
	fmt.Fprintf(b, "# BuckHashSys = %d\n", s.BuckHashSys)

	fmt.Fprintf(b, "# NextGC = %d\n", s.NextGC)
	fmt.Fprintf(b, "# PauseNs = %d\n", s.PauseNs)
	fmt.Fprintf(b, "# NumGC = %d\n", s.NumGC)
	fmt.Fprintf(b, "# EnableGC = %v\n", s.EnableGC)
	fmt.Fprintf(b, "# DebugGC = %v\n", s.DebugGC)

	fmt.Fprintf(b, "# BySize = Size * (Active = Mallocs - Frees)\n")
	fmt.Fprintf(b, "# (Excluding large blocks.)\n")
	for _, t := range s.BySize {
		if t.Mallocs > 0 {
			fmt.Fprintf(b, "#   %d * (%d = %d - %d)\n", t.Size, t.Mallocs-t.Frees, t.Mallocs, t.Frees)
		}
	}
	return b.Flush()
}

var cpu struct {
	sync.Mutex
	profiling bool
	done      chan bool
}

// StartCPUProfile enables CPU profiling for the current process.
// While profiling, the profile will be buffered and written to w.
// StartCPUProfile returns an error if profiling is already enabled.
func StartCPUProfile(w io.Writer) os.Error {
	// The runtime routines allow a variable profiling rate,
	// but in practice operating systems cannot trigger signals
	// at more than about 500 Hz, and our processing of the
	// signal is not cheap (mostly getting the stack trace).
	// 100 Hz is a reasonable choice: it is frequent enough to
	// produce useful data, rare enough not to bog down the
	// system, and a nice round number to make it easy to
	// convert sample counts to seconds.  Instead of requiring
	// each client to specify the frequency, we hard code it.
	const hz = 100

	// Avoid queueing behind StopCPUProfile.
	// Could use TryLock instead if we had it.
	if cpu.profiling {
		return fmt.Errorf("cpu profiling already in use")
	}

	cpu.Lock()
	defer cpu.Unlock()
	if cpu.done == nil {
		cpu.done = make(chan bool)
	}
	// Double-check.
	if cpu.profiling {
		return fmt.Errorf("cpu profiling already in use")
	}
	cpu.profiling = true
	runtime.SetCPUProfileRate(hz)
	go profileWriter(w)
	return nil
}

func profileWriter(w io.Writer) {
	for {
		data := runtime.CPUProfile()
		if data == nil {
			break
		}
		w.Write(data)
	}
	cpu.done <- true
}

// StopCPUProfile stops the current CPU profile, if any.
// StopCPUProfile only returns after all the writes for the
// profile have completed.
func StopCPUProfile() {
	cpu.Lock()
	defer cpu.Unlock()

	if !cpu.profiling {
		return
	}
	cpu.profiling = false
	runtime.SetCPUProfileRate(0)
	<-cpu.done
}
<|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 apis

import (
	"strings"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/sets"
)

const (
	// LabelOS is a label to indicate the operating system of the node.
	// The OS labels are promoted to GA in 1.14. kubelet applies GA labels and stop applying the beta OS labels in Kubernetes 1.19.
	LabelOS = "beta.kubernetes.io/os"
	// LabelArch is a label to indicate the architecture of the node.
	// The Arch labels are promoted to GA in 1.14. kubelet applies GA labels and stop applying the beta Arch labels in Kubernetes 1.19.
	LabelArch = "beta.kubernetes.io/arch"
)

var kubeletLabels = sets.NewString(
	v1.LabelHostname,
	v1.LabelTopologyZone,
	v1.LabelTopologyRegion,
	v1.LabelFailureDomainBetaZone,
	v1.LabelFailureDomainBetaRegion,
	v1.LabelInstanceType,
	v1.LabelInstanceTypeStable,
	v1.LabelOSStable,
	v1.LabelArchStable,

	LabelOS,
	LabelArch,

	// These are special for OpenShift:
	"node-role.kubernetes.io/master",
	"node-role.kubernetes.io/worker",
	"node-role.kubernetes.io/etcd",
)

var kubeletLabelNamespaces = sets.NewString(
	v1.LabelNamespaceSuffixKubelet,
	v1.LabelNamespaceSuffixNode,
)

// KubeletLabels returns the list of label keys kubelets are allowed to set on their own Node objects
func KubeletLabels() []string {
	return kubeletLabels.List()
}

// KubeletLabelNamespaces returns the list of label key namespaces kubelets are allowed to set on their own Node objects
func KubeletLabelNamespaces() []string {
	return kubeletLabelNamespaces.List()
}

// IsKubeletLabel returns true if the label key is one that kubelets are allowed to set on their own Node object.
// This checks if the key is in the KubeletLabels() list, or has a namespace in the KubeletLabelNamespaces() list.
func IsKubeletLabel(key string) bool {
	if kubeletLabels.Has(key) {
		return true
	}

	namespace := getLabelNamespace(key)
	for allowedNamespace := range kubeletLabelNamespaces {
		if namespace == allowedNamespace || strings.HasSuffix(namespace, "."+allowedNamespace) {
			return true
		}
	}

	return false
}

func getLabelNamespace(key string) string {
	if parts := strings.SplitN(key, "/", 2); len(parts) == 2 {
		return parts[0]
	}
	return ""
}
<commit_msg>UPSTREAM: <carry>: add control plane to allow roles<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 apis

import (
	"strings"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/sets"
)

const (
	// LabelOS is a label to indicate the operating system of the node.
	// The OS labels are promoted to GA in 1.14. kubelet applies GA labels and stop applying the beta OS labels in Kubernetes 1.19.
	LabelOS = "beta.kubernetes.io/os"
	// LabelArch is a label to indicate the architecture of the node.
	// The Arch labels are promoted to GA in 1.14. kubelet applies GA labels and stop applying the beta Arch labels in Kubernetes 1.19.
	LabelArch = "beta.kubernetes.io/arch"
)

var kubeletLabels = sets.NewString(
	v1.LabelHostname,
	v1.LabelTopologyZone,
	v1.LabelTopologyRegion,
	v1.LabelFailureDomainBetaZone,
	v1.LabelFailureDomainBetaRegion,
	v1.LabelInstanceType,
	v1.LabelInstanceTypeStable,
	v1.LabelOSStable,
	v1.LabelArchStable,

	LabelOS,
	LabelArch,

	// These are special for OpenShift:
	"node-role.kubernetes.io/control-plane",
	"node-role.kubernetes.io/master",
	"node-role.kubernetes.io/worker",
	"node-role.kubernetes.io/etcd",
)

var kubeletLabelNamespaces = sets.NewString(
	v1.LabelNamespaceSuffixKubelet,
	v1.LabelNamespaceSuffixNode,
)

// KubeletLabels returns the list of label keys kubelets are allowed to set on their own Node objects
func KubeletLabels() []string {
	return kubeletLabels.List()
}

// KubeletLabelNamespaces returns the list of label key namespaces kubelets are allowed to set on their own Node objects
func KubeletLabelNamespaces() []string {
	return kubeletLabelNamespaces.List()
}

// IsKubeletLabel returns true if the label key is one that kubelets are allowed to set on their own Node object.
// This checks if the key is in the KubeletLabels() list, or has a namespace in the KubeletLabelNamespaces() list.
func IsKubeletLabel(key string) bool {
	if kubeletLabels.Has(key) {
		return true
	}

	namespace := getLabelNamespace(key)
	for allowedNamespace := range kubeletLabelNamespaces {
		if namespace == allowedNamespace || strings.HasSuffix(namespace, "."+allowedNamespace) {
			return true
		}
	}

	return false
}

func getLabelNamespace(key string) string {
	if parts := strings.SplitN(key, "/", 2); len(parts) == 2 {
		return parts[0]
	}
	return ""
}
<|endoftext|>


<commit_before>/*
Copyright 2020 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 controlplane

import (
	"context"
	"strings"
	"testing"
	"time"

	coordinationv1 "k8s.io/api/coordination/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apiserver/pkg/features"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/client-go/kubernetes"
	featuregatetesting "k8s.io/component-base/featuregate/testing"
	kubeapiservertesting "k8s.io/kubernetes/cmd/kube-apiserver/app/testing"
	"k8s.io/kubernetes/pkg/controlplane"
	"k8s.io/kubernetes/test/integration/framework"
	"k8s.io/utils/pointer"
)

const (
	testLeaseName = "apiserver-lease-test"
)

func TestCreateLeaseOnStart(t *testing.T) {
	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.APIServerIdentity, true)()
	result := kubeapiservertesting.StartTestServerOrDie(t, nil, nil, framework.SharedEtcd())
	defer result.TearDownFn()

	kubeclient, err := kubernetes.NewForConfig(result.ClientConfig)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	t.Logf(`Waiting the kube-apiserver Lease to be created`)
	if err := wait.PollImmediate(500*time.Millisecond, 10*time.Second, func() (bool, error) {
		leases, err := kubeclient.
			CoordinationV1().
			Leases(metav1.NamespaceSystem).
			List(context.TODO(), metav1.ListOptions{LabelSelector: controlplane.KubeAPIServerIdentityLeaseLabelSelector})
		if err != nil {
			return false, err
		}
		if leases != nil && len(leases.Items) == 1 && strings.HasPrefix(leases.Items[0].Name, "kube-apiserver-") {
			return true, nil
		}
		return false, nil
	}); err != nil {
		t.Fatalf("Failed to see the kube-apiserver lease: %v", err)
	}
}

func TestLeaseGarbageCollection(t *testing.T) {
	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.APIServerIdentity, true)()
	result := kubeapiservertesting.StartTestServerOrDie(t, nil,
		// This shorten the GC check period to make the test run faster.
		// Since we are testing GC behavior on leases we create, what happens to
		// the real apiserver lease doesn't matter.
		[]string{"--identity-lease-duration-seconds=1"},
		framework.SharedEtcd())
	defer result.TearDownFn()

	kubeclient, err := kubernetes.NewForConfig(result.ClientConfig)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	expiredLease := newTestLease(time.Now().Add(-2*time.Hour), metav1.NamespaceSystem)
	t.Run("expired apiserver lease should be garbage collected",
		testLeaseGarbageCollected(t, kubeclient, expiredLease))

	freshLease := newTestLease(time.Now().Add(-2*time.Minute), metav1.NamespaceSystem)
	t.Run("fresh apiserver lease should not be garbage collected",
		testLeaseNotGarbageCollected(t, kubeclient, freshLease))

	expiredLease.Labels = nil
	t.Run("expired non-identity lease should not be garbage collected",
		testLeaseNotGarbageCollected(t, kubeclient, expiredLease))

	// identity leases (with k8s.io/component label) created in user namespaces should not be GC'ed
	expiredNonKubeSystemLease := newTestLease(time.Now().Add(-2*time.Hour), metav1.NamespaceDefault)
	t.Run("expired non-system identity lease should not be garbage collected",
		testLeaseNotGarbageCollected(t, kubeclient, expiredNonKubeSystemLease))
}

func testLeaseGarbageCollected(t *testing.T, client kubernetes.Interface, lease *coordinationv1.Lease) func(t *testing.T) {
	return func(t *testing.T) {
		ns := lease.Namespace
		if _, err := client.CoordinationV1().Leases(ns).Create(context.TODO(), lease, metav1.CreateOptions{}); err != nil {
			t.Fatalf("Unexpected error creating lease: %v", err)
		}
		if err := wait.PollImmediate(500*time.Millisecond, 5*time.Second, func() (bool, error) {
			_, err := client.CoordinationV1().Leases(ns).Get(context.TODO(), lease.Name, metav1.GetOptions{})
			if err == nil {
				return false, nil
			}
			if apierrors.IsNotFound(err) {
				return true, nil
			}
			return false, err
		}); err != nil {
			t.Fatalf("Failed to see the expired lease garbage collected: %v", err)
		}

	}
}

func testLeaseNotGarbageCollected(t *testing.T, client kubernetes.Interface, lease *coordinationv1.Lease) func(t *testing.T) {
	return func(t *testing.T) {
		ns := lease.Namespace
		if _, err := client.CoordinationV1().Leases(ns).Create(context.TODO(), lease, metav1.CreateOptions{}); err != nil {
			t.Fatalf("Unexpected error creating lease: %v", err)
		}
		if err := wait.PollImmediate(500*time.Millisecond, 5*time.Second, func() (bool, error) {
			_, err := client.CoordinationV1().Leases(ns).Get(context.TODO(), lease.Name, metav1.GetOptions{})
			if err != nil && apierrors.IsNotFound(err) {
				return true, nil
			}
			return false, nil
		}); err == nil {
			t.Fatalf("Unexpected valid lease getting garbage collected")
		}
		if _, err := client.CoordinationV1().Leases(ns).Get(context.TODO(), lease.Name, metav1.GetOptions{}); err != nil {
			t.Fatalf("Failed to retrieve valid lease: %v", err)
		}
		if err := client.CoordinationV1().Leases(ns).Delete(context.TODO(), lease.Name, metav1.DeleteOptions{}); err != nil {
			t.Fatalf("Failed to clean up valid lease: %v", err)
		}
	}
}

func newTestLease(acquireTime time.Time, namespace string) *coordinationv1.Lease {
	return &coordinationv1.Lease{
		ObjectMeta: metav1.ObjectMeta{
			Name:      testLeaseName,
			Namespace: namespace,
			Labels: map[string]string{
				controlplane.IdentityLeaseComponentLabelKey: controlplane.KubeAPIServer,
			},
		},
		Spec: coordinationv1.LeaseSpec{
			HolderIdentity:       pointer.StringPtr(testLeaseName),
			LeaseDurationSeconds: pointer.Int32Ptr(3600),
			AcquireTime:          &metav1.MicroTime{Time: acquireTime},
			RenewTime:            &metav1.MicroTime{Time: acquireTime},
		},
	}
}
<commit_msg>apiserver identity: update Lease creation integration test to validate new naming format and hostname label<commit_after>/*
Copyright 2020 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 controlplane

import (
	"context"
	"fmt"
	"hash/fnv"
	"os"
	"testing"
	"time"

	coordinationv1 "k8s.io/api/coordination/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apiserver/pkg/features"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/client-go/kubernetes"
	featuregatetesting "k8s.io/component-base/featuregate/testing"
	kubeapiservertesting "k8s.io/kubernetes/cmd/kube-apiserver/app/testing"
	"k8s.io/kubernetes/pkg/controlplane"
	"k8s.io/kubernetes/test/integration/framework"
	"k8s.io/utils/pointer"
)

const (
	testLeaseName = "apiserver-lease-test"
)

func expectedAPIServerIdentity(hostname string) string {
	h := fnv.New32a()
	h.Write([]byte(hostname))
	return "kube-apiserver-" + fmt.Sprint(h.Sum32())
}

func TestCreateLeaseOnStart(t *testing.T) {
	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.APIServerIdentity, true)()
	result := kubeapiservertesting.StartTestServerOrDie(t, nil, nil, framework.SharedEtcd())
	defer result.TearDownFn()

	kubeclient, err := kubernetes.NewForConfig(result.ClientConfig)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	hostname, err := os.Hostname()
	if err != nil {
		t.Fatalf("Unexpected error getting apiserver hostname: %v", err)
	}

	t.Logf(`Waiting the kube-apiserver Lease to be created`)
	if err := wait.PollImmediate(500*time.Millisecond, 10*time.Second, func() (bool, error) {
		leases, err := kubeclient.
			CoordinationV1().
			Leases(metav1.NamespaceSystem).
			List(context.TODO(), metav1.ListOptions{LabelSelector: controlplane.KubeAPIServerIdentityLeaseLabelSelector})
		if err != nil {
			return false, err
		}

		if leases == nil {
			return false, nil
		}

		if len(leases.Items) != 1 {
			return false, nil
		}

		lease := leases.Items[0]
		if lease.Name != expectedAPIServerIdentity(hostname) {
			return false, fmt.Errorf("unexpected apiserver identity, got: %v, expected: %v", lease.Name, expectedAPIServerIdentity(hostname))
		}

		if lease.Labels[corev1.LabelHostname] != hostname {
			return false, fmt.Errorf("unexpected hostname label, got: %v, expected: %v", lease.Labels[corev1.LabelHostname], hostname)
		}

		return true, nil
	}); err != nil {
		t.Fatalf("Failed to see the kube-apiserver lease: %v", err)
	}
}

func TestLeaseGarbageCollection(t *testing.T) {
	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.APIServerIdentity, true)()
	result := kubeapiservertesting.StartTestServerOrDie(t, nil,
		// This shorten the GC check period to make the test run faster.
		// Since we are testing GC behavior on leases we create, what happens to
		// the real apiserver lease doesn't matter.
		[]string{"--identity-lease-duration-seconds=1"},
		framework.SharedEtcd())
	defer result.TearDownFn()

	kubeclient, err := kubernetes.NewForConfig(result.ClientConfig)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	expiredLease := newTestLease(time.Now().Add(-2*time.Hour), metav1.NamespaceSystem)
	t.Run("expired apiserver lease should be garbage collected",
		testLeaseGarbageCollected(t, kubeclient, expiredLease))

	freshLease := newTestLease(time.Now().Add(-2*time.Minute), metav1.NamespaceSystem)
	t.Run("fresh apiserver lease should not be garbage collected",
		testLeaseNotGarbageCollected(t, kubeclient, freshLease))

	expiredLease.Labels = nil
	t.Run("expired non-identity lease should not be garbage collected",
		testLeaseNotGarbageCollected(t, kubeclient, expiredLease))

	// identity leases (with k8s.io/component label) created in user namespaces should not be GC'ed
	expiredNonKubeSystemLease := newTestLease(time.Now().Add(-2*time.Hour), metav1.NamespaceDefault)
	t.Run("expired non-system identity lease should not be garbage collected",
		testLeaseNotGarbageCollected(t, kubeclient, expiredNonKubeSystemLease))
}

func testLeaseGarbageCollected(t *testing.T, client kubernetes.Interface, lease *coordinationv1.Lease) func(t *testing.T) {
	return func(t *testing.T) {
		ns := lease.Namespace
		if _, err := client.CoordinationV1().Leases(ns).Create(context.TODO(), lease, metav1.CreateOptions{}); err != nil {
			t.Fatalf("Unexpected error creating lease: %v", err)
		}
		if err := wait.PollImmediate(500*time.Millisecond, 5*time.Second, func() (bool, error) {
			_, err := client.CoordinationV1().Leases(ns).Get(context.TODO(), lease.Name, metav1.GetOptions{})
			if err == nil {
				return false, nil
			}
			if apierrors.IsNotFound(err) {
				return true, nil
			}
			return false, err
		}); err != nil {
			t.Fatalf("Failed to see the expired lease garbage collected: %v", err)
		}

	}
}

func testLeaseNotGarbageCollected(t *testing.T, client kubernetes.Interface, lease *coordinationv1.Lease) func(t *testing.T) {
	return func(t *testing.T) {
		ns := lease.Namespace
		if _, err := client.CoordinationV1().Leases(ns).Create(context.TODO(), lease, metav1.CreateOptions{}); err != nil {
			t.Fatalf("Unexpected error creating lease: %v", err)
		}
		if err := wait.PollImmediate(500*time.Millisecond, 5*time.Second, func() (bool, error) {
			_, err := client.CoordinationV1().Leases(ns).Get(context.TODO(), lease.Name, metav1.GetOptions{})
			if err != nil && apierrors.IsNotFound(err) {
				return true, nil
			}
			return false, nil
		}); err == nil {
			t.Fatalf("Unexpected valid lease getting garbage collected")
		}
		if _, err := client.CoordinationV1().Leases(ns).Get(context.TODO(), lease.Name, metav1.GetOptions{}); err != nil {
			t.Fatalf("Failed to retrieve valid lease: %v", err)
		}
		if err := client.CoordinationV1().Leases(ns).Delete(context.TODO(), lease.Name, metav1.DeleteOptions{}); err != nil {
			t.Fatalf("Failed to clean up valid lease: %v", err)
		}
	}
}

func newTestLease(acquireTime time.Time, namespace string) *coordinationv1.Lease {
	return &coordinationv1.Lease{
		ObjectMeta: metav1.ObjectMeta{
			Name:      testLeaseName,
			Namespace: namespace,
			Labels: map[string]string{
				controlplane.IdentityLeaseComponentLabelKey: controlplane.KubeAPIServer,
			},
		},
		Spec: coordinationv1.LeaseSpec{
			HolderIdentity:       pointer.StringPtr(testLeaseName),
			LeaseDurationSeconds: pointer.Int32Ptr(3600),
			AcquireTime:          &metav1.MicroTime{Time: acquireTime},
			RenewTime:            &metav1.MicroTime{Time: acquireTime},
		},
	}
}
<|endoftext|>