<commit_before>package permissions

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/oauth"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/web/jsonapi"
	"github.com/labstack/echo"
	"github.com/stretchr/testify/assert"
	jwt "gopkg.in/dgrijalva/jwt-go.v3"
)

var ts *httptest.Server
var token string
var testInstance *instance.Instance
var clientID string

func TestMain(m *testing.M) {
	config.UseTestFile()

	testInstance = &instance.Instance{
		OAuthSecret: []byte("topsecret"),
		Domain:      "example.com",
	}

	err := couchdb.ResetDB(testInstance, "io.cozy.events")
	if err != nil {
		fmt.Println("Cant reset db", err)
		os.Exit(1)
	}

	err = couchdb.ResetDB(testInstance, consts.Permissions)
	if err != nil {
		fmt.Println("Cant reset db", err)
		os.Exit(1)
	}
	err = couchdb.DefineIndexes(testInstance, consts.Indexes)
	if err != nil {
		fmt.Println("Cant define index", err)
		os.Exit(1)
	}
	err = couchdb.DefineViews(testInstance, consts.Views)
	if err != nil {
		fmt.Println("cant define views", err)
		os.Exit(1)
	}

	client := oauth.Client{
		RedirectURIs: []string{"http://localhost/oauth/callback"},
		ClientName:   "test-permissions",
		SoftwareID:   "github.com/cozy/cozy-stack/web/permissions",
	}
	client.Create(testInstance)
	clientID = client.ClientID

	token, _ = crypto.NewJWT(testInstance.OAuthSecret, permissions.Claims{
		StandardClaims: jwt.StandardClaims{
			Audience: permissions.AccessTokenAudience,
			Issuer:   testInstance.Domain,
			IssuedAt: crypto.Timestamp(),
			Subject:  clientID,
		},
		Scope: "io.cozy.contacts io.cozy.files:GET",
	})

	handler := echo.New()

	group := handler.Group("/permissions")
	group.Use(injectInstance(testInstance))
	Routes(group)

	ts = httptest.NewServer(handler)
	res := m.Run()
	ts.Close()
	os.Exit(res)
}

func TestGetPermissions(t *testing.T) {
	req, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req.Header.Add("Authorization", "Bearer "+token)
	res, err := http.DefaultClient.Do(req)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if !assert.NoError(t, err) {
		return
	}
	var out map[string]interface{}
	err = json.Unmarshal(body, &out)

	assert.NoError(t, err)
	assert.Equal(t, "200 OK", res.Status, "should get a 200")
	for key, r := range out {
		rule := r.(map[string]interface{})
		if key == "rule0" {
			assert.Equal(t, "io.cozy.contacts", rule["type"])
		} else {
			assert.Equal(t, "io.cozy.files", rule["type"])
			assert.Equal(t, []interface{}{"GET"}, rule["verbs"])
		}
	}
}

func TestGetPermissionsForRevokedClient(t *testing.T) {
	tok, err := crypto.NewJWT(testInstance.OAuthSecret, permissions.Claims{
		StandardClaims: jwt.StandardClaims{
			Audience: permissions.AccessTokenAudience,
			Issuer:   testInstance.Domain,
			IssuedAt: crypto.Timestamp(),
			Subject:  "revoked-client",
		},
		Scope: "io.cozy.contacts io.cozy.files:GET",
	})
	assert.NoError(t, err)
	req, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req.Header.Add("Authorization", "Bearer "+tok)
	res, err := http.DefaultClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, res.StatusCode)
	body, err := ioutil.ReadAll(res.Body)
	assert.NoError(t, err)
	assert.Equal(t, `{"message":"Invalid JWT token"}`, string(body))
}

func TestGetPermissionsForExpiredToken(t *testing.T) {
	pastTimestamp := crypto.Timestamp() - 30*24*3600 // in seconds
	tok, err := crypto.NewJWT(testInstance.OAuthSecret, permissions.Claims{
		StandardClaims: jwt.StandardClaims{
			Audience: permissions.AccessTokenAudience,
			Issuer:   testInstance.Domain,
			IssuedAt: pastTimestamp,
			Subject:  clientID,
		},
		Scope: "io.cozy.contacts io.cozy.files:GET",
	})
	assert.NoError(t, err)
	req, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req.Header.Add("Authorization", "Bearer "+tok)
	res, err := http.DefaultClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, res.StatusCode)
	body, err := ioutil.ReadAll(res.Body)
	assert.NoError(t, err)
	assert.Equal(t, `{"message":"Expired token"}`, string(body))
}

func TestBadPermissionsBearer(t *testing.T) {
	req, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req.Header.Add("Authorization", "Bearer garbage")
	res, err := http.DefaultClient.Do(req)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	assert.Equal(t, res.StatusCode, http.StatusBadRequest)
}

func TestCreateSubPermission(t *testing.T) {
	reqbody := strings.NewReader(`{
"data": {
	"type": "io.cozy.permissions",
	"attributes": {
		"permissions": {
			"whatever": {
				"type":   "io.cozy.files",
				"verbs":  ["GET"],
				"values": ["io.cozy.music"]
			}
		}
	}
}
	}`)
	req, _ := http.NewRequest("POST", ts.URL+"/permissions?codes=bob,alice", reqbody)
	req.Header.Add("Authorization", "Bearer "+token)
	req.Header.Add("Content-Type", "application/json")
	res, err := http.DefaultClient.Do(req)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if !assert.NoError(t, err) {
		return
	}
	var out map[string]interface{}
	err = json.Unmarshal(body, &out)
	if !assert.NoError(t, err) {
		return
	}

	data := out["data"].(map[string]interface{})
	attrs := data["attributes"].(map[string]interface{})
	codes := attrs["codes"].(map[string]interface{})

	aCode := codes["alice"].(string)
	bCode := codes["bob"].(string)

	assert.NotEqual(t, aCode, token)
	assert.NotEqual(t, bCode, token)
	assert.NotEqual(t, aCode, bCode)

	req2, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req2.Header.Add("Authorization", "Bearer "+aCode)
	res2, err := http.DefaultClient.Do(req2)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	body2, err := ioutil.ReadAll(res2.Body)
	if !assert.NoError(t, err) {
		return
	}
	var out2 map[string]interface{}
	err = json.Unmarshal(body2, &out2)

	assert.NoError(t, err)
	assert.Equal(t, "200 OK", res2.Status, "should get a 200")
	assert.Len(t, out2, 1)
	assert.Equal(t, "io.cozy.files", out2["whatever"].(map[string]interface{})["type"])

}

func TestListPermission(t *testing.T) {

	ev1, _ := createTestEvent(testInstance)
	ev2, _ := createTestEvent(testInstance)
	ev3, _ := createTestEvent(testInstance)

	parent, _ := permissions.GetForOauth(&permissions.Claims{
		StandardClaims: jwt.StandardClaims{
			Audience: permissions.AccessTokenAudience,
			Issuer:   testInstance.Domain,
			IssuedAt: crypto.Timestamp(),
			Subject:  clientID,
		},
		Scope: "io.cozy.events",
	})
	p1 := permissions.Set{
		permissions.Rule{
			Type:   "io.cozy.events",
			Verbs:  permissions.Verbs(permissions.DELETE, permissions.PATCH),
			Values: []string{ev1.ID()},
		}}
	p2 := permissions.Set{
		permissions.Rule{
			Type:   "io.cozy.events",
			Verbs:  permissions.Verbs(permissions.GET),
			Values: []string{ev2.ID()},
		}}

	codes := map[string]string{"bob": "secret"}
	permissions.CreateShareSet(testInstance, parent, codes, &p1)
	permissions.CreateShareSet(testInstance, parent, codes, &p2)

	reqbody := strings.NewReader(`{
"data": [
{ "type": "io.cozy.events", "id": "` + ev1.ID() + `" },
{ "type": "io.cozy.events", "id": "` + ev2.ID() + `" },
{ "type": "io.cozy.events", "id": "non-existing-id" },
{ "type": "io.cozy.events", "id": "another-fake-id" },
{ "type": "io.cozy.events", "id": "` + ev3.ID() + `" }
]	}`)

	req, _ := http.NewRequest("POST", ts.URL+"/permissions/exists", reqbody)
	req.Header.Add("Authorization", "Bearer "+token)
	req.Header.Add("Content-Type", "application/json")
	res, err := http.DefaultClient.Do(req)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if !assert.NoError(t, err) {
		return
	}
	var out jsonapi.Document
	err = json.Unmarshal(body, &out)
	if !assert.NoError(t, err) {
		return
	}
	var results []refAndVerb
	err = json.Unmarshal(*out.Data, &results)
	if !assert.NoError(t, err) {
		return
	}

	assert.Len(t, results, 2)
	for _, result := range results {
		assert.Equal(t, "io.cozy.events", result.DocType)
		if result.ID == ev1.ID() {
			assert.Equal(t, "PATCH,DELETE", result.Verbs.String())
		} else {
			assert.Equal(t, ev2.ID(), result.ID)
			assert.Equal(t, "GET", result.Verbs.String())
		}
	}

}

func injectInstance(i *instance.Instance) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			c.Set("instance", i)
			return next(c)
		}
	}
}

func createTestEvent(i *instance.Instance) (*couchdb.JSONDoc, error) {
	e := &couchdb.JSONDoc{
		Type: "io.cozy.events",
		M:    map[string]interface{}{"test": "value"},
	}
	err := couchdb.CreateDoc(i, e)
	return e, err
}
<commit_msg>Fix tests for web/permissions<commit_after>package permissions

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/oauth"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/web/jsonapi"
	"github.com/labstack/echo"
	"github.com/stretchr/testify/assert"
	jwt "gopkg.in/dgrijalva/jwt-go.v3"
)

var ts *httptest.Server
var token string
var testInstance *instance.Instance
var clientID string

func TestMain(m *testing.M) {
	config.UseTestFile()

	testInstance = &instance.Instance{
		OAuthSecret: []byte("topsecret"),
		Domain:      "example.com",
	}

	err := couchdb.ResetDB(testInstance, "io.cozy.events")
	if err != nil {
		fmt.Println("Cant reset db", err)
		os.Exit(1)
	}

	err = couchdb.ResetDB(testInstance, consts.Permissions)
	if err != nil {
		fmt.Println("Cant reset db", err)
		os.Exit(1)
	}
	err = couchdb.DefineIndexes(testInstance, consts.IndexesByDoctype(consts.Permissions))
	if err != nil {
		fmt.Println("Cant define index", err)
		os.Exit(1)
	}
	err = couchdb.DefineViews(testInstance, consts.ViewsByDoctype(consts.Permissions))
	if err != nil {
		fmt.Println("cant define views", err)
		os.Exit(1)
	}

	client := oauth.Client{
		RedirectURIs: []string{"http://localhost/oauth/callback"},
		ClientName:   "test-permissions",
		SoftwareID:   "github.com/cozy/cozy-stack/web/permissions",
	}
	client.Create(testInstance)
	clientID = client.ClientID

	token, _ = crypto.NewJWT(testInstance.OAuthSecret, permissions.Claims{
		StandardClaims: jwt.StandardClaims{
			Audience: permissions.AccessTokenAudience,
			Issuer:   testInstance.Domain,
			IssuedAt: crypto.Timestamp(),
			Subject:  clientID,
		},
		Scope: "io.cozy.contacts io.cozy.files:GET",
	})

	handler := echo.New()

	group := handler.Group("/permissions")
	group.Use(injectInstance(testInstance))
	Routes(group)

	ts = httptest.NewServer(handler)
	res := m.Run()
	ts.Close()
	os.Exit(res)
}

func TestGetPermissions(t *testing.T) {
	req, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req.Header.Add("Authorization", "Bearer "+token)
	res, err := http.DefaultClient.Do(req)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if !assert.NoError(t, err) {
		return
	}
	var out map[string]interface{}
	err = json.Unmarshal(body, &out)

	assert.NoError(t, err)
	assert.Equal(t, "200 OK", res.Status, "should get a 200")
	for key, r := range out {
		rule := r.(map[string]interface{})
		if key == "rule0" {
			assert.Equal(t, "io.cozy.contacts", rule["type"])
		} else {
			assert.Equal(t, "io.cozy.files", rule["type"])
			assert.Equal(t, []interface{}{"GET"}, rule["verbs"])
		}
	}
}

func TestGetPermissionsForRevokedClient(t *testing.T) {
	tok, err := crypto.NewJWT(testInstance.OAuthSecret, permissions.Claims{
		StandardClaims: jwt.StandardClaims{
			Audience: permissions.AccessTokenAudience,
			Issuer:   testInstance.Domain,
			IssuedAt: crypto.Timestamp(),
			Subject:  "revoked-client",
		},
		Scope: "io.cozy.contacts io.cozy.files:GET",
	})
	assert.NoError(t, err)
	req, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req.Header.Add("Authorization", "Bearer "+tok)
	res, err := http.DefaultClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, res.StatusCode)
	body, err := ioutil.ReadAll(res.Body)
	assert.NoError(t, err)
	assert.Equal(t, `{"message":"Invalid JWT token"}`, string(body))
}

func TestGetPermissionsForExpiredToken(t *testing.T) {
	pastTimestamp := crypto.Timestamp() - 30*24*3600 // in seconds
	tok, err := crypto.NewJWT(testInstance.OAuthSecret, permissions.Claims{
		StandardClaims: jwt.StandardClaims{
			Audience: permissions.AccessTokenAudience,
			Issuer:   testInstance.Domain,
			IssuedAt: pastTimestamp,
			Subject:  clientID,
		},
		Scope: "io.cozy.contacts io.cozy.files:GET",
	})
	assert.NoError(t, err)
	req, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req.Header.Add("Authorization", "Bearer "+tok)
	res, err := http.DefaultClient.Do(req)
	assert.NoError(t, err)
	assert.Equal(t, 400, res.StatusCode)
	body, err := ioutil.ReadAll(res.Body)
	assert.NoError(t, err)
	assert.Equal(t, `{"message":"Expired token"}`, string(body))
}

func TestBadPermissionsBearer(t *testing.T) {
	req, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req.Header.Add("Authorization", "Bearer garbage")
	res, err := http.DefaultClient.Do(req)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	assert.Equal(t, res.StatusCode, http.StatusBadRequest)
}

func TestCreateSubPermission(t *testing.T) {
	reqbody := strings.NewReader(`{
"data": {
	"type": "io.cozy.permissions",
	"attributes": {
		"permissions": {
			"whatever": {
				"type":   "io.cozy.files",
				"verbs":  ["GET"],
				"values": ["io.cozy.music"]
			}
		}
	}
}
	}`)
	req, _ := http.NewRequest("POST", ts.URL+"/permissions?codes=bob,alice", reqbody)
	req.Header.Add("Authorization", "Bearer "+token)
	req.Header.Add("Content-Type", "application/json")
	res, err := http.DefaultClient.Do(req)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if !assert.NoError(t, err) {
		return
	}
	var out map[string]interface{}
	err = json.Unmarshal(body, &out)
	if !assert.NoError(t, err) {
		return
	}

	data := out["data"].(map[string]interface{})
	attrs := data["attributes"].(map[string]interface{})
	codes := attrs["codes"].(map[string]interface{})

	aCode := codes["alice"].(string)
	bCode := codes["bob"].(string)

	assert.NotEqual(t, aCode, token)
	assert.NotEqual(t, bCode, token)
	assert.NotEqual(t, aCode, bCode)

	req2, _ := http.NewRequest("GET", ts.URL+"/permissions/self", nil)
	req2.Header.Add("Authorization", "Bearer "+aCode)
	res2, err := http.DefaultClient.Do(req2)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	body2, err := ioutil.ReadAll(res2.Body)
	if !assert.NoError(t, err) {
		return
	}
	var out2 map[string]interface{}
	err = json.Unmarshal(body2, &out2)

	assert.NoError(t, err)
	assert.Equal(t, "200 OK", res2.Status, "should get a 200")
	assert.Len(t, out2, 1)
	assert.Equal(t, "io.cozy.files", out2["whatever"].(map[string]interface{})["type"])

}

func TestListPermission(t *testing.T) {

	ev1, _ := createTestEvent(testInstance)
	ev2, _ := createTestEvent(testInstance)
	ev3, _ := createTestEvent(testInstance)

	parent, _ := permissions.GetForOauth(&permissions.Claims{
		StandardClaims: jwt.StandardClaims{
			Audience: permissions.AccessTokenAudience,
			Issuer:   testInstance.Domain,
			IssuedAt: crypto.Timestamp(),
			Subject:  clientID,
		},
		Scope: "io.cozy.events",
	})
	p1 := permissions.Set{
		permissions.Rule{
			Type:   "io.cozy.events",
			Verbs:  permissions.Verbs(permissions.DELETE, permissions.PATCH),
			Values: []string{ev1.ID()},
		}}
	p2 := permissions.Set{
		permissions.Rule{
			Type:   "io.cozy.events",
			Verbs:  permissions.Verbs(permissions.GET),
			Values: []string{ev2.ID()},
		}}

	codes := map[string]string{"bob": "secret"}
	permissions.CreateShareSet(testInstance, parent, codes, &p1)
	permissions.CreateShareSet(testInstance, parent, codes, &p2)

	reqbody := strings.NewReader(`{
"data": [
{ "type": "io.cozy.events", "id": "` + ev1.ID() + `" },
{ "type": "io.cozy.events", "id": "` + ev2.ID() + `" },
{ "type": "io.cozy.events", "id": "non-existing-id" },
{ "type": "io.cozy.events", "id": "another-fake-id" },
{ "type": "io.cozy.events", "id": "` + ev3.ID() + `" }
]	}`)

	req, _ := http.NewRequest("POST", ts.URL+"/permissions/exists", reqbody)
	req.Header.Add("Authorization", "Bearer "+token)
	req.Header.Add("Content-Type", "application/json")
	res, err := http.DefaultClient.Do(req)
	if !assert.NoError(t, err) {
		return
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if !assert.NoError(t, err) {
		return
	}
	var out jsonapi.Document
	err = json.Unmarshal(body, &out)
	if !assert.NoError(t, err) {
		return
	}
	var results []refAndVerb
	err = json.Unmarshal(*out.Data, &results)
	if !assert.NoError(t, err) {
		return
	}

	assert.Len(t, results, 2)
	for _, result := range results {
		assert.Equal(t, "io.cozy.events", result.DocType)
		if result.ID == ev1.ID() {
			assert.Equal(t, "PATCH,DELETE", result.Verbs.String())
		} else {
			assert.Equal(t, ev2.ID(), result.ID)
			assert.Equal(t, "GET", result.Verbs.String())
		}
	}

}

func injectInstance(i *instance.Instance) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			c.Set("instance", i)
			return next(c)
		}
	}
}

func createTestEvent(i *instance.Instance) (*couchdb.JSONDoc, error) {
	e := &couchdb.JSONDoc{
		Type: "io.cozy.events",
		M:    map[string]interface{}{"test": "value"},
	}
	err := couchdb.CreateDoc(i, e)
	return e, err
}
<|endoftext|>


<commit_before>package sftp

import (
	"encoding"
	"io"
	"sync"

	"github.com/pkg/errors"
)

// conn implements a bidirectional channel on which client and server
// connections are multiplexed.
type conn struct {
	io.Reader
	io.WriteCloser
	// this is the same allocator used in packet manager
	alloc      *allocator
	sync.Mutex // used to serialise writes to sendPacket
	// sendPacketTest is needed to replicate packet issues in testing
	sendPacketTest func(w io.Writer, m encoding.BinaryMarshaler) error
}

// the orderID is used in server mode if the allocator is enabled.
// For the client mode just pass 0
func (c *conn) recvPacket(orderID uint32) (uint8, []byte, error) {
	return recvPacket(c, c.alloc, orderID)
}

func (c *conn) sendPacket(m encoding.BinaryMarshaler) error {
	c.Lock()
	defer c.Unlock()
	if c.sendPacketTest != nil {
		return c.sendPacketTest(c, m)
	}
	return sendPacket(c, m)
}

func (c *conn) Close() error {
	c.Lock()
	defer c.Unlock()
	return c.WriteCloser.Close()
}

type clientConn struct {
	conn
	wg sync.WaitGroup

	sync.Mutex                          // protects inflight
	inflight   map[uint32]chan<- result // outstanding requests

	closed chan struct{}
	err    error
}

// Wait blocks until the conn has shut down, and return the error
// causing the shutdown. It can be called concurrently from multiple
// goroutines.
func (c *clientConn) Wait() error {
	<-c.closed
	return c.err
}

// Close closes the SFTP session.
func (c *clientConn) Close() error {
	defer c.wg.Wait()
	return c.conn.Close()
}

func (c *clientConn) loop() {
	defer c.wg.Done()
	err := c.recv()
	if err != nil {
		c.broadcastErr(err)
	}
}

// recv continuously reads from the server and forwards responses to the
// appropriate channel.
func (c *clientConn) recv() error {
	defer c.conn.Close()

	for {
		typ, data, err := c.recvPacket(0)
		if err != nil {
			return err
		}
		sid, _ := unmarshalUint32(data)

		ch, ok := c.getChannel(sid)
		if !ok {
			// This is an unexpected occurrence. Send the error
			// back to all listeners so that they terminate
			// gracefully.
			return errors.Errorf("sid: %v not fond", sid)
		}

		ch <- result{typ: typ, data: data}
	}
}

func (c *clientConn) putChannel(ch chan<- result, sid uint32) bool {
	c.Lock()
	defer c.Unlock()

	select {
	case <-c.closed:
		// already closed with broadcastErr, return error on chan.
		ch <- result{err: c.err}
		return false
	default:
	}

	c.inflight[sid] = ch
	return true
}

func (c *clientConn) getChannel(sid uint32) (chan<- result, bool) {
	c.Lock()
	defer c.Unlock()

	ch, ok := c.inflight[sid]
	delete(c.inflight, sid)

	return ch, ok
}

// result captures the result of receiving the a packet from the server
type result struct {
	typ  byte
	data []byte
	err  error
}

type idmarshaler interface {
	id() uint32
	encoding.BinaryMarshaler
}

func (c *clientConn) sendPacket(ch chan result, p idmarshaler) (byte, []byte, error) {
	if cap(ch) < 1 {
		ch = make(chan result, 1)
	}

	c.dispatchRequest(ch, p)
	s := <-ch
	return s.typ, s.data, s.err
}

// dispatchRequest should ideally only be called by race-detection tests outside of this file,
// where you have to ensure two packets are in flight sequentially after each other.
func (c *clientConn) dispatchRequest(ch chan<- result, p idmarshaler) {
	sid := p.id()

	if !c.putChannel(ch, sid) {
		// already closed.
		return
	}

	if err := c.conn.sendPacket(p); err != nil {
		if ch, ok := c.getChannel(sid); ok {
			ch <- result{err: err}
		}
	}
}

// broadcastErr sends an error to all goroutines waiting for a response.
func (c *clientConn) broadcastErr(err error) {
	c.Lock()
	defer c.Unlock()

	bcastRes := result{err: errors.New("unexpected server disconnect")}
	for sid, ch := range c.inflight {
		ch <- bcastRes

		// Replace the chan in inflight,
		// we have hijacked this chan,
		// and this guarantees always-only-once sending.
		c.inflight[sid] = make(chan<- result, 1)
	}

	c.err = err
	close(c.closed)
}

type serverConn struct {
	conn
}

func (s *serverConn) sendError(p ider, err error) error {
	return s.sendPacket(statusFromError(p, err))
}
<commit_msg>a new server disconnect error point rather than EOF<commit_after>package sftp

import (
	"encoding"
	"io"
	"sync"

	"github.com/pkg/errors"
)

// conn implements a bidirectional channel on which client and server
// connections are multiplexed.
type conn struct {
	io.Reader
	io.WriteCloser
	// this is the same allocator used in packet manager
	alloc      *allocator
	sync.Mutex // used to serialise writes to sendPacket
	// sendPacketTest is needed to replicate packet issues in testing
	sendPacketTest func(w io.Writer, m encoding.BinaryMarshaler) error
}

// the orderID is used in server mode if the allocator is enabled.
// For the client mode just pass 0
func (c *conn) recvPacket(orderID uint32) (uint8, []byte, error) {
	return recvPacket(c, c.alloc, orderID)
}

func (c *conn) sendPacket(m encoding.BinaryMarshaler) error {
	c.Lock()
	defer c.Unlock()
	if c.sendPacketTest != nil {
		return c.sendPacketTest(c, m)
	}
	return sendPacket(c, m)
}

func (c *conn) Close() error {
	c.Lock()
	defer c.Unlock()
	return c.WriteCloser.Close()
}

type clientConn struct {
	conn
	wg sync.WaitGroup

	sync.Mutex                          // protects inflight
	inflight   map[uint32]chan<- result // outstanding requests

	closed chan struct{}
	err    error
}

// Wait blocks until the conn has shut down, and return the error
// causing the shutdown. It can be called concurrently from multiple
// goroutines.
func (c *clientConn) Wait() error {
	<-c.closed
	return c.err
}

// Close closes the SFTP session.
func (c *clientConn) Close() error {
	defer c.wg.Wait()
	return c.conn.Close()
}

func (c *clientConn) loop() {
	defer c.wg.Done()
	err := c.recv()
	if err != nil {
		c.broadcastErr(err)
	}
}

// recv continuously reads from the server and forwards responses to the
// appropriate channel.
func (c *clientConn) recv() error {
	defer c.conn.Close()

	for {
		typ, data, err := c.recvPacket(0)
		if err != nil {
			return err
		}
		sid, _ := unmarshalUint32(data)

		ch, ok := c.getChannel(sid)
		if !ok {
			// This is an unexpected occurrence. Send the error
			// back to all listeners so that they terminate
			// gracefully.
			return errors.Errorf("sid: %v not fond", sid)
		}

		ch <- result{typ: typ, data: data}
	}
}

func (c *clientConn) putChannel(ch chan<- result, sid uint32) bool {
	c.Lock()
	defer c.Unlock()

	select {
	case <-c.closed:
		// already closed with broadcastErr, return error on chan.
		ch <- result{err: errors.New("unexpected server disconnect")}
		return false
	default:
	}

	c.inflight[sid] = ch
	return true
}

func (c *clientConn) getChannel(sid uint32) (chan<- result, bool) {
	c.Lock()
	defer c.Unlock()

	ch, ok := c.inflight[sid]
	delete(c.inflight, sid)

	return ch, ok
}

// result captures the result of receiving the a packet from the server
type result struct {
	typ  byte
	data []byte
	err  error
}

type idmarshaler interface {
	id() uint32
	encoding.BinaryMarshaler
}

func (c *clientConn) sendPacket(ch chan result, p idmarshaler) (byte, []byte, error) {
	if cap(ch) < 1 {
		ch = make(chan result, 1)
	}

	c.dispatchRequest(ch, p)
	s := <-ch
	return s.typ, s.data, s.err
}

// dispatchRequest should ideally only be called by race-detection tests outside of this file,
// where you have to ensure two packets are in flight sequentially after each other.
func (c *clientConn) dispatchRequest(ch chan<- result, p idmarshaler) {
	sid := p.id()

	if !c.putChannel(ch, sid) {
		// already closed.
		return
	}

	if err := c.conn.sendPacket(p); err != nil {
		if ch, ok := c.getChannel(sid); ok {
			ch <- result{err: err}
		}
	}
}

// broadcastErr sends an error to all goroutines waiting for a response.
func (c *clientConn) broadcastErr(err error) {
	c.Lock()
	defer c.Unlock()

	bcastRes := result{err: errors.New("unexpected server disconnect")}
	for sid, ch := range c.inflight {
		ch <- bcastRes

		// Replace the chan in inflight,
		// we have hijacked this chan,
		// and this guarantees always-only-once sending.
		c.inflight[sid] = make(chan<- result, 1)
	}

	c.err = err
	close(c.closed)
}

type serverConn struct {
	conn
}

func (s *serverConn) sendError(p ider, err error) error {
	return s.sendPacket(statusFromError(p, err))
}
<|endoftext|>


<commit_before>package netlink

import (
	"errors"
	"io"
	"math/rand"
	"os"
	"sync/atomic"

	"golang.org/x/net/bpf"
)

// Error messages which can be returned by Validate.
var (
	errMismatchedSequence = errors.New("mismatched sequence in netlink reply")
	errMismatchedPID      = errors.New("mismatched PID in netlink reply")
	errShortErrorMessage  = errors.New("not enough data for netlink error code")
)

// Errors which can be returned by a Socket that does not implement
// all exposed methods of Conn.
var (
	errReadWriteCloserNotSupported = errors.New("raw read/write/closer not supported")
	errMulticastGroupsNotSupported = errors.New("multicast groups not supported")
	errBPFFiltersNotSupported      = errors.New("BPF filters not supported")
	errOptionsNotSupported         = errors.New("options not supported")
)

// A Conn is a connection to netlink.  A Conn can be used to send and
// receives messages to and from netlink.
//
// A Conn is safe for concurrent use, but to avoid contention in
// high-throughput applications, the caller should almost certainly create a
// pool of Conns and distribute them among workers.
type Conn struct {
	// sock is the operating system-specific implementation of
	// a netlink sockets connection.
	sock Socket

	// seq is an atomically incremented integer used to provide sequence
	// numbers when Conn.Send is called.
	seq *uint32

	// pid is the PID assigned by netlink.
	pid uint32

	// d provides debugging capabilities for a Conn if not nil.
	d *debugger
}

// A Socket is an operating-system specific implementation of netlink
// sockets used by Conn.
type Socket interface {
	Close() error
	Send(m Message) error
	SendMessages(m []Message) error
	Receive() ([]Message, error)
}

// Dial dials a connection to netlink, using the specified netlink family.
// Config specifies optional configuration for Conn.  If config is nil, a default
// configuration will be used.
func Dial(family int, config *Config) (*Conn, error) {
	// Use OS-specific dial() to create Socket
	c, pid, err := dial(family, config)
	if err != nil {
		return nil, err
	}

	return NewConn(c, pid), nil
}

// NewConn creates a Conn using the specified Socket and PID for netlink
// communications.
//
// NewConn is primarily useful for tests. Most applications should use
// Dial instead.
func NewConn(c Socket, pid uint32) *Conn {
	seq := rand.Uint32()

	// Configure a debugger if arguments are set.
	var d *debugger
	if len(debugArgs) > 0 {
		d = newDebugger(debugArgs)
	}

	return &Conn{
		sock: c,
		seq:  &seq,
		pid:  pid,
		d:    d,
	}
}

// debug executes fn with the debugger if the debugger is not nil.
func (c *Conn) debug(fn func(d *debugger)) {
	if c.d == nil {
		return
	}

	fn(c.d)
}

// Close closes the connection.
func (c *Conn) Close() error {
	return c.sock.Close()
}

// Execute sends a single Message to netlink using Conn.Send, receives one or more
// replies using Conn.Receive, and then checks the validity of the replies against
// the request using Validate.
//
// See the documentation of Conn.Send, Conn.Receive, and Validate for details about
// each function.
func (c *Conn) Execute(m Message) ([]Message, error) {
	req, err := c.Send(m)
	if err != nil {
		return nil, err
	}

	replies, err := c.Receive()
	if err != nil {
		return nil, err
	}

	if err := Validate(req, replies); err != nil {
		return nil, err
	}

	return replies, nil
}

func (c *Conn) fixMsg(m *Message, ml int) {
	if m.Header.Length == 0 {
		m.Header.Length = uint32(nlmsgAlign(ml))
	}

	if m.Header.Sequence == 0 {
		m.Header.Sequence = c.nextSequence()
	}

	if m.Header.PID == 0 {
		m.Header.PID = c.pid
	}
}

// SendMessages sends multiple Messages to netlink. The handling of
// m.Header.Length, Sequence and PID is the same as when calling Send.
func (c *Conn) SendMessages(messages []Message) ([]Message, error) {
	for idx, m := range messages {
		ml := nlmsgLength(len(m.Data))

		// TODO(mdlayher): fine-tune this limit.
		if ml > (1024 * 32) {
			return nil, errors.New("netlink message data too large")
		}

		c.fixMsg(&messages[idx], ml)
	}

	c.debug(func(d *debugger) {
		for _, m := range messages {
			d.debugf(1, "send msgs: %+v", m)
		}
	})

	if err := c.sock.SendMessages(messages); err != nil {
		c.debug(func(d *debugger) {
			d.debugf(1, "send msgs: err: %v", err)
		})

		return nil, err
	}

	return messages, nil
}

// Send sends a single Message to netlink.  In most cases, m.Header's Length,
// Sequence, and PID fields should be set to 0, so they can be populated
// automatically before the Message is sent.  On success, Send returns a copy
// of the Message with all parameters populated, for later validation.
//
// If m.Header.Length is 0, it will be automatically populated using the
// correct length for the Message, including its payload.
//
// If m.Header.Sequence is 0, it will be automatically populated using the
// next sequence number for this connection.
//
// If m.Header.PID is 0, it will be automatically populated using a PID
// assigned by netlink.
func (c *Conn) Send(m Message) (Message, error) {
	ml := nlmsgLength(len(m.Data))

	// TODO(mdlayher): fine-tune this limit.
	if ml > (1024 * 32) {
		return Message{}, errors.New("netlink message data too large")
	}

	c.fixMsg(&m, ml)

	c.debug(func(d *debugger) {
		d.debugf(1, "send: %+v", m)
	})

	if err := c.sock.Send(m); err != nil {
		c.debug(func(d *debugger) {
			d.debugf(1, "send: err: %v", err)
		})

		return Message{}, err
	}

	return m, nil
}

// Receive receives one or more messages from netlink.  Multi-part messages are
// handled transparently and returned as a single slice of Messages, with the
// final empty "multi-part done" message removed.
//
// If any of the messages indicate a netlink error, that error will be returned.
func (c *Conn) Receive() ([]Message, error) {
	msgs, err := c.receive()
	if err != nil {
		c.debug(func(d *debugger) {
			d.debugf(1, "recv: err: %v", err)
		})

		return nil, err
	}

	c.debug(func(d *debugger) {
		for _, m := range msgs {
			d.debugf(1, "recv: %+v", m)
		}
	})

	// When using nltest, it's possible for zero messages to be returned by receive.
	if len(msgs) == 0 {
		return msgs, nil
	}

	// Trim the final message with multi-part done indicator if
	// present.
	if m := msgs[len(msgs)-1]; m.Header.Flags&HeaderFlagsMulti != 0 && m.Header.Type == HeaderTypeDone {
		return msgs[:len(msgs)-1], nil
	}

	return msgs, nil
}

// receive is the internal implementation of Conn.Receive, which can be called
// recursively to handle multi-part messages.
func (c *Conn) receive() ([]Message, error) {
	var res []Message
	for {
		msgs, err := c.sock.Receive()
		if err != nil {
			return nil, err
		}

		// If this message is multi-part, we will need to perform an recursive call
		// to continue draining the socket
		var multi bool

		for _, m := range msgs {
			if err := checkMessage(m); err != nil {
				return nil, err
			}

			// Does this message indicate a multi-part message?
			if m.Header.Flags&HeaderFlagsMulti == 0 {
				// No, check the next messages.
				continue
			}

			// Does this message indicate the last message in a series of
			// multi-part messages from a single read?
			multi = m.Header.Type != HeaderTypeDone
		}

		res = append(res, msgs...)

		if !multi {
			// No more messages coming.
			return res, nil
		}
	}
}

// An fder is a Socket that supports retrieving its raw file descriptor.
type fder interface {
	Socket
	FD() int
}

var _ io.ReadWriteCloser = &fileReadWriteCloser{}

// A fileReadWriteCloser is a limited *os.File which only allows access to its
// Read and Write methods.
type fileReadWriteCloser struct {
	f *os.File
}

// Read implements io.ReadWriteCloser.
func (rwc *fileReadWriteCloser) Read(b []byte) (int, error) { return rwc.f.Read(b) }

// Write implements io.ReadWriteCloser.
func (rwc *fileReadWriteCloser) Write(b []byte) (int, error) { return rwc.f.Write(b) }

// Close implements io.ReadWriteCloser.
func (rwc *fileReadWriteCloser) Close() error { return rwc.f.Close() }

// ReadWriteCloser returns a raw io.ReadWriteCloser backed by the connection
// of the Conn.
//
// ReadWriteCloser is intended for advanced use cases, such as those that do
// not involve standard netlink message passing.
//
// Once invoked, it is the caller's responsibility to ensure that operations
// performed using Conn and the raw io.ReadWriteCloser do not conflict with
// each other.  In almost all scenarios, only one of the two should be used.
func (c *Conn) ReadWriteCloser() (io.ReadWriteCloser, error) {
	fc, ok := c.sock.(fder)
	if !ok {
		return nil, errReadWriteCloserNotSupported
	}

	return &fileReadWriteCloser{
		// Backing the io.ReadWriteCloser with an *os.File enables easy reading
		// and writing without more system call boilerplate.
		f: os.NewFile(uintptr(fc.FD()), "netlink"),
	}, nil
}

// A groupJoinLeaver is a Socket that supports joining and leaving
// netlink multicast groups.
type groupJoinLeaver interface {
	Socket
	JoinGroup(group uint32) error
	LeaveGroup(group uint32) error
}

// JoinGroup joins a netlink multicast group by its ID.
func (c *Conn) JoinGroup(group uint32) error {
	gc, ok := c.sock.(groupJoinLeaver)
	if !ok {
		return errMulticastGroupsNotSupported
	}

	return gc.JoinGroup(group)
}

// LeaveGroup leaves a netlink multicast group by its ID.
func (c *Conn) LeaveGroup(group uint32) error {
	gc, ok := c.sock.(groupJoinLeaver)
	if !ok {
		return errMulticastGroupsNotSupported
	}

	return gc.LeaveGroup(group)
}

// A bpfSetter is a Socket that supports setting and removing BPF filters.
type bpfSetter interface {
	Socket
	bpf.Setter
	RemoveBPF() error
}

// SetBPF attaches an assembled BPF program to a Conn.
func (c *Conn) SetBPF(filter []bpf.RawInstruction) error {
	bc, ok := c.sock.(bpfSetter)
	if !ok {
		return errBPFFiltersNotSupported
	}

	return bc.SetBPF(filter)
}

// RemoveBPF removes a BPF filter from a Conn.
func (c *Conn) RemoveBPF() error {
	s, ok := c.sock.(bpfSetter)
	if !ok {
		return errBPFFiltersNotSupported
	}

	return s.RemoveBPF()
}

// A ConnOption is a boolean option that may be set for a Conn.
type ConnOption int

// Possible ConnOption values.  These constants are equivalent to the Linux
// setsockopt boolean options for netlink sockets.
const (
	PacketInfo ConnOption = iota
	BroadcastError
	NoENOBUFS
	ListenAllNSID
	CapAcknowledge
)

// An optionSetter is a Socket that supports setting netlink options.
type optionSetter interface {
	Socket
	SetOption(option ConnOption, enable bool) error
}

// SetOption enables or disables a netlink socket option for the Conn.
func (c *Conn) SetOption(option ConnOption, enable bool) error {
	fc, ok := c.sock.(optionSetter)
	if !ok {
		return errOptionsNotSupported
	}

	return fc.SetOption(option, enable)
}

// nextSequence atomically increments Conn's sequence number and returns
// the incremented value.
func (c *Conn) nextSequence() uint32 {
	return atomic.AddUint32(c.seq, 1)
}

// Validate validates one or more reply Messages against a request Message,
// ensuring that they contain matching sequence numbers and PIDs.
func Validate(request Message, replies []Message) error {
	for _, m := range replies {
		// Check for mismatched sequence, unless:
		//   - request had no sequence, meaning we are probably validating
		//     a multicast reply
		if m.Header.Sequence != request.Header.Sequence && request.Header.Sequence != 0 {
			return errMismatchedSequence
		}

		// Check for mismatched PID, unless:
		//   - request had no PID, meaning we are either:
		//     - validating a multicast reply
		//     - netlink has not yet assigned us a PID
		//   - response had no PID, meaning it's from the kernel as a multicast reply
		if m.Header.PID != request.Header.PID && request.Header.PID != 0 && m.Header.PID != 0 {
			return errMismatchedPID
		}
	}

	return nil
}

// Config contains options for a Conn.
type Config struct {
	// Groups is a bitmask which specifies multicast groups. If set to 0,
	// no multicast group subscriptions will be made.
	Groups uint32

	// Experimental: do not lock the internal system call handling goroutine
	// to its OS thread.  This may result in a speed-up of system call handling,
	// but may cause unexpected behavior when sending and receiving a large number
	// of messages.
	//
	// This should almost certainly be set to false, but if you come up with a
	// valid reason for using this, please file an issue at
	// https://github.com/mdlayher/netlink to discuss your thoughts.
	NoLockThread bool
}
<commit_msg>netlink: ensure random sequence numbers by seeding local generator<commit_after>package netlink

import (
	"errors"
	"io"
	"math/rand"
	"os"
	"sync/atomic"
	"time"

	"golang.org/x/net/bpf"
)

// Error messages which can be returned by Validate.
var (
	errMismatchedSequence = errors.New("mismatched sequence in netlink reply")
	errMismatchedPID      = errors.New("mismatched PID in netlink reply")
	errShortErrorMessage  = errors.New("not enough data for netlink error code")
)

// Errors which can be returned by a Socket that does not implement
// all exposed methods of Conn.
var (
	errReadWriteCloserNotSupported = errors.New("raw read/write/closer not supported")
	errMulticastGroupsNotSupported = errors.New("multicast groups not supported")
	errBPFFiltersNotSupported      = errors.New("BPF filters not supported")
	errOptionsNotSupported         = errors.New("options not supported")
)

// A Conn is a connection to netlink.  A Conn can be used to send and
// receives messages to and from netlink.
//
// A Conn is safe for concurrent use, but to avoid contention in
// high-throughput applications, the caller should almost certainly create a
// pool of Conns and distribute them among workers.
type Conn struct {
	// sock is the operating system-specific implementation of
	// a netlink sockets connection.
	sock Socket

	// seq is an atomically incremented integer used to provide sequence
	// numbers when Conn.Send is called.
	seq *uint32

	// pid is the PID assigned by netlink.
	pid uint32

	// d provides debugging capabilities for a Conn if not nil.
	d *debugger
}

// A Socket is an operating-system specific implementation of netlink
// sockets used by Conn.
type Socket interface {
	Close() error
	Send(m Message) error
	SendMessages(m []Message) error
	Receive() ([]Message, error)
}

// Dial dials a connection to netlink, using the specified netlink family.
// Config specifies optional configuration for Conn.  If config is nil, a default
// configuration will be used.
func Dial(family int, config *Config) (*Conn, error) {
	// Use OS-specific dial() to create Socket
	c, pid, err := dial(family, config)
	if err != nil {
		return nil, err
	}

	return NewConn(c, pid), nil
}

// NewConn creates a Conn using the specified Socket and PID for netlink
// communications.
//
// NewConn is primarily useful for tests. Most applications should use
// Dial instead.
func NewConn(c Socket, pid uint32) *Conn {
	// Seed the sequence number using a random number generator.
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	seq := r.Uint32()

	// Configure a debugger if arguments are set.
	var d *debugger
	if len(debugArgs) > 0 {
		d = newDebugger(debugArgs)
	}

	return &Conn{
		sock: c,
		seq:  &seq,
		pid:  pid,
		d:    d,
	}
}

// debug executes fn with the debugger if the debugger is not nil.
func (c *Conn) debug(fn func(d *debugger)) {
	if c.d == nil {
		return
	}

	fn(c.d)
}

// Close closes the connection.
func (c *Conn) Close() error {
	return c.sock.Close()
}

// Execute sends a single Message to netlink using Conn.Send, receives one or more
// replies using Conn.Receive, and then checks the validity of the replies against
// the request using Validate.
//
// See the documentation of Conn.Send, Conn.Receive, and Validate for details about
// each function.
func (c *Conn) Execute(m Message) ([]Message, error) {
	req, err := c.Send(m)
	if err != nil {
		return nil, err
	}

	replies, err := c.Receive()
	if err != nil {
		return nil, err
	}

	if err := Validate(req, replies); err != nil {
		return nil, err
	}

	return replies, nil
}

func (c *Conn) fixMsg(m *Message, ml int) {
	if m.Header.Length == 0 {
		m.Header.Length = uint32(nlmsgAlign(ml))
	}

	if m.Header.Sequence == 0 {
		m.Header.Sequence = c.nextSequence()
	}

	if m.Header.PID == 0 {
		m.Header.PID = c.pid
	}
}

// SendMessages sends multiple Messages to netlink. The handling of
// m.Header.Length, Sequence and PID is the same as when calling Send.
func (c *Conn) SendMessages(messages []Message) ([]Message, error) {
	for idx, m := range messages {
		ml := nlmsgLength(len(m.Data))

		// TODO(mdlayher): fine-tune this limit.
		if ml > (1024 * 32) {
			return nil, errors.New("netlink message data too large")
		}

		c.fixMsg(&messages[idx], ml)
	}

	c.debug(func(d *debugger) {
		for _, m := range messages {
			d.debugf(1, "send msgs: %+v", m)
		}
	})

	if err := c.sock.SendMessages(messages); err != nil {
		c.debug(func(d *debugger) {
			d.debugf(1, "send msgs: err: %v", err)
		})

		return nil, err
	}

	return messages, nil
}

// Send sends a single Message to netlink.  In most cases, m.Header's Length,
// Sequence, and PID fields should be set to 0, so they can be populated
// automatically before the Message is sent.  On success, Send returns a copy
// of the Message with all parameters populated, for later validation.
//
// If m.Header.Length is 0, it will be automatically populated using the
// correct length for the Message, including its payload.
//
// If m.Header.Sequence is 0, it will be automatically populated using the
// next sequence number for this connection.
//
// If m.Header.PID is 0, it will be automatically populated using a PID
// assigned by netlink.
func (c *Conn) Send(m Message) (Message, error) {
	ml := nlmsgLength(len(m.Data))

	// TODO(mdlayher): fine-tune this limit.
	if ml > (1024 * 32) {
		return Message{}, errors.New("netlink message data too large")
	}

	c.fixMsg(&m, ml)

	c.debug(func(d *debugger) {
		d.debugf(1, "send: %+v", m)
	})

	if err := c.sock.Send(m); err != nil {
		c.debug(func(d *debugger) {
			d.debugf(1, "send: err: %v", err)
		})

		return Message{}, err
	}

	return m, nil
}

// Receive receives one or more messages from netlink.  Multi-part messages are
// handled transparently and returned as a single slice of Messages, with the
// final empty "multi-part done" message removed.
//
// If any of the messages indicate a netlink error, that error will be returned.
func (c *Conn) Receive() ([]Message, error) {
	msgs, err := c.receive()
	if err != nil {
		c.debug(func(d *debugger) {
			d.debugf(1, "recv: err: %v", err)
		})

		return nil, err
	}

	c.debug(func(d *debugger) {
		for _, m := range msgs {
			d.debugf(1, "recv: %+v", m)
		}
	})

	// When using nltest, it's possible for zero messages to be returned by receive.
	if len(msgs) == 0 {
		return msgs, nil
	}

	// Trim the final message with multi-part done indicator if
	// present.
	if m := msgs[len(msgs)-1]; m.Header.Flags&HeaderFlagsMulti != 0 && m.Header.Type == HeaderTypeDone {
		return msgs[:len(msgs)-1], nil
	}

	return msgs, nil
}

// receive is the internal implementation of Conn.Receive, which can be called
// recursively to handle multi-part messages.
func (c *Conn) receive() ([]Message, error) {
	var res []Message
	for {
		msgs, err := c.sock.Receive()
		if err != nil {
			return nil, err
		}

		// If this message is multi-part, we will need to perform an recursive call
		// to continue draining the socket
		var multi bool

		for _, m := range msgs {
			if err := checkMessage(m); err != nil {
				return nil, err
			}

			// Does this message indicate a multi-part message?
			if m.Header.Flags&HeaderFlagsMulti == 0 {
				// No, check the next messages.
				continue
			}

			// Does this message indicate the last message in a series of
			// multi-part messages from a single read?
			multi = m.Header.Type != HeaderTypeDone
		}

		res = append(res, msgs...)

		if !multi {
			// No more messages coming.
			return res, nil
		}
	}
}

// An fder is a Socket that supports retrieving its raw file descriptor.
type fder interface {
	Socket
	FD() int
}

var _ io.ReadWriteCloser = &fileReadWriteCloser{}

// A fileReadWriteCloser is a limited *os.File which only allows access to its
// Read and Write methods.
type fileReadWriteCloser struct {
	f *os.File
}

// Read implements io.ReadWriteCloser.
func (rwc *fileReadWriteCloser) Read(b []byte) (int, error) { return rwc.f.Read(b) }

// Write implements io.ReadWriteCloser.
func (rwc *fileReadWriteCloser) Write(b []byte) (int, error) { return rwc.f.Write(b) }

// Close implements io.ReadWriteCloser.
func (rwc *fileReadWriteCloser) Close() error { return rwc.f.Close() }

// ReadWriteCloser returns a raw io.ReadWriteCloser backed by the connection
// of the Conn.
//
// ReadWriteCloser is intended for advanced use cases, such as those that do
// not involve standard netlink message passing.
//
// Once invoked, it is the caller's responsibility to ensure that operations
// performed using Conn and the raw io.ReadWriteCloser do not conflict with
// each other.  In almost all scenarios, only one of the two should be used.
func (c *Conn) ReadWriteCloser() (io.ReadWriteCloser, error) {
	fc, ok := c.sock.(fder)
	if !ok {
		return nil, errReadWriteCloserNotSupported
	}

	return &fileReadWriteCloser{
		// Backing the io.ReadWriteCloser with an *os.File enables easy reading
		// and writing without more system call boilerplate.
		f: os.NewFile(uintptr(fc.FD()), "netlink"),
	}, nil
}

// A groupJoinLeaver is a Socket that supports joining and leaving
// netlink multicast groups.
type groupJoinLeaver interface {
	Socket
	JoinGroup(group uint32) error
	LeaveGroup(group uint32) error
}

// JoinGroup joins a netlink multicast group by its ID.
func (c *Conn) JoinGroup(group uint32) error {
	gc, ok := c.sock.(groupJoinLeaver)
	if !ok {
		return errMulticastGroupsNotSupported
	}

	return gc.JoinGroup(group)
}

// LeaveGroup leaves a netlink multicast group by its ID.
func (c *Conn) LeaveGroup(group uint32) error {
	gc, ok := c.sock.(groupJoinLeaver)
	if !ok {
		return errMulticastGroupsNotSupported
	}

	return gc.LeaveGroup(group)
}

// A bpfSetter is a Socket that supports setting and removing BPF filters.
type bpfSetter interface {
	Socket
	bpf.Setter
	RemoveBPF() error
}

// SetBPF attaches an assembled BPF program to a Conn.
func (c *Conn) SetBPF(filter []bpf.RawInstruction) error {
	bc, ok := c.sock.(bpfSetter)
	if !ok {
		return errBPFFiltersNotSupported
	}

	return bc.SetBPF(filter)
}

// RemoveBPF removes a BPF filter from a Conn.
func (c *Conn) RemoveBPF() error {
	s, ok := c.sock.(bpfSetter)
	if !ok {
		return errBPFFiltersNotSupported
	}

	return s.RemoveBPF()
}

// A ConnOption is a boolean option that may be set for a Conn.
type ConnOption int

// Possible ConnOption values.  These constants are equivalent to the Linux
// setsockopt boolean options for netlink sockets.
const (
	PacketInfo ConnOption = iota
	BroadcastError
	NoENOBUFS
	ListenAllNSID
	CapAcknowledge
)

// An optionSetter is a Socket that supports setting netlink options.
type optionSetter interface {
	Socket
	SetOption(option ConnOption, enable bool) error
}

// SetOption enables or disables a netlink socket option for the Conn.
func (c *Conn) SetOption(option ConnOption, enable bool) error {
	fc, ok := c.sock.(optionSetter)
	if !ok {
		return errOptionsNotSupported
	}

	return fc.SetOption(option, enable)
}

// nextSequence atomically increments Conn's sequence number and returns
// the incremented value.
func (c *Conn) nextSequence() uint32 {
	return atomic.AddUint32(c.seq, 1)
}

// Validate validates one or more reply Messages against a request Message,
// ensuring that they contain matching sequence numbers and PIDs.
func Validate(request Message, replies []Message) error {
	for _, m := range replies {
		// Check for mismatched sequence, unless:
		//   - request had no sequence, meaning we are probably validating
		//     a multicast reply
		if m.Header.Sequence != request.Header.Sequence && request.Header.Sequence != 0 {
			return errMismatchedSequence
		}

		// Check for mismatched PID, unless:
		//   - request had no PID, meaning we are either:
		//     - validating a multicast reply
		//     - netlink has not yet assigned us a PID
		//   - response had no PID, meaning it's from the kernel as a multicast reply
		if m.Header.PID != request.Header.PID && request.Header.PID != 0 && m.Header.PID != 0 {
			return errMismatchedPID
		}
	}

	return nil
}

// Config contains options for a Conn.
type Config struct {
	// Groups is a bitmask which specifies multicast groups. If set to 0,
	// no multicast group subscriptions will be made.
	Groups uint32

	// Experimental: do not lock the internal system call handling goroutine
	// to its OS thread.  This may result in a speed-up of system call handling,
	// but may cause unexpected behavior when sending and receiving a large number
	// of messages.
	//
	// This should almost certainly be set to false, but if you come up with a
	// valid reason for using this, please file an issue at
	// https://github.com/mdlayher/netlink to discuss your thoughts.
	NoLockThread bool
}
<|endoftext|>


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

package nftables

import (
	"fmt"
	"sync"

	"github.com/google/nftables/expr"
	"github.com/mdlayher/netlink"
	"github.com/mdlayher/netlink/nltest"
	"golang.org/x/sys/unix"
)

// A Conn represents a netlink connection of the nftables family.
//
// All methods return their input, so that variables can be defined from string
// literals when desired.
//
// Commands are buffered. Flush sends all buffered commands in a single batch.
type Conn struct {
	TestDial nltest.Func // for testing only; passed to nltest.Dial
	NetNS    int         // Network namespace netlink will interact with.
	sync.Mutex
	messages []netlink.Message
	err      error
}

// Flush sends all buffered commands in a single batch to nftables.
func (cc *Conn) Flush() error {
	cc.Lock()
	defer func() {
		cc.messages = nil
		cc.Unlock()
	}()
	if len(cc.messages) == 0 {
		// Messages were already programmed, returning nil
		return nil
	}
	if cc.err != nil {
		return cc.err // serialization error
	}
	conn, err := cc.dialNetlink()
	if err != nil {
		return err
	}

	defer conn.Close()

	if _, err := conn.SendMessages(batch(cc.messages)); err != nil {
		return fmt.Errorf("SendMessages: %w", err)
	}

	if _, err := conn.Receive(); err != nil {
		return fmt.Errorf("Receive:  %w", err)
	}

	return nil
}

// FlushRuleset flushes the entire ruleset. See also
// https://wiki.nftables.org/wiki-nftables/index.php/Operations_at_ruleset_level
func (cc *Conn) FlushRuleset() {
	cc.Lock()
	defer cc.Unlock()
	cc.messages = append(cc.messages, netlink.Message{
		Header: netlink.Header{
			Type:  netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_DELTABLE),
			Flags: netlink.Request | netlink.Acknowledge | netlink.Create,
		},
		Data: extraHeader(0, 0),
	})
}

func (cc *Conn) dialNetlink() (*netlink.Conn, error) {
	if cc.TestDial != nil {
		return nltest.Dial(cc.TestDial), nil
	}
	return netlink.Dial(unix.NETLINK_NETFILTER, &netlink.Config{NetNS: cc.NetNS})
}

func (cc *Conn) setErr(err error) {
	if cc.err != nil {
		return
	}
	cc.err = err
}

func (cc *Conn) marshalAttr(attrs []netlink.Attribute) []byte {
	b, err := netlink.MarshalAttributes(attrs)
	if err != nil {
		cc.setErr(err)
		return nil
	}
	return b
}

func (cc *Conn) marshalExpr(e expr.Any) []byte {
	b, err := expr.Marshal(e)
	if err != nil {
		cc.setErr(err)
		return nil
	}
	return b
}

func batch(messages []netlink.Message) []netlink.Message {
	batch := []netlink.Message{
		{
			Header: netlink.Header{
				Type:  netlink.HeaderType(unix.NFNL_MSG_BATCH_BEGIN),
				Flags: netlink.Request,
			},
			Data: extraHeader(0, unix.NFNL_SUBSYS_NFTABLES),
		},
	}

	batch = append(batch, messages...)

	batch = append(batch, netlink.Message{
		Header: netlink.Header{
			Type:  netlink.HeaderType(unix.NFNL_MSG_BATCH_END),
			Flags: netlink.Request,
		},
		Data: extraHeader(0, unix.NFNL_SUBSYS_NFTABLES),
	})

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

package nftables

import (
	"fmt"
	"sync"

	"github.com/google/nftables/expr"
	"github.com/mdlayher/netlink"
	"github.com/mdlayher/netlink/nltest"
	"golang.org/x/sys/unix"
)

// A Conn represents a netlink connection of the nftables family.
//
// All methods return their input, so that variables can be defined from string
// literals when desired.
//
// Commands are buffered. Flush sends all buffered commands in a single batch.
type Conn struct {
	TestDial nltest.Func // for testing only; passed to nltest.Dial
	NetNS    int         // Network namespace netlink will interact with.
	sync.Mutex
	messages []netlink.Message
	err      error
}

// Flush sends all buffered commands in a single batch to nftables.
func (cc *Conn) Flush() error {
	cc.Lock()
	defer func() {
		cc.messages = nil
		cc.Unlock()
	}()
	if len(cc.messages) == 0 {
		// Messages were already programmed, returning nil
		return nil
	}
	if cc.err != nil {
		return cc.err // serialization error
	}
	conn, err := cc.dialNetlink()
	if err != nil {
		return err
	}

	defer conn.Close()

	if _, err := conn.SendMessages(batch(cc.messages)); err != nil {
		return fmt.Errorf("SendMessages: %w", err)
	}

	if _, err := conn.Receive(); err != nil {
		return fmt.Errorf("Receive: %w", err)
	}

	return nil
}

// FlushRuleset flushes the entire ruleset. See also
// https://wiki.nftables.org/wiki-nftables/index.php/Operations_at_ruleset_level
func (cc *Conn) FlushRuleset() {
	cc.Lock()
	defer cc.Unlock()
	cc.messages = append(cc.messages, netlink.Message{
		Header: netlink.Header{
			Type:  netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_DELTABLE),
			Flags: netlink.Request | netlink.Acknowledge | netlink.Create,
		},
		Data: extraHeader(0, 0),
	})
}

func (cc *Conn) dialNetlink() (*netlink.Conn, error) {
	if cc.TestDial != nil {
		return nltest.Dial(cc.TestDial), nil
	}
	return netlink.Dial(unix.NETLINK_NETFILTER, &netlink.Config{NetNS: cc.NetNS})
}

func (cc *Conn) setErr(err error) {
	if cc.err != nil {
		return
	}
	cc.err = err
}

func (cc *Conn) marshalAttr(attrs []netlink.Attribute) []byte {
	b, err := netlink.MarshalAttributes(attrs)
	if err != nil {
		cc.setErr(err)
		return nil
	}
	return b
}

func (cc *Conn) marshalExpr(e expr.Any) []byte {
	b, err := expr.Marshal(e)
	if err != nil {
		cc.setErr(err)
		return nil
	}
	return b
}

func batch(messages []netlink.Message) []netlink.Message {
	batch := []netlink.Message{
		{
			Header: netlink.Header{
				Type:  netlink.HeaderType(unix.NFNL_MSG_BATCH_BEGIN),
				Flags: netlink.Request,
			},
			Data: extraHeader(0, unix.NFNL_SUBSYS_NFTABLES),
		},
	}

	batch = append(batch, messages...)

	batch = append(batch, netlink.Message{
		Header: netlink.Header{
			Type:  netlink.HeaderType(unix.NFNL_MSG_BATCH_END),
			Flags: netlink.Request,
		},
		Data: extraHeader(0, unix.NFNL_SUBSYS_NFTABLES),
	})

	return batch
}
<|endoftext|>


<commit_before>// +build e2e

/*
Copyright 2019 The Knative 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 ingress

import (
	"fmt"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/gorilla/websocket"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/sets"
	"knative.dev/serving/pkg/apis/networking/v1alpha1"
	"knative.dev/serving/test"
)

// TestWebsocket verifies that websockets may be used via a simple Ingress.
func TestWebsocket(t *testing.T) {
	t.Parallel()
	clients := test.Setup(t)

	const suffix = "- pong"
	name, port, cancel := CreateWebsocketService(t, clients, suffix)
	defer cancel()

	domain := name + ".example.com"

	// Create a simple Ingress over the Service.
	_, dialCtx, cancel := CreateIngressReadyDialContext(t, clients, v1alpha1.IngressSpec{
		Rules: []v1alpha1.IngressRule{{
			Hosts:      []string{domain},
			Visibility: v1alpha1.IngressVisibilityExternalIP,
			HTTP: &v1alpha1.HTTPIngressRuleValue{
				Paths: []v1alpha1.HTTPIngressPath{{
					Splits: []v1alpha1.IngressBackendSplit{{
						IngressBackend: v1alpha1.IngressBackend{
							ServiceName:      name,
							ServiceNamespace: test.ServingNamespace,
							ServicePort:      intstr.FromInt(port),
						},
					}},
				}},
			},
		}},
	})
	defer cancel()

	dialer := websocket.Dialer{
		NetDialContext:   dialCtx,
		Proxy:            http.ProxyFromEnvironment,
		HandshakeTimeout: 45 * time.Second,
	}

	u := url.URL{Scheme: "ws", Host: domain, Path: "/"}
	conn, _, err := dialer.Dial(u.String(), http.Header{"Host": {domain}})
	if err != nil {
		t.Fatalf("Dial() = %v", err)
	}
	defer conn.Close()

	for i := 0; i < 100; i++ {
		checkWebsocketRoundTrip(t, conn, suffix)
	}
}

// TestWebsocketSplit verifies that websockets may be used across a traffic split.
func TestWebsocketSplit(t *testing.T) {
	t.Parallel()
	clients := test.Setup(t)

	const suffixBlue = "- blue"
	blueName, bluePort, cancel := CreateWebsocketService(t, clients, suffixBlue)
	defer cancel()

	const suffixGreen = "- green"
	greenName, greenPort, cancel := CreateWebsocketService(t, clients, suffixGreen)
	defer cancel()

	// The suffixes we expect to see.
	want := sets.NewString(suffixBlue, suffixGreen)

	// Create a simple Ingress over the Service.
	name := test.ObjectNameForTest(t)
	domain := name + ".example.com"
	_, dialCtx, cancel := CreateIngressReadyDialContext(t, clients, v1alpha1.IngressSpec{
		Rules: []v1alpha1.IngressRule{{
			Hosts:      []string{domain},
			Visibility: v1alpha1.IngressVisibilityExternalIP,
			HTTP: &v1alpha1.HTTPIngressRuleValue{
				Paths: []v1alpha1.HTTPIngressPath{{
					Splits: []v1alpha1.IngressBackendSplit{{
						IngressBackend: v1alpha1.IngressBackend{
							ServiceName:      blueName,
							ServiceNamespace: test.ServingNamespace,
							ServicePort:      intstr.FromInt(bluePort),
						},
						Percent: 50,
					}, {
						IngressBackend: v1alpha1.IngressBackend{
							ServiceName:      greenName,
							ServiceNamespace: test.ServingNamespace,
							ServicePort:      intstr.FromInt(greenPort),
						},
						Percent: 50,
					}},
				}},
			},
		}},
	})
	defer cancel()

	dialer := websocket.Dialer{
		NetDialContext:   dialCtx,
		Proxy:            http.ProxyFromEnvironment,
		HandshakeTimeout: 45 * time.Second,
	}
	u := url.URL{Scheme: "ws", Host: domain, Path: "/"}

	got := sets.NewString()
	for i := 0; i < 10; i++ {
		conn, _, err := dialer.Dial(u.String(), http.Header{"Host": {domain}})
		if err != nil {
			t.Fatalf("Dial() = %v", err)
		}
		defer conn.Close()

		suffix := findWebsocketSuffix(t, conn)
		if suffix == "" {
			continue
		}
		got.Insert(suffix)

		for j := 0; j < 10; j++ {
			checkWebsocketRoundTrip(t, conn, suffix)
		}
	}

	if !want.Equal(got) {
		t.Errorf("(-want, +got) = %s", cmp.Diff(want, got))
	}
}

func findWebsocketSuffix(t *testing.T, conn *websocket.Conn) string {
	// Establish the suffix that corresponds to this socket.
	message := fmt.Sprintf("ping - %d", rand.Intn(1000))
	if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
		t.Errorf("WriteMessage() = %v", err)
		return ""
	}

	_, recv, err := conn.ReadMessage()
	if err != nil {
		t.Errorf("ReadMessage() = %v", err)
		return ""
	}
	gotMsg := string(recv)
	if !strings.HasPrefix(gotMsg, message) {
		t.Errorf("ReadMessage() = %s, wanted %s prefix", gotMsg, message)
		return ""
	}
	return strings.TrimSpace(strings.TrimPrefix(gotMsg, message))
}

func checkWebsocketRoundTrip(t *testing.T, conn *websocket.Conn, suffix string) {
	message := fmt.Sprintf("ping - %d", rand.Intn(1000))
	if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
		t.Errorf("WriteMessage() = %v", err)
		return
	}

	// Read back the echoed message and compared with sent.
	if _, recv, err := conn.ReadMessage(); err != nil {
		t.Errorf("ReadMessage() = %v", err)
	} else if got, want := string(recv), message+" "+suffix; got != want {
		t.Errorf("ReadMessage() = %s, wanted %s", got, want)
	}
}
<commit_msg>Harden websocket split test. (#7707)<commit_after>// +build e2e

/*
Copyright 2019 The Knative 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 ingress

import (
	"fmt"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/gorilla/websocket"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/sets"
	"knative.dev/serving/pkg/apis/networking/v1alpha1"
	"knative.dev/serving/test"
)

// TestWebsocket verifies that websockets may be used via a simple Ingress.
func TestWebsocket(t *testing.T) {
	t.Parallel()
	clients := test.Setup(t)

	const suffix = "- pong"
	name, port, cancel := CreateWebsocketService(t, clients, suffix)
	defer cancel()

	domain := name + ".example.com"

	// Create a simple Ingress over the Service.
	_, dialCtx, cancel := CreateIngressReadyDialContext(t, clients, v1alpha1.IngressSpec{
		Rules: []v1alpha1.IngressRule{{
			Hosts:      []string{domain},
			Visibility: v1alpha1.IngressVisibilityExternalIP,
			HTTP: &v1alpha1.HTTPIngressRuleValue{
				Paths: []v1alpha1.HTTPIngressPath{{
					Splits: []v1alpha1.IngressBackendSplit{{
						IngressBackend: v1alpha1.IngressBackend{
							ServiceName:      name,
							ServiceNamespace: test.ServingNamespace,
							ServicePort:      intstr.FromInt(port),
						},
					}},
				}},
			},
		}},
	})
	defer cancel()

	dialer := websocket.Dialer{
		NetDialContext:   dialCtx,
		Proxy:            http.ProxyFromEnvironment,
		HandshakeTimeout: 45 * time.Second,
	}

	u := url.URL{Scheme: "ws", Host: domain, Path: "/"}
	conn, _, err := dialer.Dial(u.String(), http.Header{"Host": {domain}})
	if err != nil {
		t.Fatalf("Dial() = %v", err)
	}
	defer conn.Close()

	for i := 0; i < 100; i++ {
		checkWebsocketRoundTrip(t, conn, suffix)
	}
}

// TestWebsocketSplit verifies that websockets may be used across a traffic split.
func TestWebsocketSplit(t *testing.T) {
	t.Parallel()
	clients := test.Setup(t)

	const suffixBlue = "- blue"
	blueName, bluePort, cancel := CreateWebsocketService(t, clients, suffixBlue)
	defer cancel()

	const suffixGreen = "- green"
	greenName, greenPort, cancel := CreateWebsocketService(t, clients, suffixGreen)
	defer cancel()

	// The suffixes we expect to see.
	want := sets.NewString(suffixBlue, suffixGreen)

	// Create a simple Ingress over the Service.
	name := test.ObjectNameForTest(t)
	domain := name + ".example.com"
	_, dialCtx, cancel := CreateIngressReadyDialContext(t, clients, v1alpha1.IngressSpec{
		Rules: []v1alpha1.IngressRule{{
			Hosts:      []string{domain},
			Visibility: v1alpha1.IngressVisibilityExternalIP,
			HTTP: &v1alpha1.HTTPIngressRuleValue{
				Paths: []v1alpha1.HTTPIngressPath{{
					Splits: []v1alpha1.IngressBackendSplit{{
						IngressBackend: v1alpha1.IngressBackend{
							ServiceName:      blueName,
							ServiceNamespace: test.ServingNamespace,
							ServicePort:      intstr.FromInt(bluePort),
						},
						Percent: 50,
					}, {
						IngressBackend: v1alpha1.IngressBackend{
							ServiceName:      greenName,
							ServiceNamespace: test.ServingNamespace,
							ServicePort:      intstr.FromInt(greenPort),
						},
						Percent: 50,
					}},
				}},
			},
		}},
	})
	defer cancel()

	dialer := websocket.Dialer{
		NetDialContext:   dialCtx,
		Proxy:            http.ProxyFromEnvironment,
		HandshakeTimeout: 45 * time.Second,
	}
	u := url.URL{Scheme: "ws", Host: domain, Path: "/"}

	const maxRequests = 100
	got := sets.NewString()
	for i := 0; i < maxRequests; i++ {
		conn, _, err := dialer.Dial(u.String(), http.Header{"Host": {domain}})
		if err != nil {
			t.Fatalf("Dial() = %v", err)
		}
		defer conn.Close()

		suffix := findWebsocketSuffix(t, conn)
		if suffix == "" {
			continue
		}
		got.Insert(suffix)

		for j := 0; j < 10; j++ {
			checkWebsocketRoundTrip(t, conn, suffix)
		}

		if want.Equal(got) {
			// Short circuit if we've seen all splits.
			return
		}
	}

	// Us getting here means we haven't seen splits.
	t.Errorf("(over %d requests) (-want, +got) = %s", maxRequests, cmp.Diff(want, got))
}

func findWebsocketSuffix(t *testing.T, conn *websocket.Conn) string {
	// Establish the suffix that corresponds to this socket.
	message := fmt.Sprintf("ping - %d", rand.Intn(1000))
	if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
		t.Errorf("WriteMessage() = %v", err)
		return ""
	}

	_, recv, err := conn.ReadMessage()
	if err != nil {
		t.Errorf("ReadMessage() = %v", err)
		return ""
	}
	gotMsg := string(recv)
	if !strings.HasPrefix(gotMsg, message) {
		t.Errorf("ReadMessage() = %s, wanted %s prefix", gotMsg, message)
		return ""
	}
	return strings.TrimSpace(strings.TrimPrefix(gotMsg, message))
}

func checkWebsocketRoundTrip(t *testing.T, conn *websocket.Conn, suffix string) {
	message := fmt.Sprintf("ping - %d", rand.Intn(1000))
	if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
		t.Errorf("WriteMessage() = %v", err)
		return
	}

	// Read back the echoed message and compared with sent.
	if _, recv, err := conn.ReadMessage(); err != nil {
		t.Errorf("ReadMessage() = %v", err)
	} else if got, want := string(recv), message+" "+suffix; got != want {
		t.Errorf("ReadMessage() = %s, wanted %s", got, want)
	}
}
<|endoftext|>


<commit_before>package copy

import (
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

const (
	// tmpPermissionForDirectory makes the destination directory writable,
	// so that stuff can be copied recursively even if any original directory is NOT writable.
	// See https://github.com/otiai10/copy/pull/9 for more information.
	tmpPermissionForDirectory = os.FileMode(0755)
)

// Copy copies src to dest, doesn't matter if src is a directory or a file
func Copy(src, dest string, opt ...Options) error {
	info, err := os.Lstat(src)
	if err != nil {
		return err
	}
	return copy(src, dest, info, assure(opt...))
}

// copy dispatches copy-funcs according to the mode.
// Because this "copy" could be called recursively,
// "info" MUST be given here, NOT nil.
func copy(src, dest string, info os.FileInfo, opt Options) error {

	if opt.Skip(src) {
		return nil
	}

	if info.Mode()&os.ModeSymlink != 0 {
		return onsymlink(src, dest, info, opt)
	}

	if info.IsDir() {
		return dcopy(src, dest, info, opt)
	}
	return fcopy(src, dest, info)
}

// fcopy is for just a file,
// with considering existence of parent directory
// and file permission.
func fcopy(src, dest string, info os.FileInfo) (err error) {

	if err := os.MkdirAll(filepath.Dir(dest), os.ModePerm); err != nil {
		return err
	}

	f, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer fclose(f, &err)

	if err = os.Chmod(f.Name(), info.Mode()); err != nil {
		return err
	}

	s, err := os.Open(src)
	if err != nil {
		return err
	}
	defer fclose(s, &err)

	_, err = io.Copy(f, s)
	return err
}

// dcopy is for a directory,
// with scanning contents inside the directory
// and pass everything to "copy" recursively.
func dcopy(srcdir, destdir string, info os.FileInfo, opt Options) (err error) {

	originalMode := info.Mode()

	// Make dest dir with 0755 so that everything writable.
	if err := os.MkdirAll(destdir, tmpPermissionForDirectory); err != nil {
		return err
	}
	// Recover dir mode with original one.
	defer chmod(destdir, originalMode, &err)

	contents, err := ioutil.ReadDir(srcdir)
	if err != nil {
		return err
	}

	for _, content := range contents {
		cs, cd := filepath.Join(srcdir, content.Name()), filepath.Join(destdir, content.Name())
		if err := copy(cs, cd, content, opt); err != nil {
			// If any error, exit immediately
			return err
		}
	}

	return nil
}

func onsymlink(src, dest string, info os.FileInfo, opt Options) error {

	switch opt.OnSymlink(src) {
	case Shallow:
		return lcopy(src, dest)
	case Deep:
		orig, err := os.Readlink(src)
		if err != nil {
			return err
		}
		info, err = os.Lstat(orig)
		if err != nil {
			return err
		}
		return copy(orig, dest, info, opt)
	case Skip:
		fallthrough
	default:
		return nil // do nothing
	}
}

// lcopy is for a symlink,
// with just creating a new symlink by replicating src symlink.
func lcopy(src, dest string) error {
	src, err := os.Readlink(src)
	if err != nil {
		return err
	}
	return os.Symlink(src, dest)
}

// fclose ANYHOW closes file,
// with asiging error occured BUT respecting the error already reported.
func fclose(f *os.File, reported *error) {
	if err := f.Close(); *reported == nil {
		*reported = err
	}
}

// chmod ANYHOW changes file mode,
// with asiging error occured BUT respecting the error already reported.
func chmod(dir string, mode os.FileMode, reported *error) {
	if err := os.Chmod(dir, mode); *reported == nil {
		*reported = err
	}
}

// assure Options struct, should be called only once.
// All optional values MUST NOT BE nil/zero after assured.
func assure(opts ...Options) Options {
	if len(opts) == 0 {
		return DefaultOptions
	}
	if opts[0].OnSymlink == nil {
		opts[0].OnSymlink = DefaultOptions.OnSymlink
	}
	if opts[0].Skip == nil {
		opts[0].Skip = DefaultOptions.Skip
	}
	return opts[0]
}
<commit_msg>Fix wording<commit_after>package copy

import (
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

const (
	// tmpPermissionForDirectory makes the destination directory writable,
	// so that stuff can be copied recursively even if any original directory is NOT writable.
	// See https://github.com/otiai10/copy/pull/9 for more information.
	tmpPermissionForDirectory = os.FileMode(0755)
)

// Copy copies src to dest, doesn't matter if src is a directory or a file.
func Copy(src, dest string, opt ...Options) error {
	info, err := os.Lstat(src)
	if err != nil {
		return err
	}
	return copy(src, dest, info, assure(opt...))
}

// copy dispatches copy-funcs according to the mode.
// Because this "copy" could be called recursively,
// "info" MUST be given here, NOT nil.
func copy(src, dest string, info os.FileInfo, opt Options) error {

	if opt.Skip(src) {
		return nil
	}

	if info.Mode()&os.ModeSymlink != 0 {
		return onsymlink(src, dest, info, opt)
	}

	if info.IsDir() {
		return dcopy(src, dest, info, opt)
	}
	return fcopy(src, dest, info)
}

// fcopy is for just a file,
// with considering existence of parent directory
// and file permission.
func fcopy(src, dest string, info os.FileInfo) (err error) {

	if err := os.MkdirAll(filepath.Dir(dest), os.ModePerm); err != nil {
		return err
	}

	f, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer fclose(f, &err)

	if err = os.Chmod(f.Name(), info.Mode()); err != nil {
		return err
	}

	s, err := os.Open(src)
	if err != nil {
		return err
	}
	defer fclose(s, &err)

	_, err = io.Copy(f, s)
	return err
}

// dcopy is for a directory,
// with scanning contents inside the directory
// and pass everything to "copy" recursively.
func dcopy(srcdir, destdir string, info os.FileInfo, opt Options) (err error) {

	originalMode := info.Mode()

	// Make dest dir with 0755 so that everything writable.
	if err := os.MkdirAll(destdir, tmpPermissionForDirectory); err != nil {
		return err
	}
	// Recover dir mode with original one.
	defer chmod(destdir, originalMode, &err)

	contents, err := ioutil.ReadDir(srcdir)
	if err != nil {
		return err
	}

	for _, content := range contents {
		cs, cd := filepath.Join(srcdir, content.Name()), filepath.Join(destdir, content.Name())
		if err := copy(cs, cd, content, opt); err != nil {
			// If any error, exit immediately
			return err
		}
	}

	return nil
}

func onsymlink(src, dest string, info os.FileInfo, opt Options) error {

	switch opt.OnSymlink(src) {
	case Shallow:
		return lcopy(src, dest)
	case Deep:
		orig, err := os.Readlink(src)
		if err != nil {
			return err
		}
		info, err = os.Lstat(orig)
		if err != nil {
			return err
		}
		return copy(orig, dest, info, opt)
	case Skip:
		fallthrough
	default:
		return nil // do nothing
	}
}

// lcopy is for a symlink,
// with just creating a new symlink by replicating src symlink.
func lcopy(src, dest string) error {
	src, err := os.Readlink(src)
	if err != nil {
		return err
	}
	return os.Symlink(src, dest)
}

// fclose ANYHOW closes file,
// with asiging error raised during Close,
// BUT respecting the error already reported.
func fclose(f *os.File, reported *error) {
	if err := f.Close(); *reported == nil {
		*reported = err
	}
}

// chmod ANYHOW changes file mode,
// with asiging error raised during Chmod,
// BUT respecting the error already reported.
func chmod(dir string, mode os.FileMode, reported *error) {
	if err := os.Chmod(dir, mode); *reported == nil {
		*reported = err
	}
}

// assure Options struct, should be called only once.
// All optional values MUST NOT BE nil/zero after assured.
func assure(opts ...Options) Options {
	if len(opts) == 0 {
		return DefaultOptions
	}
	if opts[0].OnSymlink == nil {
		opts[0].OnSymlink = DefaultOptions.OnSymlink
	}
	if opts[0].Skip == nil {
		opts[0].Skip = DefaultOptions.Skip
	}
	return opts[0]
}
<|endoftext|>


<commit_before>package blob

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/orisano/httpc"
	"github.com/pkg/errors"
)

const (
	APIVersion = "2017-04-17"
)

var DefaultUserAgent = fmt.Sprintf("Go/%s (%s-%s) go-ams.blob", runtime.Version(), runtime.GOARCH, runtime.GOOS)

type SASClient struct {
	rb         *httpc.RequestBuilder
	httpClient *http.Client

	userAgent string
	logger    *log.Logger
	debug     bool

	TimeNow func() time.Time
}

func NewSASClient(rawurl string, opts ...clientOption) (*SASClient, error) {
	options := &clientOptions{
		Client:    http.DefaultClient,
		UserAgent: DefaultUserAgent,
		Logger:    log.New(ioutil.Discard, "", log.Lshortfile),
		Debug:     false,
	}
	for _, opt := range opts {
		opt(options)
	}

	h := make(http.Header)
	h.Set("x-ms-version", APIVersion)
	h.Set("User-Agent", options.UserAgent)

	rb, err := httpc.NewRequestBuilder(rawurl, h)
	if err != nil {
		return nil, err
	}

	if options.Debug {
		httpc.InjectDebugTransport(options.Client, os.Stderr)
	}

	return &SASClient{
		rb:         rb,
		httpClient: options.Client,
		userAgent:  options.UserAgent,
		logger:     options.Logger,
		debug:      options.Debug,
		TimeNow:    time.Now,
	}, nil
}

func withDate(t time.Time) httpc.RequestOption {
	return func(o *httpc.RequestOptions) error {
		o.Header.Set("Date", t.UTC().Format(time.RFC3339))
		return nil
	}
}

func (c *SASClient) PutBlob(ctx context.Context, blob io.Reader, blockID string) error {
	if ctx == nil {
		return errors.New("missing ctx")
	}
	if blob == nil {
		return errors.New("missing blob")
	}
	if len(blockID) == 0 {
		return errors.New("missing blockID")
	}
	req, err := c.rb.NewRequest(ctx, http.MethodPut, "",
		withDate(c.TimeNow()),
		httpc.SetHeaderField("x-ms-blob-type", "BlockBlob"),
		httpc.AddQuery("comp", "block"),
		httpc.AddQuery("blockid", blockID),
		httpc.WithBinary(blob),
		httpc.EnforceContentLength,
	)
	if err != nil {
		return errors.Wrap(err, "failed to construct http request")
	}
	c.logger.Print("[INFO] put blob ...")
	resp, err := httpc.Retry(c.httpClient, req)
	if err != nil {
		return errors.Wrap(err, "failed to http request")
	}
	defer resp.Body.Close()

	if got := resp.StatusCode; got != http.StatusCreated {
		return errors.Errorf("unexpected status code. expected: %v, but got: %v", http.StatusCreated, got)
	}

	c.logger.Print("[INFO] completed")
	return nil
}

func (c *SASClient) PutBlockList(ctx context.Context, blockList []string) error {
	if ctx == nil {
		return errors.New("missing ctx")
	}
	if len(blockList) == 0 {
		return errors.New("missing blockList")
	}
	req, err := c.rb.NewRequest(ctx, http.MethodPut, "",
		withDate(c.TimeNow()),
		httpc.AddQuery("comp", "blocklist"),
		httpc.WithXML(&BlockList{Blocks: blockList}),
		httpc.EnforceContentLength,
	)
	if err != nil {
		return errors.Wrap(err, "failed to construct http request")
	}
	c.logger.Print("[INFO] put block list ...")
	resp, err := httpc.Retry(c.httpClient, req)
	if err != nil {
		return errors.Wrap(err, "failed to http request")
	}
	defer resp.Body.Close()

	if got := resp.StatusCode; got != http.StatusCreated {
		return errors.Errorf("unexpected status code. expected: %v, but got: %v", http.StatusCreated, got)
	}

	c.logger.Print("[INFO] completed")
	return nil
}
<commit_msg>feat: add Upload method<commit_after>package blob

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/orisano/httpc"
	"github.com/pkg/errors"
)

const (
	APIVersion = "2017-04-17"
)

var DefaultUserAgent = fmt.Sprintf("Go/%s (%s-%s) go-ams.blob", runtime.Version(), runtime.GOARCH, runtime.GOOS)

type SASClient struct {
	rb         *httpc.RequestBuilder
	httpClient *http.Client

	userAgent string
	logger    *log.Logger
	debug     bool

	TimeNow func() time.Time
}

func NewSASClient(rawurl string, opts ...clientOption) (*SASClient, error) {
	options := &clientOptions{
		Client:    http.DefaultClient,
		UserAgent: DefaultUserAgent,
		Logger:    log.New(ioutil.Discard, "", log.Lshortfile),
		Debug:     false,
	}
	for _, opt := range opts {
		opt(options)
	}

	h := make(http.Header)
	h.Set("x-ms-version", APIVersion)
	h.Set("User-Agent", options.UserAgent)

	rb, err := httpc.NewRequestBuilder(rawurl, h)
	if err != nil {
		return nil, err
	}

	if options.Debug {
		httpc.InjectDebugTransport(options.Client, os.Stderr)
	}

	return &SASClient{
		rb:         rb,
		httpClient: options.Client,
		userAgent:  options.UserAgent,
		logger:     options.Logger,
		debug:      options.Debug,
		TimeNow:    time.Now,
	}, nil
}

func withDate(t time.Time) httpc.RequestOption {
	return func(o *httpc.RequestOptions) error {
		o.Header.Set("Date", t.UTC().Format(time.RFC3339))
		return nil
	}
}

func (c *SASClient) PutBlob(ctx context.Context, blob io.Reader, blockID string) error {
	if ctx == nil {
		return errors.New("missing ctx")
	}
	if blob == nil {
		return errors.New("missing blob")
	}
	if len(blockID) == 0 {
		return errors.New("missing blockID")
	}
	req, err := c.rb.NewRequest(ctx, http.MethodPut, "",
		withDate(c.TimeNow()),
		httpc.SetHeaderField("x-ms-blob-type", "BlockBlob"),
		httpc.AddQuery("comp", "block"),
		httpc.AddQuery("blockid", blockID),
		httpc.WithBinary(blob),
		httpc.EnforceContentLength,
	)
	if err != nil {
		return errors.Wrap(err, "failed to construct http request")
	}
	c.logger.Print("[INFO] put blob ...")
	resp, err := httpc.Retry(c.httpClient, req)
	if err != nil {
		return errors.Wrap(err, "failed to http request")
	}
	defer resp.Body.Close()

	if got := resp.StatusCode; got != http.StatusCreated {
		return errors.Errorf("unexpected status code. expected: %v, but got: %v", http.StatusCreated, got)
	}

	c.logger.Print("[INFO] completed")
	return nil
}

func (c *SASClient) PutBlockList(ctx context.Context, blockList []string) error {
	if ctx == nil {
		return errors.New("missing ctx")
	}
	if len(blockList) == 0 {
		return errors.New("missing blockList")
	}
	req, err := c.rb.NewRequest(ctx, http.MethodPut, "",
		withDate(c.TimeNow()),
		httpc.AddQuery("comp", "blocklist"),
		httpc.WithXML(&BlockList{Blocks: blockList}),
		httpc.EnforceContentLength,
	)
	if err != nil {
		return errors.Wrap(err, "failed to construct http request")
	}
	c.logger.Print("[INFO] put block list ...")
	resp, err := httpc.Retry(c.httpClient, req)
	if err != nil {
		return errors.Wrap(err, "failed to http request")
	}
	defer resp.Body.Close()

	if got := resp.StatusCode; got != http.StatusCreated {
		return errors.Errorf("unexpected status code. expected: %v, but got: %v", http.StatusCreated, got)
	}

	c.logger.Print("[INFO] completed")
	return nil
}

type UploadError struct {
	errs []error
	m    *sync.Mutex
}

func newUploadError() *UploadError {
	return &UploadError{
		m: new(sync.Mutex),
	}
}

func (e *UploadError) add(err error) {
	if err != nil {
		e.m.Lock()
		e.errs = append(e.errs, err)
		e.m.Unlock()
	}
}

func (e *UploadError) Error() string {
	return fmt.Sprintf("failed to upload: (%d error occurred)", len(e.errs))
}

func (e *UploadError) Errors() []error {
	return e.errs
}

func (c *SASClient) Upload(ctx context.Context, r io.Reader, chunkSize int64, workers uint) (int64, error) {
	if ctx == nil {
		return 0, errors.New("missing ctx")
	}
	if r == nil {
		return 0, errors.New("missing r")
	}
	if chunkSize <= 0 {
		return 0, errors.New("chunkSize must be greater than 0")
	}
	if workers == 0 {
		return 0, errors.New("workers must be greater then 0")
	}

	type job struct {
		blockID string
		r       io.Reader
	}

	jobs := make(chan job, workers)
	uploadErr := newUploadError()
	wg := new(sync.WaitGroup)
	wg.Add(int(workers))
	for w := uint(0); w < workers; w++ {
		go func() {
			for job := range jobs {
				err := c.PutBlob(ctx, job.r, job.blockID)
				if err != nil {
					uploadErr.add(err)
				}
			}
			wg.Done()
		}()
	}

	var blockList []string
	contentLength := int64(0)
	bi := 0
	for {
		bi++
		b := bytes.NewBuffer(make([]byte, 0, chunkSize))
		n, err := io.CopyN(b, r, chunkSize)
		if n == 0 {
			break
		}
		if err != nil {
			close(jobs)
			wg.Done()
			return 0, errors.Wrap(err, "failed to read r")
		}

		contentLength += n
		blockID := buildBlockID(bi)
		blockList = append(blockList, blockID)
		jobs <- job{blockID, b}

		if n < chunkSize {
			break
		}
	}
	close(jobs)
	wg.Done()

	if uploadErr.Errors() != nil {
		return 0, uploadErr
	}

	if err := c.PutBlockList(ctx, blockList); err != nil {
		return 0, errors.Wrap(err, "failed to put blocklist")
	}

	return contentLength, nil
}
<|endoftext|>


<commit_before>package signed

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/go/canonical/json"
	"github.com/docker/notary/tuf/data"
	"github.com/docker/notary/tuf/utils"
)

// Various basic signing errors
var (
	ErrNoSignatures = errors.New("tuf: data has no signatures")
	ErrInvalid      = errors.New("tuf: signature verification failed")
	ErrWrongType    = errors.New("tuf: meta file has wrong type")
)

// IsExpired checks if the given time passed before the present time
func IsExpired(t time.Time) bool {
	return t.Before(time.Now())
}

// VerifyExpiry returns ErrExpired if the metadata is expired
func VerifyExpiry(s *data.SignedCommon, role data.RoleName) error {
	if IsExpired(s.Expires) {
		logrus.Errorf("Metadata for %s expired", role)
		return ErrExpired{Role: role, Expired: s.Expires.Format("Mon Jan 2 15:04:05 MST 2006")}
	}
	return nil
}

// VerifyVersion returns ErrLowVersion if the metadata version is lower than the min version
func VerifyVersion(s *data.SignedCommon, minVersion int) error {
	if s.Version < minVersion {
		return ErrLowVersion{Actual: s.Version, Current: minVersion}
	}
	return nil
}

// VerifySignatures checks the we have sufficient valid signatures for the given role
func VerifySignatures(s *data.Signed, roleData data.BaseRole) error {
	if len(s.Signatures) == 0 {
		return ErrNoSignatures
	}

	if roleData.Threshold < 1 {
		return ErrRoleThreshold{}
	}
	logrus.Debugf("%s role has key IDs: %s", roleData.Name, strings.Join(roleData.ListKeyIDs(), ","))

	// remarshal the signed part so we can verify the signature, since the signature has
	// to be of a canonically marshalled signed object
	var decoded map[string]interface{}
	if err := json.Unmarshal(*s.Signed, &decoded); err != nil {
		return err
	}
	msg, err := json.MarshalCanonical(decoded)
	if err != nil {
		return err
	}

	valid := make(map[string]struct{})
	for i := range s.Signatures {
		sig := &(s.Signatures[i])
		logrus.Debug("verifying signature for key ID: ", sig.KeyID)
		key, ok := roleData.Keys[sig.KeyID]
		if !ok {
			logrus.Debugf("continuing b/c keyid lookup was nil: %s\n", sig.KeyID)
			continue
		}
		// Check that the signature key ID actually matches the content ID of the key
		if key.ID() != sig.KeyID {
			return ErrInvalidKeyID{}
		}
		if err := VerifySignature(msg, sig, key); err != nil {
			logrus.Debugf("continuing b/c %s", err.Error())
			continue
		}
		valid[sig.KeyID] = struct{}{}
	}
	if len(valid) < roleData.Threshold {
		return ErrRoleThreshold{
			Msg: fmt.Sprintf("valid signatures did not meet threshold for %s", roleData.Name),
		}
	}

	return nil
}

// VerifySignature checks a single signature and public key against a payload
// If the signature is verified, the signature's is valid field will actually
// be mutated to be equal to the boolean true
func VerifySignature(msg []byte, sig *data.Signature, pk data.PublicKey) error {
	// method lookup is consistent due to Unmarshal JSON doing lower case for us.
	method := sig.Method
	verifier, ok := Verifiers[method]
	if !ok {
		return fmt.Errorf("signing method is not supported: %s", sig.Method)
	}

	if err := verifier.Verify(pk, sig.Signature, msg); err != nil {
		return fmt.Errorf("signature was invalid")
	}
	sig.IsValid = true
	return nil
}

// VerifyPublicKeyMatchesPrivateKey checks if the private key and the public keys forms valid key pairs.
// This should work with both ecdsa-x509 certificate PublicKey as well as ecdsa PublicKey
func VerifyPublicKeyMatchesPrivateKey(privKey data.PrivateKey, pubKey data.PublicKey) error {
	if pubKey.Algorithm() != privKey.Algorithm() {
		return fmt.Errorf("public key and private key has different algorithms")
	}

	privKeyID := privKey.ID()
	pubKeyID, err := utils.CanonicalKeyID(pubKey)

	if err != nil {
		return fmt.Errorf("could not verify key pair: %v", err)
	}
	if pubKeyID != privKeyID {
		return fmt.Errorf("private key did not match public key")
	}
	return nil
}
<commit_msg>fix broken key pair check<commit_after>package signed

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/go/canonical/json"
	"github.com/docker/notary/tuf/data"
	"github.com/docker/notary/tuf/utils"
)

// Various basic signing errors
var (
	ErrNoSignatures = errors.New("tuf: data has no signatures")
	ErrInvalid      = errors.New("tuf: signature verification failed")
	ErrWrongType    = errors.New("tuf: meta file has wrong type")
)

// IsExpired checks if the given time passed before the present time
func IsExpired(t time.Time) bool {
	return t.Before(time.Now())
}

// VerifyExpiry returns ErrExpired if the metadata is expired
func VerifyExpiry(s *data.SignedCommon, role data.RoleName) error {
	if IsExpired(s.Expires) {
		logrus.Errorf("Metadata for %s expired", role)
		return ErrExpired{Role: role, Expired: s.Expires.Format("Mon Jan 2 15:04:05 MST 2006")}
	}
	return nil
}

// VerifyVersion returns ErrLowVersion if the metadata version is lower than the min version
func VerifyVersion(s *data.SignedCommon, minVersion int) error {
	if s.Version < minVersion {
		return ErrLowVersion{Actual: s.Version, Current: minVersion}
	}
	return nil
}

// VerifySignatures checks the we have sufficient valid signatures for the given role
func VerifySignatures(s *data.Signed, roleData data.BaseRole) error {
	if len(s.Signatures) == 0 {
		return ErrNoSignatures
	}

	if roleData.Threshold < 1 {
		return ErrRoleThreshold{}
	}
	logrus.Debugf("%s role has key IDs: %s", roleData.Name, strings.Join(roleData.ListKeyIDs(), ","))

	// remarshal the signed part so we can verify the signature, since the signature has
	// to be of a canonically marshalled signed object
	var decoded map[string]interface{}
	if err := json.Unmarshal(*s.Signed, &decoded); err != nil {
		return err
	}
	msg, err := json.MarshalCanonical(decoded)
	if err != nil {
		return err
	}

	valid := make(map[string]struct{})
	for i := range s.Signatures {
		sig := &(s.Signatures[i])
		logrus.Debug("verifying signature for key ID: ", sig.KeyID)
		key, ok := roleData.Keys[sig.KeyID]
		if !ok {
			logrus.Debugf("continuing b/c keyid lookup was nil: %s\n", sig.KeyID)
			continue
		}
		// Check that the signature key ID actually matches the content ID of the key
		if key.ID() != sig.KeyID {
			return ErrInvalidKeyID{}
		}
		if err := VerifySignature(msg, sig, key); err != nil {
			logrus.Debugf("continuing b/c %s", err.Error())
			continue
		}
		valid[sig.KeyID] = struct{}{}
	}
	if len(valid) < roleData.Threshold {
		return ErrRoleThreshold{
			Msg: fmt.Sprintf("valid signatures did not meet threshold for %s", roleData.Name),
		}
	}

	return nil
}

// VerifySignature checks a single signature and public key against a payload
// If the signature is verified, the signature's is valid field will actually
// be mutated to be equal to the boolean true
func VerifySignature(msg []byte, sig *data.Signature, pk data.PublicKey) error {
	// method lookup is consistent due to Unmarshal JSON doing lower case for us.
	method := sig.Method
	verifier, ok := Verifiers[method]
	if !ok {
		return fmt.Errorf("signing method is not supported: %s", sig.Method)
	}

	if err := verifier.Verify(pk, sig.Signature, msg); err != nil {
		return fmt.Errorf("signature was invalid")
	}
	sig.IsValid = true
	return nil
}

// VerifyPublicKeyMatchesPrivateKey checks if the private key and the public keys forms valid key pairs.
// This should work with both ecdsa-x509 certificate PublicKey as well as ecdsa PublicKey
func VerifyPublicKeyMatchesPrivateKey(privKey data.PrivateKey, pubKey data.PublicKey) error {
	pubKeyID, err := utils.CanonicalKeyID(pubKey)
	privKeyID := privKey.ID()

	if err != nil {
		return fmt.Errorf("could not verify key pair: %v", err)
	}
	if pubKeyID != privKeyID {
		return fmt.Errorf("private key did not match public key")
	}
	return nil
}
<|endoftext|>


<commit_before>// Package compose aims to provide simple "helper" methods to ease the use of
// compose (through libcompose) in (integration) tests using the go-check package.
package compose

import (
	"github.com/go-check/check"

	"github.com/docker/engine-api/types"
	"github.com/libkermit/docker/compose"
)

// Project holds compose related project attributes
type Project struct {
	project *compose.Project
}

// CreateProject creates a compose project with the given name based on the
// specified compose files
func CreateProject(c *check.C, name string, composeFiles ...string) *Project {
	project, err := compose.CreateProject(name, composeFiles...)
	c.Assert(err, check.IsNil,
		check.Commentf("error while creating compose project %s from %v", name, composeFiles))
	return &Project{
		project: project,
	}
}

// Start creates and starts the compose project.
func (p *Project) Start(c *check.C) {
	c.Assert(p.project.Start(), check.IsNil,
		check.Commentf("error while starting compose project"))
}

// Stop shuts down and clean the project
func (p *Project) Stop(c *check.C) {
	c.Assert(p.project.Stop(), check.IsNil,
		check.Commentf("error while stopping compose project"))
}

// Scale scale a service up
func (p *Project) Scale(c *check.C, service string, count int) {
	c.Assert(p.project.Scale(service, count), check.IsNil,
		check.Commentf("error while scaling the service '%s'", service))
}

// Containers lists containers for a given services.
func (p *Project) Containers(c *check.C, service string) []types.ContainerJSON {
	containers, err := p.project.Containers(service)
	c.Assert(err, check.IsNil,
		check.Commentf("error while getting the containers for service '%s'", service))
	return containers
}

// Container return the one and only container for a given services.
// It fails if there is more than one container for the service.
func (p *Project) Container(c *check.C, service string) types.ContainerJSON {
	container, err := p.project.Container(service)
	c.Assert(err, check.IsNil,
		check.Commentf("error while getting the container for service '%s'", service))
	return container
}
<commit_msg>Fix import -_-<commit_after>// Package compose aims to provide simple "helper" methods to ease the use of
// compose (through libcompose) in (integration) tests using the go-check package.
package compose

import (
	"github.com/go-check/check"

	"github.com/docker/engine-api/types"
	"github.com/libkermit/compose"
)

// Project holds compose related project attributes
type Project struct {
	project *compose.Project
}

// CreateProject creates a compose project with the given name based on the
// specified compose files
func CreateProject(c *check.C, name string, composeFiles ...string) *Project {
	project, err := compose.CreateProject(name, composeFiles...)
	c.Assert(err, check.IsNil,
		check.Commentf("error while creating compose project %s from %v", name, composeFiles))
	return &Project{
		project: project,
	}
}

// Start creates and starts the compose project.
func (p *Project) Start(c *check.C) {
	c.Assert(p.project.Start(), check.IsNil,
		check.Commentf("error while starting compose project"))
}

// Stop shuts down and clean the project
func (p *Project) Stop(c *check.C) {
	c.Assert(p.project.Stop(), check.IsNil,
		check.Commentf("error while stopping compose project"))
}

// Scale scale a service up
func (p *Project) Scale(c *check.C, service string, count int) {
	c.Assert(p.project.Scale(service, count), check.IsNil,
		check.Commentf("error while scaling the service '%s'", service))
}

// Containers lists containers for a given services.
func (p *Project) Containers(c *check.C, service string) []types.ContainerJSON {
	containers, err := p.project.Containers(service)
	c.Assert(err, check.IsNil,
		check.Commentf("error while getting the containers for service '%s'", service))
	return containers
}

// Container return the one and only container for a given services.
// It fails if there is more than one container for the service.
func (p *Project) Container(c *check.C, service string) types.ContainerJSON {
	container, err := p.project.Container(service)
	c.Assert(err, check.IsNil,
		check.Commentf("error while getting the container for service '%s'", service))
	return container
}
<|endoftext|>