<commit_before>package layer

import (
	"github.com/nbio/st"
	"gopkg.in/vinxi/vinxi.v0/utils"
	"net/http"
	"testing"
)

type vinxiHandler struct{}

func (vh vinxiHandler) HandleHTTP(w http.ResponseWriter, r *http.Request, h http.Handler) {
	w.Header().Set("foo", "bar")
	h.ServeHTTP(w, r)
}

func TestAdaptMiddlewareFunc(t *testing.T) {
	middlewareFunc := func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("foo", "bar")
			h.ServeHTTP(w, r)
		})
	}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Expect(t, w.Code, 502)
}

func TestAdaptNegroniInterface(t *testing.T) {
	middlewareFunc := func(w http.ResponseWriter, r *http.Request, h http.Handler) {
		w.Header().Set("foo", "bar")
		h.ServeHTTP(w, r)
	}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Expect(t, w.Code, 502)
}

func TestStandardHttpHandlerInterface(t *testing.T) {
	middlewareFunc := func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("foo", "bar")
	}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Reject(t, w.Code, 502)
}

func TestStandardHttpHandler(t *testing.T) {
	middlewareFunc := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("foo", "bar")
	})

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Reject(t, w.Code, 502)
}

func TestVinciHandler(t *testing.T) {
	middlewareFunc := vinxiHandler{}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Expect(t, w.Code, 502)
}
<commit_msg>test(layer): Add missing tests for adapter<commit_after>package layer

import (
	"github.com/nbio/st"
	"gopkg.in/vinxi/vinxi.v0/utils"
	"net/http"
	"testing"
)

type vinxiHandler struct{}

func (vh vinxiHandler) HandleHTTP(w http.ResponseWriter, r *http.Request, h http.Handler) {
	w.Header().Set("foo", "bar")
	h.ServeHTTP(w, r)
}

func TestAdaptMiddlewareFunc(t *testing.T) {
	middlewareFunc := func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("foo", "bar")
			h.ServeHTTP(w, r)
		})
	}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Expect(t, w.Code, 502)
}

func TestAdaptMiddlewareHandlerFunc(t *testing.T) {
	middlewareFunc := func(h http.Handler) func(http.ResponseWriter, *http.Request) {
		return func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("foo", "bar")
			h.ServeHTTP(w, r)
		}
	}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Expect(t, w.Code, 502)
}

func TestAdaptNegroniInterface(t *testing.T) {
	middlewareFunc := func(w http.ResponseWriter, r *http.Request, h http.Handler) {
		w.Header().Set("foo", "bar")
		h.ServeHTTP(w, r)
	}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Expect(t, w.Code, 502)
}

func TestStandardHttpHandlerInterface(t *testing.T) {
	middlewareFunc := func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("foo", "bar")
	}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Reject(t, w.Code, 502)
}

func TestStandardHttpHandler(t *testing.T) {
	middlewareFunc := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("foo", "bar")
	})

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Reject(t, w.Code, 502)
}

func TestVinciHandler(t *testing.T) {
	middlewareFunc := vinxiHandler{}

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(middlewareFunc)
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Expect(t, w.Code, 502)
}

type partialHandler struct{}

func (ph *partialHandler) HandleHTTP(w http.ResponseWriter, r *http.Request, h http.Handler) {
	w.Header().Set("foo", "bar")
	h.ServeHTTP(w, r)
}

func TestPartialHandler(t *testing.T) {

	w := utils.NewWriterStub()
	req := &http.Request{}

	adaptedFunc := AdaptFunc(&partialHandler{})
	adaptedFunc(FinalHandler).ServeHTTP(w, req)

	st.Expect(t, w.Header().Get("foo"), "bar")
	st.Expect(t, w.Code, 502)
}
<|endoftext|>


<commit_before>package boltql

import (
	"bytes"
	"errors"
	"fmt"
	"sort"

	"github.com/boltdb/bolt"
	"github.com/gobs/typedbuffer"
)

var (
	NO_TABLE         = bolt.ErrBucketNotFound
	NO_INDEX         = bolt.ErrBucketNotFound
	ALREADY_EXISTS   = bolt.ErrBucketExists
	NO_SCHEMA        = errors.New("no schema for table")
	SCHEMA_CORRUPTED = errors.New("schema corrupted")
	NO_KEY           = errors.New("key not found")
	BAD_VALUES       = errors.New("bad values")

	// this is just a marker for auto-increment fields
	AUTOINCREMENT = &struct{}{}
)

//
// A DataStore is the main interface to a BoltDB database
//
type DataStore bolt.DB

//
// A DataRecord is the interface for elements that can be stored in a table.
// A DataRecord needs to implement two methods:
//
// ToFieldList() should convert the record fields into a list of values, preserving their order
// (i.e. the same fields should always be in the same position)
//
// FromFieldList() should fill the record with values from the input list. The order of values
// is the same as what was returned from ToFieldList()
//
type DataRecord interface {
	ToFieldList() []interface{}
	FromFieldList([]interface{})
}

//
// Open the database (create if it doesn't exist)
//
func Open(dbfile string) (*DataStore, error) {
	db, err := bolt.Open(dbfile, 0666, nil)
	if err != nil {
		return nil, err
	}

	return (*DataStore)(db), nil
}

//
// Close the database
//
func (d *DataStore) Close() error {
	db := (*bolt.DB)(d)
	return db.Close()
}

func (d *DataStore) SetBulk(b bool) {
	db := (*bolt.DB)(d)
	db.NoSync = b
}

func indices(name string) []byte {
	return []byte(name + "_idx")
}

func schema(name string) []byte {
	return []byte(name)
}

//
// A Table is a container for the table name and indices
//
type Table struct {
	name    string
	indices map[string]indexinfo

	d *DataStore
}

//
// Implement the Stringer interface
//
func (t *Table) String() string {
	return fmt.Sprintf("Table{name: %q, indices: %v}", t.name, t.indices)
}

type indexinfo struct {
	nilFirst bool
	iplist   []indexpos
}

type indexpos struct {
	field uint
	pos   uint
}

type iplist []indexpos

func makeIndexPos(fields []uint64) iplist {
	ip := make(iplist, len(fields))

	for i, f := range fields {
		ip[i].field = uint(f)
		ip[i].pos = uint(i)
	}

	sort.Sort(ip)
	return ip
}

func (ip iplist) Len() int {
	return len(ip)
}

func (ip iplist) Less(i, j int) bool {
	return ip[i].field < ip[j].field
}

func (ip iplist) Swap(i, j int) {
	ip[i], ip[j] = ip[j], ip[i]
}

//
// Create table if doesn't exist
//
func (d *DataStore) CreateTable(name string) (*Table, error) {
	db := (*bolt.DB)(d)

	err := db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucket(schema(name))
		if err != nil {
			return err
		}

		return nil
	})

	if err == nil {
		return &Table{name: name, indices: map[string]indexinfo{}, d: d}, nil
	} else {
		return nil, err
	}
}

//
// Get existing Table
//
func (d *DataStore) GetTable(name string) (*Table, error) {
	db := (*bolt.DB)(d)
	table := Table{name: name, indices: map[string]indexinfo{}, d: d}

	err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(schema(name))
		if b == nil {
			return NO_TABLE
		}

		b.ForEach(func(k, v []byte) error {
			name := string(k)

			nilFirst, rest, err := typedbuffer.Decode(v)
			if err != nil {
				return SCHEMA_CORRUPTED
			}
			fields, err := typedbuffer.DecodeUintArray(rest)
			if err != nil {
				return SCHEMA_CORRUPTED
			}

			table.indices[name] = indexinfo{
				nilFirst: nilFirst.(bool),
				iplist:   makeIndexPos(fields),
			}

			return nil
		})

		return nil
	})

	if err == nil {
		return &table, nil
	} else {
		return nil, err
	}
}

//
// Create an index given the name (index) and a list of field positions
// used to create a composite key.
//
// nilFirst specifies if nil values should sort first (lowest possible value) or last (highest possible value)
//
// The field position should corrispond to the entries in DataRecord ToFieldList() and FromFieldList()
//
func (t *Table) CreateIndex(index string, nilFirst bool, fields ...uint64) error {
	db := (*bolt.DB)(t.d)

	err := db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(schema(t.name))
		if b == nil {
			return NO_TABLE
		}

		b1, err := typedbuffer.Encode(nilFirst)
		if err != nil {
			return BAD_VALUES
		}
		b2, err := typedbuffer.Encode(fields)
		if err != nil {
			return BAD_VALUES
		}

		enc := append(b1, b2...)
		if err := b.Put([]byte(index), enc); err != nil {
			return err
		}

		if _, err := tx.CreateBucket(indices(index)); err != nil {
			return err
		}

		return nil
	})

	if err == nil {
		t.indices[index] = indexinfo{
			nilFirst: nilFirst,
			iplist:   makeIndexPos(fields),
		}
	}

	return err
}

func (info indexinfo) marshalKeyValue(fields []interface{}) (key, value []byte, err error) {
	if len(info.iplist) == 0 {
		return
	}

	vkey := make([]interface{}, len(info.iplist))
	vval := make([]interface{}, 0)

	kk, lk := 0, len(info.iplist)

	for fi, fv := range fields {
		if kk < lk && uint(fi) == info.iplist[kk].field {
			vkey[info.iplist[kk].pos] = fv
			kk += 1
		} else {
			vval = append(vval, fv)
		}
	}

	if len(vkey) > 0 {
		if key, err = typedbuffer.EncodeNils(info.nilFirst, vkey...); err != nil {
			return
		}
	}

	if len(vval) > 0 {
		value, err = typedbuffer.EncodeNils(info.nilFirst, vval...)
	}

	return
}

func (info indexinfo) unmarshalKeyValue(k, v []byte) ([]interface{}, error) {
	vkey, err := typedbuffer.DecodeAll(false, k)
	if err != nil {
		return nil, err
	}

	vval, err := typedbuffer.DecodeAll(false, v)
	if err != nil {
		return nil, err
	}

	lkey := len(vkey)
	lval := len(vval)

	fields := []interface{}{}

	var ival interface{}

	kk, lk := 0, len(info.iplist)

	for i := 0; i < lkey+lval; i++ {
		if kk < lk && uint(i) == info.iplist[kk].field {
			ival = vkey[info.iplist[kk].pos]
			kk += 1
		} else {
			ival = vval[0]
			vval = vval[1:]
		}

		fields = append(fields, ival)
	}

	return fields, nil
}

//
// Add a record to the table, updating all indices.
// If a record with the same key exists, it's updated.
//
func (t *Table) Put(rec DataRecord) (uint64, error) {
	db := (*bolt.DB)(t.d)

	var key uint64

	err := db.Update(func(tx *bolt.Tx) (err error) {
		b := tx.Bucket([]byte(t.name))
		if b == nil {
			return NO_TABLE
		}

		fields := rec.ToFieldList()

		for i := range fields {
			if fields[i] == AUTOINCREMENT {
				fields[i], err = b.NextSequence()
				if err != nil {
					return err
				}
			}
		}

		for index, info := range t.indices {
			ib := tx.Bucket(indices(index))
			if ib == nil {
				return NO_TABLE
			}

			key, val, err := info.marshalKeyValue(fields)
			if err != nil {
				return err
			}

			if key == nil {
				continue
			}

			if err := ib.Put(key, val); err != nil {
				return err
			}
		}

		return nil
	})

	return key, err
}

//
// Get a record from the table, given the index and the key
//
func (t *Table) Get(index string, key, res DataRecord) error {
	db := (*bolt.DB)(t.d)

	err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(indices(index))
		if b == nil {
			return NO_INDEX
		}

		c := b.Cursor()

		info := t.indices[index]

		sk, _, err := info.marshalKeyValue(key.ToFieldList())
		if err != nil {
			return err
		}

		if sk == nil {
			return NO_KEY
		}

		resk, resv := c.Seek(sk)
		if !bytes.Equal(sk, resk) {
			return NO_KEY
		}

		fields, err := info.unmarshalKeyValue(resk, resv)
		if err != nil {
			return err
		}

		res.FromFieldList(fields)
		return nil
	})

	return err
}

//
// Delete a record from the table, given the index and the key
//
func (t *Table) Delete(index string, key DataRecord) error {
	db := (*bolt.DB)(t.d)

	err := db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(indices(index))
		if b == nil {
			return NO_INDEX
		}

		info := t.indices[index]

		sk, _, err := info.marshalKeyValue(key.ToFieldList())
		if err != nil {
			return err
		}

		if sk == nil {
			return NO_KEY
		}

		c := b.Cursor()
		k, v := c.Seek(sk)

		// Seek will return the next key if there is no match
		// so make sure we check we got the right record

		if bytes.Equal(sk, k) {
			if err := c.Delete(); err != nil {
				return err
			}

			fields, err := info.unmarshalKeyValue(k, v)
			if err != nil {
				return err
			}

			for i, info := range t.indices {
				if i == index {
					// already done
					continue
				}

				b := tx.Bucket(indices(i))
				if b == nil {
					continue
				}

				// could use marshalKeyValue() instead

				vkey := make([]interface{}, len(info.iplist))
				for _, ip := range info.iplist {
					vkey[ip.pos] = fields[ip.field]
				}

				dkey, err := typedbuffer.Encode(vkey...)
				if err != nil {
					return err
				}

				if err := b.Delete(dkey); err != nil {
					return err
				}
			}
		}

		return nil
	})

	return err
}

//
// Get all records sorted by index keys (ascending or descending)
// Call user function with record content or error
//
func (t *Table) Scan(index string, ascending bool, start, res DataRecord, callback func(DataRecord, error) bool) error {
	db := (*bolt.DB)(t.d)

	return db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(indices(index))
		if b == nil {
			return NO_INDEX
		}

		c := b.Cursor()

		info := t.indices[index]

		var k, v []byte

		if start != nil {
			key, _, err := info.marshalKeyValue(start.ToFieldList())
			if err != nil {
				return err
			}

			if key != nil {
				k, v = c.Seek(key)
				if !ascending && !bytes.Equal(key, k) {
					// if descending and keys don't match we want to start from the first key
					// in range (previous)

					k, v = c.Prev()
				}
			}
		}

		if k == nil {
			if ascending {
				k, v = c.First()
			} else {
				k, v = c.Last()
			}
		}

		var next func() (key []byte, value []byte)

		if ascending {
			next = c.Next
		} else {
			next = c.Prev
		}

		for ; k != nil; k, v = next() {
			fields, err := info.unmarshalKeyValue(k, v)
			if err != nil {
				return err
			}

			res.FromFieldList(fields)

			if !callback(res, err) {
				break
			}
		}

		return nil
	})
}

//
// Scan through all records in an index. Calls specified callback with key and value (as []byte)
//
func (t *Table) ForEach(index string, callback func(k, v []byte) error) error {
	db := (*bolt.DB)(t.d)

	return db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(t.name))
		if len(index) > 0 {
			b = tx.Bucket(indices(index))
		}

		if b == nil {
			return NO_INDEX
		}

		return b.ForEach(callback)
	})
}
<commit_msg>some documentation for marshal/unmarshal fields<commit_after>package boltql

import (
	"bytes"
	"errors"
	"fmt"
	"sort"

	"github.com/boltdb/bolt"
	"github.com/gobs/typedbuffer"
)

var (
	NO_TABLE         = bolt.ErrBucketNotFound
	NO_INDEX         = bolt.ErrBucketNotFound
	ALREADY_EXISTS   = bolt.ErrBucketExists
	NO_SCHEMA        = errors.New("no schema for table")
	SCHEMA_CORRUPTED = errors.New("schema corrupted")
	NO_KEY           = errors.New("key not found")
	BAD_VALUES       = errors.New("bad values")

	// this is just a marker for auto-increment fields
	AUTOINCREMENT = &struct{}{}
)

//
// A DataStore is the main interface to a BoltDB database
//
type DataStore bolt.DB

//
// A DataRecord is the interface for elements that can be stored in a table.
// A DataRecord needs to implement two methods:
//
// ToFieldList() should convert the record fields into a list of values, preserving their order
// (i.e. the same fields should always be in the same position)
//
// FromFieldList() should fill the record with values from the input list. The order of values
// is the same as what was returned from ToFieldList()
//
type DataRecord interface {
	ToFieldList() []interface{}
	FromFieldList([]interface{})
}

//
// Open the database (create if it doesn't exist)
//
func Open(dbfile string) (*DataStore, error) {
	db, err := bolt.Open(dbfile, 0666, nil)
	if err != nil {
		return nil, err
	}

	return (*DataStore)(db), nil
}

//
// Close the database
//
func (d *DataStore) Close() error {
	db := (*bolt.DB)(d)
	return db.Close()
}

func (d *DataStore) SetBulk(b bool) {
	db := (*bolt.DB)(d)
	db.NoSync = b
}

func indices(name string) []byte {
	return []byte(name + "_idx")
}

func schema(name string) []byte {
	return []byte(name)
}

//
// A Table is a container for the table name and indices
//
type Table struct {
	name    string
	indices map[string]indexinfo

	d *DataStore
}

//
// Implement the Stringer interface
//
func (t *Table) String() string {
	return fmt.Sprintf("Table{name: %q, indices: %v}", t.name, t.indices)
}

type indexinfo struct {
	nilFirst bool
	iplist   []indexpos
}

type indexpos struct {
	field uint
	pos   uint
}

type iplist []indexpos

func makeIndexPos(fields []uint64) iplist {
	ip := make(iplist, len(fields))

	for i, f := range fields {
		ip[i].field = uint(f)
		ip[i].pos = uint(i)
	}

	sort.Sort(ip)
	return ip
}

func (ip iplist) Len() int {
	return len(ip)
}

func (ip iplist) Less(i, j int) bool {
	return ip[i].field < ip[j].field
}

func (ip iplist) Swap(i, j int) {
	ip[i], ip[j] = ip[j], ip[i]
}

//
// Create table if doesn't exist
//
func (d *DataStore) CreateTable(name string) (*Table, error) {
	db := (*bolt.DB)(d)

	err := db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucket(schema(name))
		if err != nil {
			return err
		}

		return nil
	})

	if err == nil {
		return &Table{name: name, indices: map[string]indexinfo{}, d: d}, nil
	} else {
		return nil, err
	}
}

//
// Get existing Table
//
func (d *DataStore) GetTable(name string) (*Table, error) {
	db := (*bolt.DB)(d)
	table := Table{name: name, indices: map[string]indexinfo{}, d: d}

	err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(schema(name))
		if b == nil {
			return NO_TABLE
		}

		b.ForEach(func(k, v []byte) error {
			name := string(k)

			nilFirst, rest, err := typedbuffer.Decode(v)
			if err != nil {
				return SCHEMA_CORRUPTED
			}
			fields, err := typedbuffer.DecodeUintArray(rest)
			if err != nil {
				return SCHEMA_CORRUPTED
			}

			table.indices[name] = indexinfo{
				nilFirst: nilFirst.(bool),
				iplist:   makeIndexPos(fields),
			}

			return nil
		})

		return nil
	})

	if err == nil {
		return &table, nil
	} else {
		return nil, err
	}
}

//
// Create an index given the name (index) and a list of field positions
// used to create a composite key.
//
// nilFirst specifies if nil values should sort first (lowest possible value) or last (highest possible value)
//
// The field position should corrispond to the entries in DataRecord ToFieldList() and FromFieldList()
//
func (t *Table) CreateIndex(index string, nilFirst bool, fields ...uint64) error {
	db := (*bolt.DB)(t.d)

	err := db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(schema(t.name))
		if b == nil {
			return NO_TABLE
		}

		b1, err := typedbuffer.Encode(nilFirst)
		if err != nil {
			return BAD_VALUES
		}
		b2, err := typedbuffer.Encode(fields)
		if err != nil {
			return BAD_VALUES
		}

		enc := append(b1, b2...)
		if err := b.Put([]byte(index), enc); err != nil {
			return err
		}

		if _, err := tx.CreateBucket(indices(index)); err != nil {
			return err
		}

		return nil
	})

	if err == nil {
		t.indices[index] = indexinfo{
			nilFirst: nilFirst,
			iplist:   makeIndexPos(fields),
		}
	}

	return err
}

//
// marshal an array of fields into a key and value pair of encoded values
//
// the key is a composed key of the fields described in info.iplist (field number and order)
// the value is a collection of the remaning fields
//
func (info indexinfo) marshalKeyValue(fields []interface{}) (key, value []byte, err error) {
	if len(info.iplist) == 0 {
		return
	}

	vkey := make([]interface{}, len(info.iplist))
	vval := make([]interface{}, 0)

	kk, lk := 0, len(info.iplist)

	for fi, fv := range fields {
		if kk < lk && uint(fi) == info.iplist[kk].field {
			vkey[info.iplist[kk].pos] = fv
			kk += 1
		} else {
			vval = append(vval, fv)
		}
	}

	if len(vkey) > 0 {
		if key, err = typedbuffer.EncodeNils(info.nilFirst, vkey...); err != nil {
			return
		}
	}

	if len(vval) > 0 {
		value, err = typedbuffer.EncodeNils(info.nilFirst, vval...)
	}

	return
}

//
// unmarshal key, value into a list of decoded fields
//
func (info indexinfo) unmarshalKeyValue(k, v []byte) ([]interface{}, error) {
	vkey, err := typedbuffer.DecodeAll(false, k)
	if err != nil {
		return nil, err
	}

	vval, err := typedbuffer.DecodeAll(false, v)
	if err != nil {
		return nil, err
	}

	lkey := len(vkey)
	lval := len(vval)

	fields := []interface{}{}

	var ival interface{}

	kk, lk := 0, len(info.iplist)

	for i := 0; i < lkey+lval; i++ {
		if kk < lk && uint(i) == info.iplist[kk].field {
			ival = vkey[info.iplist[kk].pos]
			kk += 1
		} else {
			ival = vval[0]
			vval = vval[1:]
		}

		fields = append(fields, ival)
	}

	return fields, nil
}

//
// Add a record to the table, updating all indices.
// If a record with the same key exists, it's updated.
//
func (t *Table) Put(rec DataRecord) (uint64, error) {
	db := (*bolt.DB)(t.d)

	var key uint64

	err := db.Update(func(tx *bolt.Tx) (err error) {
		b := tx.Bucket([]byte(t.name))
		if b == nil {
			return NO_TABLE
		}

		fields := rec.ToFieldList()

		for i := range fields {
			if fields[i] == AUTOINCREMENT {
				fields[i], err = b.NextSequence()
				if err != nil {
					return err
				}
			}
		}

		for index, info := range t.indices {
			ib := tx.Bucket(indices(index))
			if ib == nil {
				return NO_TABLE
			}

			key, val, err := info.marshalKeyValue(fields)
			if err != nil {
				return err
			}

			if key == nil {
				continue
			}

			if err := ib.Put(key, val); err != nil {
				return err
			}
		}

		return nil
	})

	return key, err
}

//
// Get a record from the table, given the index and the key
//
func (t *Table) Get(index string, key, res DataRecord) error {
	db := (*bolt.DB)(t.d)

	err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(indices(index))
		if b == nil {
			return NO_INDEX
		}

		c := b.Cursor()

		info := t.indices[index]

		sk, _, err := info.marshalKeyValue(key.ToFieldList())
		if err != nil {
			return err
		}

		if sk == nil {
			return NO_KEY
		}

		resk, resv := c.Seek(sk)
		if !bytes.Equal(sk, resk) {
			return NO_KEY
		}

		fields, err := info.unmarshalKeyValue(resk, resv)
		if err != nil {
			return err
		}

		res.FromFieldList(fields)
		return nil
	})

	return err
}

//
// Delete a record from the table, given the index and the key
//
func (t *Table) Delete(index string, key DataRecord) error {
	db := (*bolt.DB)(t.d)

	err := db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(indices(index))
		if b == nil {
			return NO_INDEX
		}

		info := t.indices[index]

		sk, _, err := info.marshalKeyValue(key.ToFieldList())
		if err != nil {
			return err
		}

		if sk == nil {
			return NO_KEY
		}

		c := b.Cursor()
		k, v := c.Seek(sk)

		// Seek will return the next key if there is no match
		// so make sure we check we got the right record

		if bytes.Equal(sk, k) {
			if err := c.Delete(); err != nil {
				return err
			}

			fields, err := info.unmarshalKeyValue(k, v)
			if err != nil {
				return err
			}

			for i, info := range t.indices {
				if i == index {
					// already done
					continue
				}

				b := tx.Bucket(indices(i))
				if b == nil {
					continue
				}

				// could use marshalKeyValue() instead

				vkey := make([]interface{}, len(info.iplist))
				for _, ip := range info.iplist {
					vkey[ip.pos] = fields[ip.field]
				}

				dkey, err := typedbuffer.Encode(vkey...)
				if err != nil {
					return err
				}

				if err := b.Delete(dkey); err != nil {
					return err
				}
			}
		}

		return nil
	})

	return err
}

//
// Get all records sorted by index keys (ascending or descending)
// Call user function with record content or error
//
func (t *Table) Scan(index string, ascending bool, start, res DataRecord, callback func(DataRecord, error) bool) error {
	db := (*bolt.DB)(t.d)

	return db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(indices(index))
		if b == nil {
			return NO_INDEX
		}

		c := b.Cursor()

		info := t.indices[index]

		var k, v []byte

		if start != nil {
			key, _, err := info.marshalKeyValue(start.ToFieldList())
			if err != nil {
				return err
			}

			if key != nil {
				k, v = c.Seek(key)
				if !ascending && !bytes.Equal(key, k) {
					// if descending and keys don't match we want to start from the first key
					// in range (previous)

					k, v = c.Prev()
				}
			}
		}

		if k == nil {
			if ascending {
				k, v = c.First()
			} else {
				k, v = c.Last()
			}
		}

		var next func() (key []byte, value []byte)

		if ascending {
			next = c.Next
		} else {
			next = c.Prev
		}

		for ; k != nil; k, v = next() {
			fields, err := info.unmarshalKeyValue(k, v)
			if err != nil {
				return err
			}

			res.FromFieldList(fields)

			if !callback(res, err) {
				break
			}
		}

		return nil
	})
}

//
// Scan through all records in an index. Calls specified callback with key and value (as []byte, not decoded)
//
func (t *Table) ForEach(index string, callback func(k, v []byte) error) error {
	db := (*bolt.DB)(t.d)

	return db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(t.name))
		if len(index) > 0 {
			b = tx.Bucket(indices(index))
		}

		if b == nil {
			return NO_INDEX
		}

		return b.ForEach(callback)
	})
}
<|endoftext|>


<commit_before>// Package rest implements a simple REST wrapper
//
// All methods are safe for concurrent calling.
package rest

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/url"
	"sync"

	"github.com/ncw/rclone/fs"
	"github.com/pkg/errors"
)

// Client contains the info to sustain the API
type Client struct {
	mu           sync.RWMutex
	c            *http.Client
	rootURL      string
	errorHandler func(resp *http.Response) error
	headers      map[string]string
}

// NewClient takes an oauth http.Client and makes a new api instance
func NewClient(c *http.Client) *Client {
	api := &Client{
		c:            c,
		errorHandler: defaultErrorHandler,
		headers:      make(map[string]string),
	}
	return api
}

// ReadBody reads resp.Body into result, closing the body
func ReadBody(resp *http.Response) (result []byte, err error) {
	defer fs.CheckClose(resp.Body, &err)
	return ioutil.ReadAll(resp.Body)
}

// defaultErrorHandler doesn't attempt to parse the http body, just
// returns it in the error message
func defaultErrorHandler(resp *http.Response) (err error) {
	body, err := ReadBody(resp)
	if err != nil {
		return errors.Wrap(err, "error reading error out of body")
	}
	return errors.Errorf("HTTP error %v (%v) returned body: %q", resp.StatusCode, resp.Status, body)
}

// SetErrorHandler sets the handler to decode an error response when
// the HTTP status code is not 2xx.  The handler should close resp.Body.
func (api *Client) SetErrorHandler(fn func(resp *http.Response) error) *Client {
	api.mu.Lock()
	defer api.mu.Unlock()
	api.errorHandler = fn
	return api
}

// SetRoot sets the default RootURL.  You can override this on a per
// call basis using the RootURL field in Opts.
func (api *Client) SetRoot(RootURL string) *Client {
	api.mu.Lock()
	defer api.mu.Unlock()
	api.rootURL = RootURL
	return api
}

// SetHeader sets a header for all requests
func (api *Client) SetHeader(key, value string) *Client {
	api.mu.Lock()
	defer api.mu.Unlock()
	api.headers[key] = value
	return api
}

// Opts contains parameters for Call, CallJSON etc
type Opts struct {
	Method                string // GET, POST etc
	Path                  string // relative to RootURL
	RootURL               string // override RootURL passed into SetRoot()
	Body                  io.Reader
	NoResponse            bool // set to close Body
	ContentType           string
	ContentLength         *int64
	ContentRange          string
	ExtraHeaders          map[string]string
	UserName              string // username for Basic Auth
	Password              string // password for Basic Auth
	Options               []fs.OpenOption
	IgnoreStatus          bool       // if set then we don't check error status or parse error body
	MultipartMetadataName string     // set the following 3 vars
	MultipartContentName  string     // and Body and pass in request
	MultipartFileName     string     // for multipart upload
	Parameters            url.Values // any parameters for the final URL
}

// Copy creates a copy of the options
func (o *Opts) Copy() *Opts {
	newOpts := *o
	return &newOpts
}

// DecodeJSON decodes resp.Body into result
func DecodeJSON(resp *http.Response, result interface{}) (err error) {
	defer fs.CheckClose(resp.Body, &err)
	decoder := json.NewDecoder(resp.Body)
	return decoder.Decode(result)
}

// ClientWithHeaderReset makes a new http client which resets the
// headers passed in on redirect
//
// FIXME This is now unecessary with go1.8
func ClientWithHeaderReset(c *http.Client, headers map[string]string) *http.Client {
	if len(headers) == 0 {
		return c
	}
	clientCopy := *c
	clientCopy.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		if len(via) >= 10 {
			return errors.New("stopped after 10 redirects")
		}
		// Reset the headers in the new request
		for k, v := range headers {
			if v != "" {
				req.Header.Set(k, v)
			}
		}
		return nil
	}
	return &clientCopy
}

// Call makes the call and returns the http.Response
//
// if err != nil then resp.Body will need to be closed
//
// it will return resp if at all possible, even if err is set
func (api *Client) Call(opts *Opts) (resp *http.Response, err error) {
	api.mu.RLock()
	defer api.mu.RUnlock()
	if opts == nil {
		return nil, errors.New("call() called with nil opts")
	}
	url := api.rootURL
	if opts.RootURL != "" {
		url = opts.RootURL
	}
	if url == "" {
		return nil, errors.New("RootURL not set")
	}
	url += opts.Path
	if opts.Parameters != nil && len(opts.Parameters) > 0 {
		url += "?" + opts.Parameters.Encode()
	}
	req, err := http.NewRequest(opts.Method, url, opts.Body)
	if err != nil {
		return
	}
	headers := make(map[string]string)
	// Set default headers
	for k, v := range api.headers {
		headers[k] = v
	}
	if opts.ContentType != "" {
		headers["Content-Type"] = opts.ContentType
	}
	if opts.ContentLength != nil {
		req.ContentLength = *opts.ContentLength
	}
	if opts.ContentRange != "" {
		headers["Content-Range"] = opts.ContentRange
	}
	// Set any extra headers
	if opts.ExtraHeaders != nil {
		for k, v := range opts.ExtraHeaders {
			headers[k] = v
		}
	}
	// add any options to the headers
	fs.OpenOptionAddHeaders(opts.Options, headers)
	// Now set the headers
	for k, v := range headers {
		if v != "" {
			req.Header.Add(k, v)
		}
	}
	if opts.UserName != "" || opts.Password != "" {
		req.SetBasicAuth(opts.UserName, opts.Password)
	}
	c := ClientWithHeaderReset(api.c, headers)
	api.mu.RUnlock()
	resp, err = c.Do(req)
	api.mu.RLock()
	if err != nil {
		return nil, err
	}
	if !opts.IgnoreStatus {
		if resp.StatusCode < 200 || resp.StatusCode > 299 {
			return resp, api.errorHandler(resp)
		}
	}
	if opts.NoResponse {
		return resp, resp.Body.Close()
	}
	return resp, nil
}

// CallJSON runs Call and decodes the body as a JSON object into response (if not nil)
//
// If request is not nil then it will be JSON encoded as the body of the request
//
// It will return resp if at all possible, even if err is set
func (api *Client) CallJSON(opts *Opts, request interface{}, response interface{}) (resp *http.Response, err error) {
	var requestBody []byte
	// Marshal the request if given
	if request != nil {
		opts = opts.Copy()
		requestBody, err = json.Marshal(request)
		opts.ContentType = "application/json"
		if err != nil {
			return nil, err
		}
		// Set the body up as a JSON object if required
		if opts.Body == nil {
			opts.Body = bytes.NewBuffer(requestBody)
		}
	}
	errChan := make(chan error, 1)
	isMultipart := opts.MultipartMetadataName != "" && opts.Body != nil && request != nil
	if isMultipart {
		bodyReader, bodyWriter := io.Pipe()
		writer := multipart.NewWriter(bodyWriter)
		opts.ContentType = writer.FormDataContentType()
		in := opts.Body
		opts.Body = bodyReader
		go func() {
			defer func() { _ = bodyWriter.Close() }()
			var err error

			// Create the first part
			err = writer.WriteField(opts.MultipartMetadataName, string(requestBody))
			if err != nil {
				errChan <- err
				return
			}

			// Add the file part
			part, err := writer.CreateFormFile(opts.MultipartContentName, opts.MultipartFileName)
			if err != nil {
				errChan <- err
				return
			}

			// Copy it in
			if _, err := io.Copy(part, in); err != nil {
				errChan <- err
				return
			}
			errChan <- writer.Close()
		}()
	}
	resp, err = api.Call(opts)
	if err != nil {
		return resp, err
	}
	if response == nil || opts.NoResponse {
		return resp, nil
	}
	if isMultipart {
		err = <-errChan
		if err != nil {
			return resp, err
		}
	}
	err = DecodeJSON(resp, response)
	return resp, err
}
<commit_msg>rest: Add TransferEncoding and Close parameters<commit_after>// Package rest implements a simple REST wrapper
//
// All methods are safe for concurrent calling.
package rest

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/url"
	"sync"

	"github.com/ncw/rclone/fs"
	"github.com/pkg/errors"
)

// Client contains the info to sustain the API
type Client struct {
	mu           sync.RWMutex
	c            *http.Client
	rootURL      string
	errorHandler func(resp *http.Response) error
	headers      map[string]string
}

// NewClient takes an oauth http.Client and makes a new api instance
func NewClient(c *http.Client) *Client {
	api := &Client{
		c:            c,
		errorHandler: defaultErrorHandler,
		headers:      make(map[string]string),
	}
	return api
}

// ReadBody reads resp.Body into result, closing the body
func ReadBody(resp *http.Response) (result []byte, err error) {
	defer fs.CheckClose(resp.Body, &err)
	return ioutil.ReadAll(resp.Body)
}

// defaultErrorHandler doesn't attempt to parse the http body, just
// returns it in the error message
func defaultErrorHandler(resp *http.Response) (err error) {
	body, err := ReadBody(resp)
	if err != nil {
		return errors.Wrap(err, "error reading error out of body")
	}
	return errors.Errorf("HTTP error %v (%v) returned body: %q", resp.StatusCode, resp.Status, body)
}

// SetErrorHandler sets the handler to decode an error response when
// the HTTP status code is not 2xx.  The handler should close resp.Body.
func (api *Client) SetErrorHandler(fn func(resp *http.Response) error) *Client {
	api.mu.Lock()
	defer api.mu.Unlock()
	api.errorHandler = fn
	return api
}

// SetRoot sets the default RootURL.  You can override this on a per
// call basis using the RootURL field in Opts.
func (api *Client) SetRoot(RootURL string) *Client {
	api.mu.Lock()
	defer api.mu.Unlock()
	api.rootURL = RootURL
	return api
}

// SetHeader sets a header for all requests
func (api *Client) SetHeader(key, value string) *Client {
	api.mu.Lock()
	defer api.mu.Unlock()
	api.headers[key] = value
	return api
}

// Opts contains parameters for Call, CallJSON etc
type Opts struct {
	Method                string // GET, POST etc
	Path                  string // relative to RootURL
	RootURL               string // override RootURL passed into SetRoot()
	Body                  io.Reader
	NoResponse            bool // set to close Body
	ContentType           string
	ContentLength         *int64
	ContentRange          string
	ExtraHeaders          map[string]string
	UserName              string // username for Basic Auth
	Password              string // password for Basic Auth
	Options               []fs.OpenOption
	IgnoreStatus          bool       // if set then we don't check error status or parse error body
	MultipartMetadataName string     // set the following 3 vars
	MultipartContentName  string     // and Body and pass in request
	MultipartFileName     string     // for multipart upload
	Parameters            url.Values // any parameters for the final URL
	TransferEncoding      []string   // transfer encoding, set to "identity" to disable chunked encoding
	Close                 bool       // set to close the connection after this transaction
}

// Copy creates a copy of the options
func (o *Opts) Copy() *Opts {
	newOpts := *o
	return &newOpts
}

// DecodeJSON decodes resp.Body into result
func DecodeJSON(resp *http.Response, result interface{}) (err error) {
	defer fs.CheckClose(resp.Body, &err)
	decoder := json.NewDecoder(resp.Body)
	return decoder.Decode(result)
}

// ClientWithHeaderReset makes a new http client which resets the
// headers passed in on redirect
//
// FIXME This is now unecessary with go1.8
func ClientWithHeaderReset(c *http.Client, headers map[string]string) *http.Client {
	if len(headers) == 0 {
		return c
	}
	clientCopy := *c
	clientCopy.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		if len(via) >= 10 {
			return errors.New("stopped after 10 redirects")
		}
		// Reset the headers in the new request
		for k, v := range headers {
			if v != "" {
				req.Header.Set(k, v)
			}
		}
		return nil
	}
	return &clientCopy
}

// Call makes the call and returns the http.Response
//
// if err != nil then resp.Body will need to be closed
//
// it will return resp if at all possible, even if err is set
func (api *Client) Call(opts *Opts) (resp *http.Response, err error) {
	api.mu.RLock()
	defer api.mu.RUnlock()
	if opts == nil {
		return nil, errors.New("call() called with nil opts")
	}
	url := api.rootURL
	if opts.RootURL != "" {
		url = opts.RootURL
	}
	if url == "" {
		return nil, errors.New("RootURL not set")
	}
	url += opts.Path
	if opts.Parameters != nil && len(opts.Parameters) > 0 {
		url += "?" + opts.Parameters.Encode()
	}
	req, err := http.NewRequest(opts.Method, url, opts.Body)
	if err != nil {
		return
	}
	headers := make(map[string]string)
	// Set default headers
	for k, v := range api.headers {
		headers[k] = v
	}
	if opts.ContentType != "" {
		headers["Content-Type"] = opts.ContentType
	}
	if opts.ContentLength != nil {
		req.ContentLength = *opts.ContentLength
	}
	if opts.ContentRange != "" {
		headers["Content-Range"] = opts.ContentRange
	}
	if len(opts.TransferEncoding) != 0 {
		req.TransferEncoding = opts.TransferEncoding
	}
	if opts.Close {
		req.Close = true
	}
	// Set any extra headers
	if opts.ExtraHeaders != nil {
		for k, v := range opts.ExtraHeaders {
			headers[k] = v
		}
	}
	// add any options to the headers
	fs.OpenOptionAddHeaders(opts.Options, headers)
	// Now set the headers
	for k, v := range headers {
		if v != "" {
			req.Header.Add(k, v)
		}
	}
	if opts.UserName != "" || opts.Password != "" {
		req.SetBasicAuth(opts.UserName, opts.Password)
	}
	c := ClientWithHeaderReset(api.c, headers)
	api.mu.RUnlock()
	resp, err = c.Do(req)
	api.mu.RLock()
	if err != nil {
		return nil, err
	}
	if !opts.IgnoreStatus {
		if resp.StatusCode < 200 || resp.StatusCode > 299 {
			return resp, api.errorHandler(resp)
		}
	}
	if opts.NoResponse {
		return resp, resp.Body.Close()
	}
	return resp, nil
}

// CallJSON runs Call and decodes the body as a JSON object into response (if not nil)
//
// If request is not nil then it will be JSON encoded as the body of the request
//
// It will return resp if at all possible, even if err is set
func (api *Client) CallJSON(opts *Opts, request interface{}, response interface{}) (resp *http.Response, err error) {
	var requestBody []byte
	// Marshal the request if given
	if request != nil {
		opts = opts.Copy()
		requestBody, err = json.Marshal(request)
		opts.ContentType = "application/json"
		if err != nil {
			return nil, err
		}
		// Set the body up as a JSON object if required
		if opts.Body == nil {
			opts.Body = bytes.NewBuffer(requestBody)
		}
	}
	errChan := make(chan error, 1)
	isMultipart := opts.MultipartMetadataName != "" && opts.Body != nil && request != nil
	if isMultipart {
		bodyReader, bodyWriter := io.Pipe()
		writer := multipart.NewWriter(bodyWriter)
		opts.ContentType = writer.FormDataContentType()
		in := opts.Body
		opts.Body = bodyReader
		go func() {
			defer func() { _ = bodyWriter.Close() }()
			var err error

			// Create the first part
			err = writer.WriteField(opts.MultipartMetadataName, string(requestBody))
			if err != nil {
				errChan <- err
				return
			}

			// Add the file part
			part, err := writer.CreateFormFile(opts.MultipartContentName, opts.MultipartFileName)
			if err != nil {
				errChan <- err
				return
			}

			// Copy it in
			if _, err := io.Copy(part, in); err != nil {
				errChan <- err
				return
			}
			errChan <- writer.Close()
		}()
	}
	resp, err = api.Call(opts)
	if err != nil {
		return resp, err
	}
	if response == nil || opts.NoResponse {
		return resp, nil
	}
	if isMultipart {
		err = <-errChan
		if err != nil {
			return resp, err
		}
	}
	err = DecodeJSON(resp, response)
	return resp, err
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces

import (
	"strings"

	"github.com/juju/errors"
	"github.com/juju/juju/api/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/worker"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"github.com/juju/utils/set"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.discoverspaces")

type discoverspacesWorker struct {
	api      *discoverspaces.API
	tomb     tomb.Tomb
	observer *worker.EnvironObserver
}

// NewWorker returns a worker
func NewWorker(api *discoverspaces.API) worker.Worker {
	dw := &discoverspacesWorker{
		api: api,
	}
	go func() {
		defer dw.tomb.Done()
		dw.tomb.Kill(dw.loop())
	}()
	return dw
}

func (dw *discoverspacesWorker) Kill() {
	dw.tomb.Kill(nil)
}

func (dw *discoverspacesWorker) Wait() error {
	return dw.tomb.Wait()
}

func (dw *discoverspacesWorker) loop() (err error) {
	dw.observer, err = worker.NewEnvironObserver(dw.api)
	if err != nil {
		return err
	}
	defer func() {
		obsErr := worker.Stop(dw.observer)
		if err == nil {
			err = obsErr
		}
	}()
	environ := dw.observer.Environ()
	networkingEnviron, ok := environs.SupportsNetworking(environ)

	if ok {
		// TODO: (mfoord) API should be switched off until this is
		// completed.
		err = dw.handleSubnets(networkingEnviron)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// TODO(mfoord): we'll have a watcher here checking if we need to
	// update the spaces/subnets definition.
	dying := dw.tomb.Dying()
	for {
		select {
		case <-dying:
			return nil
		}
	}
	return nil
}

func (dw *discoverspacesWorker) handleSubnets(env environs.NetworkingEnviron) error {
	ok, err := env.SupportsSpaceDiscovery()
	if err != nil {
		return errors.Trace(err)
	}
	if !ok {
		// Nothing to do.
		return nil
	}
	providerSpaces, err := env.Spaces()
	if err != nil {
		return errors.Trace(err)
	}
	listSpacesResult, err := dw.api.ListSpaces()
	if err != nil {
		return errors.Trace(err)
	}

	stateSubnets, err := dw.api.ListSubnets(params.SubnetsFilters{})
	if err != nil {
		return errors.Trace(err)
	}
	stateSubnetIds := make(set.Strings)
	for _, subnet := range stateSubnets.Results {
		stateSubnetIds.Add(subnet.ProviderId)
	}
	stateSpaceMap := make(map[string]params.ProviderSpace)
	spaceNames := make(set.Strings)
	for _, space := range listSpacesResult.Results {
		stateSpaceMap[space.ProviderId] = space
		spaceNames.Add(space.Name)
	}

	// TODO(mfoord): we need to delete spaces and subnets that no longer
	// exist, so long as they're not in use.
	for _, space := range providerSpaces {
		spaceName := string(space.ProviderId)
		spaceName = strings.Replace(spaceName, " ", "-", -1)
		spaceName = strings.ToLower(spaceName)
		if !names.IsValidSpace(spaceName) {
			// XXX generate a valid name here
			logger.Errorf("invalid space name %v", spaceName)
			return errors.Errorf("invalid space name: %q", spaceName)
		}
		spaceTag := names.NewSpaceTag(spaceName)
		_, ok := stateSpaceMap[string(space.ProviderId)]
		if !ok {
			// XXX skip spaces with no subnets(?)
			// We need to create the space.
			args := params.CreateSpacesParams{
				Spaces: []params.CreateSpaceParams{{
					Public:   false,
					SpaceTag: spaceTag.String(),
				}}}
			// XXX check the error result too.
			_, err = dw.api.CreateSpaces(args)
			if err != nil {
				logger.Errorf("invalid creating space %v", err)
				return errors.Trace(err)
			}
		}
		// TODO(mfoord): currently no way of removing subnets, or
		// changing the space they're in, so we can only add ones we
		// don't already know about.
		logger.Errorf("Created space %v with %v subnets", spaceName, len(space.Subnets))
		for _, subnet := range space.Subnets {
			if stateSubnetIds.Contains(string(subnet.ProviderId)) {
				continue
			}
			zones := subnet.AvailabilityZones
			if len(zones) == 0 {
				zones = []string{"default"}
			}
			args := params.AddSubnetsParams{
				Subnets: []params.AddSubnetParams{{
					SubnetProviderId: string(subnet.ProviderId),
					SpaceTag:         spaceTag.String(),
					Zones:            zones,
				}}}
			// XXX check the error result too.
			logger.Tracef("Adding subnet %v", subnet.CIDR)
			result, err := dw.api.AddSubnets(args)
			if err != nil {
				logger.Errorf("invalid creating subnet %v", err)
				return errors.Trace(err)
			}

			// XXX needs doing properly (check len(result.Results
			// == 1).
			if result.Results[0].Error != nil {
				logger.Errorf("error creating subnet %v", result.Results[0].Error)
				return errors.Errorf("error creating subnet %v", result.Results[0].Error)
			}
		}
	}
	return nil
}
<commit_msg>Change log level<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces

import (
	"strings"

	"github.com/juju/errors"
	"github.com/juju/juju/api/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/worker"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"github.com/juju/utils/set"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.discoverspaces")

type discoverspacesWorker struct {
	api      *discoverspaces.API
	tomb     tomb.Tomb
	observer *worker.EnvironObserver
}

// NewWorker returns a worker
func NewWorker(api *discoverspaces.API) worker.Worker {
	dw := &discoverspacesWorker{
		api: api,
	}
	go func() {
		defer dw.tomb.Done()
		dw.tomb.Kill(dw.loop())
	}()
	return dw
}

func (dw *discoverspacesWorker) Kill() {
	dw.tomb.Kill(nil)
}

func (dw *discoverspacesWorker) Wait() error {
	return dw.tomb.Wait()
}

func (dw *discoverspacesWorker) loop() (err error) {
	dw.observer, err = worker.NewEnvironObserver(dw.api)
	if err != nil {
		return err
	}
	defer func() {
		obsErr := worker.Stop(dw.observer)
		if err == nil {
			err = obsErr
		}
	}()
	environ := dw.observer.Environ()
	networkingEnviron, ok := environs.SupportsNetworking(environ)

	if ok {
		// TODO: (mfoord) API should be switched off until this is
		// completed.
		err = dw.handleSubnets(networkingEnviron)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// TODO(mfoord): we'll have a watcher here checking if we need to
	// update the spaces/subnets definition.
	dying := dw.tomb.Dying()
	for {
		select {
		case <-dying:
			return nil
		}
	}
	return nil
}

func (dw *discoverspacesWorker) handleSubnets(env environs.NetworkingEnviron) error {
	ok, err := env.SupportsSpaceDiscovery()
	if err != nil {
		return errors.Trace(err)
	}
	if !ok {
		// Nothing to do.
		return nil
	}
	providerSpaces, err := env.Spaces()
	if err != nil {
		return errors.Trace(err)
	}
	listSpacesResult, err := dw.api.ListSpaces()
	if err != nil {
		return errors.Trace(err)
	}

	stateSubnets, err := dw.api.ListSubnets(params.SubnetsFilters{})
	if err != nil {
		return errors.Trace(err)
	}
	stateSubnetIds := make(set.Strings)
	for _, subnet := range stateSubnets.Results {
		stateSubnetIds.Add(subnet.ProviderId)
	}
	stateSpaceMap := make(map[string]params.ProviderSpace)
	spaceNames := make(set.Strings)
	for _, space := range listSpacesResult.Results {
		stateSpaceMap[space.ProviderId] = space
		spaceNames.Add(space.Name)
	}

	// TODO(mfoord): we need to delete spaces and subnets that no longer
	// exist, so long as they're not in use.
	for _, space := range providerSpaces {
		spaceName := string(space.ProviderId)
		spaceName = strings.Replace(spaceName, " ", "-", -1)
		spaceName = strings.ToLower(spaceName)
		if !names.IsValidSpace(spaceName) {
			// XXX generate a valid name here
			logger.Errorf("invalid space name %v", spaceName)
			return errors.Errorf("invalid space name: %q", spaceName)
		}
		spaceTag := names.NewSpaceTag(spaceName)
		_, ok := stateSpaceMap[string(space.ProviderId)]
		if !ok {
			// XXX skip spaces with no subnets(?)
			// We need to create the space.
			args := params.CreateSpacesParams{
				Spaces: []params.CreateSpaceParams{{
					Public:   false,
					SpaceTag: spaceTag.String(),
				}}}
			// XXX check the error result too.
			_, err = dw.api.CreateSpaces(args)
			if err != nil {
				logger.Errorf("invalid creating space %v", err)
				return errors.Trace(err)
			}
		}
		// TODO(mfoord): currently no way of removing subnets, or
		// changing the space they're in, so we can only add ones we
		// don't already know about.
		logger.Debugf("Created space %v with %v subnets", spaceName, len(space.Subnets))
		for _, subnet := range space.Subnets {
			if stateSubnetIds.Contains(string(subnet.ProviderId)) {
				continue
			}
			zones := subnet.AvailabilityZones
			if len(zones) == 0 {
				zones = []string{"default"}
			}
			args := params.AddSubnetsParams{
				Subnets: []params.AddSubnetParams{{
					SubnetProviderId: string(subnet.ProviderId),
					SpaceTag:         spaceTag.String(),
					Zones:            zones,
				}}}
			// XXX check the error result too.
			logger.Tracef("Adding subnet %v", subnet.CIDR)
			result, err := dw.api.AddSubnets(args)
			if err != nil {
				logger.Errorf("invalid creating subnet %v", err)
				return errors.Trace(err)
			}

			// XXX needs doing properly (check len(result.Results
			// == 1).
			if result.Results[0].Error != nil {
				logger.Errorf("error creating subnet %v", result.Results[0].Error)
				return errors.Errorf("error creating subnet %v", result.Results[0].Error)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package integration_test

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"code.cloudfoundry.org/cf-networking-helpers/db"
	"code.cloudfoundry.org/cf-networking-helpers/testsupport"
	"code.cloudfoundry.org/cf-networking-helpers/testsupport/metrics"
	"code.cloudfoundry.org/cf-networking-helpers/testsupport/ports"
	locketconfig "code.cloudfoundry.org/locket/cmd/locket/config"
	"code.cloudfoundry.org/locket/cmd/locket/testrunner"
	lockettestrunner "code.cloudfoundry.org/locket/cmd/locket/testrunner"
	"code.cloudfoundry.org/policy-server/api"
	"code.cloudfoundry.org/policy-server/cc_client"
	"code.cloudfoundry.org/policy-server/config"
	"code.cloudfoundry.org/policy-server/integration/helpers"
	testhelpers "code.cloudfoundry.org/test-helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

var _ = Describe("Internal API Listing security groups", func() {
	var (
		sessions                  []*gexec.Session
		asgSyncerSession          *gexec.Session
		asgSyncerConfig           config.ASGSyncerConfig
		tlsConfig                 *tls.Config
		policyServerConfs         []config.Config
		policyServerInternalConfs []config.InternalConfig
		internalConf              config.InternalConfig
		dbConf                    db.Config
		locketDBConf              db.Config
		locketProcess             ifrit.Process

		fakeMetron   metrics.FakeMetron
		mockCCServer *helpers.ConfigurableMockCCServer
	)

	BeforeEach(func() {
		fakeMetron = metrics.NewFakeMetron()

		dbConf = testsupport.GetDBConfig()
		dbConf.DatabaseName = fmt.Sprintf("internal_api_security_groups_index_test_node_%d", ports.PickAPort())

		mockCCServer = helpers.NewConfigurableMockCCServer()
		mockCCServer.Start()

		locketPort := ports.PickAPort()
		locketAddress := fmt.Sprintf("127.0.0.1:%d", locketPort)
		locketBinaryPath, err := gexec.Build("code.cloudfoundry.org/locket/cmd/locket", "-race")
		Expect(err).NotTo(HaveOccurred())
		locketDBConf = testsupport.GetDBConfig()
		locketDBConf.DatabaseName = fmt.Sprintf("internal_api_security_groups_index_test_locket_%d", locketPort)
		locketDBConnectionString, err := locketDBConf.ConnectionString()
		Expect(err).NotTo(HaveOccurred())
		testhelpers.CreateDatabase(locketDBConf)

		locketRunner := lockettestrunner.NewLocketRunner(locketBinaryPath, func(cfg *locketconfig.LocketConfig) {
			cfg.ListenAddress = locketAddress
			cfg.DatabaseDriver = dbConf.Type
			cfg.DatabaseConnectionString = locketDBConnectionString
		})

		locketProcess = ifrit.Invoke(locketRunner)
		Eventually(locketProcess.Ready()).Should(BeClosed())

		var defaultExternalConf config.Config
		var defaultInternalConf config.InternalConfig
		defaultExternalConf, defaultInternalConf, asgSyncerConfig = helpers.DefaultTestConfigWithCCServer(dbConf, fakeMetron.Address(), "fixtures", mockCCServer.URL())
		asgSyncerConfig.LocketAddress = locketAddress
		locketClientConfig := testrunner.ClientLocketConfig()
		asgSyncerConfig.LocketCACertFile = locketClientConfig.LocketCACertFile
		asgSyncerConfig.LocketClientCertFile = locketClientConfig.LocketClientCertFile
		asgSyncerConfig.LocketClientKeyFile = locketClientConfig.LocketClientKeyFile
		policyServerConfs = configurePolicyServers(defaultExternalConf, 1)
		policyServerInternalConfs = configureInternalPolicyServers(defaultInternalConf, 1)
		internalConf = policyServerInternalConfs[0]

		sessions = startPolicyAndInternalServers(policyServerConfs, policyServerInternalConfs)

		asgSyncerSession = helpers.StartAsgSyncer(policyServerAsgSyncerPath, asgSyncerConfig)

		tlsConfig = helpers.DefaultTLSConfig()
	})

	AfterEach(func() {
		asgSyncerSession.Interrupt()
		Eventually(asgSyncerSession, helpers.DEFAULT_TIMEOUT).Should(gexec.Exit())
		stopPolicyServerExternalAndInternal(sessions, policyServerConfs, policyServerInternalConfs)
		mockCCServer.Close()
		fakeMetron.Close()
		ginkgomon.Interrupt(locketProcess, 5*time.Second)
		testhelpers.RemoveDatabase(locketDBConf)
	})

	Describe("listing security groups", func() {
		BeforeEach(func() {
			mockCCServer.AddSecurityGroup(cc_client.SecurityGroupResource{
				GUID:  "sg-guid-1",
				Name:  "sg-name-1",
				Rules: []cc_client.SecurityGroupRule{{Description: "sg-rule-1", Protocol: "tcp"}},
				GloballyEnabled: cc_client.SecurityGroupGloballyEnabled{
					Staging: true,
					Running: false,
				},
				Relationships: cc_client.SecurityGroupRelationships{
					StagingSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-a"}}},
					RunningSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-b"}}},
				},
			})

			mockCCServer.AddSecurityGroup(cc_client.SecurityGroupResource{
				GUID:  "sg-guid-2",
				Name:  "sg-name-2",
				Rules: []cc_client.SecurityGroupRule{{Description: "sg-rule-2"}},
				GloballyEnabled: cc_client.SecurityGroupGloballyEnabled{
					Staging: false,
					Running: false,
				},
				Relationships: cc_client.SecurityGroupRelationships{
					StagingSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-c"}}},
					RunningSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-c"}}},
				},
			})

			mockCCServer.AddSecurityGroup(cc_client.SecurityGroupResource{
				GUID:  "sg-guid-3",
				Name:  "sg-name-3",
				Rules: []cc_client.SecurityGroupRule{{Description: "sg-rule-3", Ports: "8080"}},
				GloballyEnabled: cc_client.SecurityGroupGloballyEnabled{
					Staging: false,
					Running: false,
				},
				Relationships: cc_client.SecurityGroupRelationships{
					StagingSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-a"}}},
					RunningSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-a"}}},
				},
			})
			time.Sleep(time.Duration(asgSyncerConfig.ASGSyncInterval) * time.Second * 2)
		})

		listSecurityGroups := func(queryString string, expectedResponse api.AsgsPayload) {
			resp := helpers.MakeAndDoHTTPSRequest(
				"GET",
				fmt.Sprintf("https://%s:%d/networking/v1/internal/security_groups%s", internalConf.ListenHost, internalConf.InternalListenPort, queryString),
				nil,
				tlsConfig,
			)

			Expect(resp.StatusCode).To(Equal(http.StatusOK))
			responseString, err := ioutil.ReadAll(resp.Body)
			Expect(err).NotTo(HaveOccurred())

			var responseJson api.AsgsPayload
			err = json.Unmarshal(responseString, &responseJson)
			Expect(err).NotTo(HaveOccurred())
			Expect(responseJson.Next).To(Equal(expectedResponse.Next))
			Expect(responseJson.SecurityGroups).To(ConsistOf(expectedResponse.SecurityGroups))

			Eventually(fakeMetron.AllEvents, "5s").Should(ContainElement(
				HaveName("InternalSecurityGroupsRequestTime"),
			))

			Eventually(fakeMetron.AllEvents, "5s").Should(ContainElement(
				HaveName("SecurityGroupsStoreBySpaceGuidsSuccessTime"),
			))
		}

		sgRule1 := `{"protocol":"tcp","destination":"","ports":"","type":0,"code":0,"description":"sg-rule-1","log":false}`
		sgRule2 := `{"protocol":"","destination":"","ports":"","type":0,"code":0,"description":"sg-rule-2","log":false}`
		sgRule3 := `{"protocol":"","destination":"","ports":"8080","type":0,"code":0,"description":"sg-rule-3","log":false}`
		expectedSecurityGroup1 := api.SecurityGroup{
			Guid:              "sg-guid-1",
			Name:              "sg-name-1",
			Rules:             `[` + sgRule1 + `]`,
			StagingDefault:    true,
			RunningDefault:    false,
			StagingSpaceGuids: []string{"space-a"},
			RunningSpaceGuids: []string{"space-b"},
		}

		expectedSecurityGroup2 := api.SecurityGroup{
			Guid:              "sg-guid-2",
			Name:              "sg-name-2",
			Rules:             `[` + sgRule2 + `]`,
			StagingDefault:    false,
			RunningDefault:    false,
			StagingSpaceGuids: []string{"space-c"},
			RunningSpaceGuids: []string{"space-c"},
		}

		expectedSecurityGroup3 := api.SecurityGroup{
			Guid:              "sg-guid-3",
			Name:              "sg-name-3",
			Rules:             `[` + sgRule3 + `]`,
			StagingDefault:    false,
			RunningDefault:    false,
			StagingSpaceGuids: []string{"space-a"},
			RunningSpaceGuids: []string{"space-a"},
		}
		allResponse := api.AsgsPayload{
			Next:           0,
			SecurityGroups: []api.SecurityGroup{expectedSecurityGroup1, expectedSecurityGroup2, expectedSecurityGroup3},
		}
		globalResponse := api.AsgsPayload{
			Next:           0,
			SecurityGroups: []api.SecurityGroup{expectedSecurityGroup1},
		}
		spaceCResponse := api.AsgsPayload{
			Next:           0,
			SecurityGroups: []api.SecurityGroup{expectedSecurityGroup1, expectedSecurityGroup2},
		}

		DescribeTable("listing security groups", listSecurityGroups,
			Entry("all", "?space_guids=space-a,space-b,space-c", allResponse),
			Entry("filtered spaces", "?space_guids=space-c", spaceCResponse),
			Entry("global", "", globalResponse),
		)

		// cloud controller does not return security groups in specific order
		Describe("pagination", func() {
			It("returns all data with pagination requests", func() {
				var totalResponseSecurityGroups []api.SecurityGroup

				resp := helpers.MakeAndDoHTTPSRequest(
					"GET",
					fmt.Sprintf("https://%s:%d/networking/v1/internal/security_groups?space_guids=space-a,space-b,space-c&limit=2", internalConf.ListenHost, internalConf.InternalListenPort),
					nil,
					tlsConfig,
				)

				Expect(resp.StatusCode).To(Equal(http.StatusOK))
				responseString, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())

				var firstResponseJson api.AsgsPayload
				err = json.Unmarshal(responseString, &firstResponseJson)
				Expect(err).NotTo(HaveOccurred())

				Expect(firstResponseJson.Next).To(Equal(3))
				Expect(firstResponseJson.SecurityGroups).To(HaveLen(2))

				totalResponseSecurityGroups = append(totalResponseSecurityGroups, firstResponseJson.SecurityGroups...)

				resp = helpers.MakeAndDoHTTPSRequest(
					"GET",
					fmt.Sprintf("https://%s:%d/networking/v1/internal/security_groups?space_guids=space-a,space-b,space-c&limit=2&from=3", internalConf.ListenHost, internalConf.InternalListenPort),
					nil,
					tlsConfig,
				)

				Expect(resp.StatusCode).To(Equal(http.StatusOK))
				responseString, err = ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())

				var secondResponseJson api.AsgsPayload
				err = json.Unmarshal(responseString, &secondResponseJson)
				Expect(err).NotTo(HaveOccurred())

				Expect(secondResponseJson.Next).To(Equal(0))
				Expect(secondResponseJson.SecurityGroups).To(HaveLen(1))

				Expect(firstResponseJson.SecurityGroups).NotTo(ContainElement(secondResponseJson.SecurityGroups[0]))

				totalResponseSecurityGroups = append(totalResponseSecurityGroups, secondResponseJson.SecurityGroups...)

				Expect(totalResponseSecurityGroups).To(ConsistOf(expectedSecurityGroup1, expectedSecurityGroup2, expectedSecurityGroup3))
			})
		})
	})
})
<commit_msg>Populate CC security groups before starting ASG Syncer<commit_after>package integration_test

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"code.cloudfoundry.org/cf-networking-helpers/db"
	"code.cloudfoundry.org/cf-networking-helpers/testsupport"
	"code.cloudfoundry.org/cf-networking-helpers/testsupport/metrics"
	"code.cloudfoundry.org/cf-networking-helpers/testsupport/ports"
	locketconfig "code.cloudfoundry.org/locket/cmd/locket/config"
	"code.cloudfoundry.org/locket/cmd/locket/testrunner"
	lockettestrunner "code.cloudfoundry.org/locket/cmd/locket/testrunner"
	"code.cloudfoundry.org/policy-server/api"
	"code.cloudfoundry.org/policy-server/cc_client"
	"code.cloudfoundry.org/policy-server/config"
	"code.cloudfoundry.org/policy-server/integration/helpers"
	testhelpers "code.cloudfoundry.org/test-helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

var _ = Describe("Internal API Listing security groups", func() {
	var (
		sessions                  []*gexec.Session
		asgSyncerSession          *gexec.Session
		asgSyncerConfig           config.ASGSyncerConfig
		tlsConfig                 *tls.Config
		policyServerConfs         []config.Config
		policyServerInternalConfs []config.InternalConfig
		internalConf              config.InternalConfig
		dbConf                    db.Config
		locketDBConf              db.Config
		locketProcess             ifrit.Process

		fakeMetron   metrics.FakeMetron
		mockCCServer *helpers.ConfigurableMockCCServer
	)

	BeforeEach(func() {
		fakeMetron = metrics.NewFakeMetron()

		dbConf = testsupport.GetDBConfig()
		dbConf.DatabaseName = fmt.Sprintf("internal_api_security_groups_index_test_node_%d", ports.PickAPort())

		mockCCServer = helpers.NewConfigurableMockCCServer()
		mockCCServer.Start()

		mockCCServer.AddSecurityGroup(cc_client.SecurityGroupResource{
			GUID:  "sg-guid-1",
			Name:  "sg-name-1",
			Rules: []cc_client.SecurityGroupRule{{Description: "sg-rule-1", Protocol: "tcp"}},
			GloballyEnabled: cc_client.SecurityGroupGloballyEnabled{
				Staging: true,
				Running: false,
			},
			Relationships: cc_client.SecurityGroupRelationships{
				StagingSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-a"}}},
				RunningSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-b"}}},
			},
		})

		mockCCServer.AddSecurityGroup(cc_client.SecurityGroupResource{
			GUID:  "sg-guid-2",
			Name:  "sg-name-2",
			Rules: []cc_client.SecurityGroupRule{{Description: "sg-rule-2"}},
			GloballyEnabled: cc_client.SecurityGroupGloballyEnabled{
				Staging: false,
				Running: false,
			},
			Relationships: cc_client.SecurityGroupRelationships{
				StagingSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-c"}}},
				RunningSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-c"}}},
			},
		})

		mockCCServer.AddSecurityGroup(cc_client.SecurityGroupResource{
			GUID:  "sg-guid-3",
			Name:  "sg-name-3",
			Rules: []cc_client.SecurityGroupRule{{Description: "sg-rule-3", Ports: "8080"}},
			GloballyEnabled: cc_client.SecurityGroupGloballyEnabled{
				Staging: false,
				Running: false,
			},
			Relationships: cc_client.SecurityGroupRelationships{
				StagingSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-a"}}},
				RunningSpaces: cc_client.SecurityGroupSpaceRelationship{Data: []map[string]string{{"guid": "space-a"}}},
			},
		})

		locketPort := ports.PickAPort()
		locketAddress := fmt.Sprintf("127.0.0.1:%d", locketPort)
		locketBinaryPath, err := gexec.Build("code.cloudfoundry.org/locket/cmd/locket", "-race")
		Expect(err).NotTo(HaveOccurred())
		locketDBConf = testsupport.GetDBConfig()
		locketDBConf.DatabaseName = fmt.Sprintf("internal_api_security_groups_index_test_locket_%d", locketPort)
		locketDBConnectionString, err := locketDBConf.ConnectionString()
		Expect(err).NotTo(HaveOccurred())
		testhelpers.CreateDatabase(locketDBConf)

		locketRunner := lockettestrunner.NewLocketRunner(locketBinaryPath, func(cfg *locketconfig.LocketConfig) {
			cfg.ListenAddress = locketAddress
			cfg.DatabaseDriver = dbConf.Type
			cfg.DatabaseConnectionString = locketDBConnectionString
		})

		locketProcess = ifrit.Invoke(locketRunner)
		Eventually(locketProcess.Ready()).Should(BeClosed())

		var defaultExternalConf config.Config
		var defaultInternalConf config.InternalConfig
		defaultExternalConf, defaultInternalConf, asgSyncerConfig = helpers.DefaultTestConfigWithCCServer(dbConf, fakeMetron.Address(), "fixtures", mockCCServer.URL())
		asgSyncerConfig.LocketAddress = locketAddress
		locketClientConfig := testrunner.ClientLocketConfig()
		asgSyncerConfig.LocketCACertFile = locketClientConfig.LocketCACertFile
		asgSyncerConfig.LocketClientCertFile = locketClientConfig.LocketClientCertFile
		asgSyncerConfig.LocketClientKeyFile = locketClientConfig.LocketClientKeyFile
		policyServerConfs = configurePolicyServers(defaultExternalConf, 1)
		policyServerInternalConfs = configureInternalPolicyServers(defaultInternalConf, 1)
		internalConf = policyServerInternalConfs[0]

		sessions = startPolicyAndInternalServers(policyServerConfs, policyServerInternalConfs)

		asgSyncerSession = helpers.StartAsgSyncer(policyServerAsgSyncerPath, asgSyncerConfig)

		tlsConfig = helpers.DefaultTLSConfig()
	})

	AfterEach(func() {
		asgSyncerSession.Interrupt()
		Eventually(asgSyncerSession, helpers.DEFAULT_TIMEOUT).Should(gexec.Exit())
		stopPolicyServerExternalAndInternal(sessions, policyServerConfs, policyServerInternalConfs)
		mockCCServer.Close()
		fakeMetron.Close()
		ginkgomon.Interrupt(locketProcess, 5*time.Second)
		testhelpers.RemoveDatabase(locketDBConf)
	})

	Describe("listing security groups", func() {
		BeforeEach(func() {
			// wait for ASG Syncer to get security groups from cloud controller
			time.Sleep(time.Duration(asgSyncerConfig.ASGSyncInterval) * time.Second * 2)
		})

		listSecurityGroups := func(queryString string, expectedResponse api.AsgsPayload) {
			resp := helpers.MakeAndDoHTTPSRequest(
				"GET",
				fmt.Sprintf("https://%s:%d/networking/v1/internal/security_groups%s", internalConf.ListenHost, internalConf.InternalListenPort, queryString),
				nil,
				tlsConfig,
			)

			Expect(resp.StatusCode).To(Equal(http.StatusOK))
			responseString, err := ioutil.ReadAll(resp.Body)
			Expect(err).NotTo(HaveOccurred())

			var responseJson api.AsgsPayload
			err = json.Unmarshal(responseString, &responseJson)
			Expect(err).NotTo(HaveOccurred())
			Expect(responseJson.Next).To(Equal(expectedResponse.Next))
			Expect(responseJson.SecurityGroups).To(ConsistOf(expectedResponse.SecurityGroups))

			Eventually(fakeMetron.AllEvents, "5s").Should(ContainElement(
				HaveName("InternalSecurityGroupsRequestTime"),
			))

			Eventually(fakeMetron.AllEvents, "5s").Should(ContainElement(
				HaveName("SecurityGroupsStoreBySpaceGuidsSuccessTime"),
			))
		}

		sgRule1 := `{"protocol":"tcp","destination":"","ports":"","type":0,"code":0,"description":"sg-rule-1","log":false}`
		sgRule2 := `{"protocol":"","destination":"","ports":"","type":0,"code":0,"description":"sg-rule-2","log":false}`
		sgRule3 := `{"protocol":"","destination":"","ports":"8080","type":0,"code":0,"description":"sg-rule-3","log":false}`
		expectedSecurityGroup1 := api.SecurityGroup{
			Guid:              "sg-guid-1",
			Name:              "sg-name-1",
			Rules:             `[` + sgRule1 + `]`,
			StagingDefault:    true,
			RunningDefault:    false,
			StagingSpaceGuids: []string{"space-a"},
			RunningSpaceGuids: []string{"space-b"},
		}

		expectedSecurityGroup2 := api.SecurityGroup{
			Guid:              "sg-guid-2",
			Name:              "sg-name-2",
			Rules:             `[` + sgRule2 + `]`,
			StagingDefault:    false,
			RunningDefault:    false,
			StagingSpaceGuids: []string{"space-c"},
			RunningSpaceGuids: []string{"space-c"},
		}

		expectedSecurityGroup3 := api.SecurityGroup{
			Guid:              "sg-guid-3",
			Name:              "sg-name-3",
			Rules:             `[` + sgRule3 + `]`,
			StagingDefault:    false,
			RunningDefault:    false,
			StagingSpaceGuids: []string{"space-a"},
			RunningSpaceGuids: []string{"space-a"},
		}
		allResponse := api.AsgsPayload{
			Next:           0,
			SecurityGroups: []api.SecurityGroup{expectedSecurityGroup1, expectedSecurityGroup2, expectedSecurityGroup3},
		}
		globalResponse := api.AsgsPayload{
			Next:           0,
			SecurityGroups: []api.SecurityGroup{expectedSecurityGroup1},
		}
		spaceCResponse := api.AsgsPayload{
			Next:           0,
			SecurityGroups: []api.SecurityGroup{expectedSecurityGroup1, expectedSecurityGroup2},
		}

		DescribeTable("listing security groups", listSecurityGroups,
			Entry("all", "?space_guids=space-a,space-b,space-c", allResponse),
			Entry("filtered spaces", "?space_guids=space-c", spaceCResponse),
			Entry("global", "", globalResponse),
		)

		// cloud controller does not return security groups in specific order
		Describe("pagination", func() {
			It("returns all data with pagination requests", func() {
				var totalResponseSecurityGroups []api.SecurityGroup

				resp := helpers.MakeAndDoHTTPSRequest(
					"GET",
					fmt.Sprintf("https://%s:%d/networking/v1/internal/security_groups?space_guids=space-a,space-b,space-c&limit=2", internalConf.ListenHost, internalConf.InternalListenPort),
					nil,
					tlsConfig,
				)

				Expect(resp.StatusCode).To(Equal(http.StatusOK))
				responseString, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())

				var firstResponseJson api.AsgsPayload
				err = json.Unmarshal(responseString, &firstResponseJson)
				Expect(err).NotTo(HaveOccurred())

				Expect(firstResponseJson.Next).To(Equal(3))
				Expect(firstResponseJson.SecurityGroups).To(HaveLen(2))

				totalResponseSecurityGroups = append(totalResponseSecurityGroups, firstResponseJson.SecurityGroups...)

				resp = helpers.MakeAndDoHTTPSRequest(
					"GET",
					fmt.Sprintf("https://%s:%d/networking/v1/internal/security_groups?space_guids=space-a,space-b,space-c&limit=2&from=3", internalConf.ListenHost, internalConf.InternalListenPort),
					nil,
					tlsConfig,
				)

				Expect(resp.StatusCode).To(Equal(http.StatusOK))
				responseString, err = ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())

				var secondResponseJson api.AsgsPayload
				err = json.Unmarshal(responseString, &secondResponseJson)
				Expect(err).NotTo(HaveOccurred())

				Expect(secondResponseJson.Next).To(Equal(0))
				Expect(secondResponseJson.SecurityGroups).To(HaveLen(1))

				Expect(firstResponseJson.SecurityGroups).NotTo(ContainElement(secondResponseJson.SecurityGroups[0]))

				totalResponseSecurityGroups = append(totalResponseSecurityGroups, secondResponseJson.SecurityGroups...)

				Expect(totalResponseSecurityGroups).To(ConsistOf(expectedSecurityGroup1, expectedSecurityGroup2, expectedSecurityGroup3))
			})
		})
	})
})
<|endoftext|>


<commit_before>package reviewdog

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"

	"github.com/reviewdog/reviewdog/diff"
)

var X = 14

// Reviewdog represents review dog application which parses result of compiler
// or linter, get diff and filter the results by diff, and report filtered
// results.
type Reviewdog struct {
	toolname string
	p        Parser
	c        CommentService
	d        DiffService
}

// NewReviewdog returns a new Reviewdog.
func NewReviewdog(toolname string, p Parser, c CommentService, d DiffService) *Reviewdog {
	return &Reviewdog{p: p, c: c, d: d, toolname: toolname}
}

func RunFromResult(ctx context.Context, c CommentService, results []*CheckResult,
	filediffs []*diff.FileDiff, strip int, toolname string) error {
	return (&Reviewdog{c: c, toolname: toolname}).runFromResult(ctx, results, filediffs, strip)
}

// CheckResult represents a checked result of static analysis tools.
// :h error-file-format
type CheckResult struct {
	Path    string   // relative file path
	Lnum    int      // line number
	Col     int      // column number (1 <tab> == 1 character column)
	Message string   // error message
	Lines   []string // Original error lines (often one line)
}

// Parser is an interface which parses compilers, linters, or any tools
// results.
type Parser interface {
	Parse(r io.Reader) ([]*CheckResult, error)
}

// Comment represents a reported result as a comment.
type Comment struct {
	*CheckResult
	Body     string
	LnumDiff int
	ToolName string
}

// CommentService is an interface which posts Comment.
type CommentService interface {
	Post(context.Context, *Comment) error
}

// BulkCommentService posts comments all at once when Flush() is called.
// Flush() will be called at the end of reviewdog run.
type BulkCommentService interface {
	CommentService
	Flush(context.Context) error
}

// DiffService is an interface which get diff.
type DiffService interface {
	Diff(context.Context) ([]byte, error)
	Strip() int
}

func (w *Reviewdog) runFromResult(ctx context.Context, results []*CheckResult,
	filediffs []*diff.FileDiff, strip int) error {
	wd, err := os.Getwd()
	if err != nil {
		return err
	}

	checks := FilterCheck(results, filediffs, strip, wd)
	for _, check := range checks {
		if !check.InDiff {
			continue
		}
		comment := &Comment{
			CheckResult: check.CheckResult,
			Body:        check.Message, // TODO: format message
			LnumDiff:    check.LnumDiff,
			ToolName:    w.toolname,
		}
		if err := w.c.Post(ctx, comment); err != nil {
			return err
		}
	}

	if bulk, ok := w.c.(BulkCommentService); ok {
		return bulk.Flush(ctx)
	}

	return nil
}

// Run runs Reviewdog application.
func (w *Reviewdog) Run(ctx context.Context, r io.Reader) error {
	results, err := w.p.Parse(r)
	if err != nil {
		return fmt.Errorf("parse error: %v", err)
	}

	d, err := w.d.Diff(ctx)
	if err != nil {
		return fmt.Errorf("fail to get diff: %v", err)
	}

	filediffs, err := diff.ParseMultiFile(bytes.NewReader(d))
	if err != nil {
		return fmt.Errorf("fail to parse diff: %v", err)
	}

	return w.runFromResult(ctx, results, filediffs, w.d.Strip())
}
<commit_msg>Revert "debug"<commit_after>package reviewdog

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"

	"github.com/reviewdog/reviewdog/diff"
)

// Reviewdog represents review dog application which parses result of compiler
// or linter, get diff and filter the results by diff, and report filtered
// results.
type Reviewdog struct {
	toolname string
	p        Parser
	c        CommentService
	d        DiffService
}

// NewReviewdog returns a new Reviewdog.
func NewReviewdog(toolname string, p Parser, c CommentService, d DiffService) *Reviewdog {
	return &Reviewdog{p: p, c: c, d: d, toolname: toolname}
}

func RunFromResult(ctx context.Context, c CommentService, results []*CheckResult,
	filediffs []*diff.FileDiff, strip int, toolname string) error {
	return (&Reviewdog{c: c, toolname: toolname}).runFromResult(ctx, results, filediffs, strip)
}

// CheckResult represents a checked result of static analysis tools.
// :h error-file-format
type CheckResult struct {
	Path    string   // relative file path
	Lnum    int      // line number
	Col     int      // column number (1 <tab> == 1 character column)
	Message string   // error message
	Lines   []string // Original error lines (often one line)
}

// Parser is an interface which parses compilers, linters, or any tools
// results.
type Parser interface {
	Parse(r io.Reader) ([]*CheckResult, error)
}

// Comment represents a reported result as a comment.
type Comment struct {
	*CheckResult
	Body     string
	LnumDiff int
	ToolName string
}

// CommentService is an interface which posts Comment.
type CommentService interface {
	Post(context.Context, *Comment) error
}

// BulkCommentService posts comments all at once when Flush() is called.
// Flush() will be called at the end of reviewdog run.
type BulkCommentService interface {
	CommentService
	Flush(context.Context) error
}

// DiffService is an interface which get diff.
type DiffService interface {
	Diff(context.Context) ([]byte, error)
	Strip() int
}

func (w *Reviewdog) runFromResult(ctx context.Context, results []*CheckResult,
	filediffs []*diff.FileDiff, strip int) error {
	wd, err := os.Getwd()
	if err != nil {
		return err
	}

	checks := FilterCheck(results, filediffs, strip, wd)
	for _, check := range checks {
		if !check.InDiff {
			continue
		}
		comment := &Comment{
			CheckResult: check.CheckResult,
			Body:        check.Message, // TODO: format message
			LnumDiff:    check.LnumDiff,
			ToolName:    w.toolname,
		}
		if err := w.c.Post(ctx, comment); err != nil {
			return err
		}
	}

	if bulk, ok := w.c.(BulkCommentService); ok {
		return bulk.Flush(ctx)
	}

	return nil
}

// Run runs Reviewdog application.
func (w *Reviewdog) Run(ctx context.Context, r io.Reader) error {
	results, err := w.p.Parse(r)
	if err != nil {
		return fmt.Errorf("parse error: %v", err)
	}

	d, err := w.d.Diff(ctx)
	if err != nil {
		return fmt.Errorf("fail to get diff: %v", err)
	}

	filediffs, err := diff.ParseMultiFile(bytes.NewReader(d))
	if err != nil {
		return fmt.Errorf("fail to parse diff: %v", err)
	}

	return w.runFromResult(ctx, results, filediffs, w.d.Strip())
}
<|endoftext|>


<commit_before>package resource

import (
	"fmt"

	"github.com/dnaeon/gru/graph"
)

// Collection type is a map which keys are the
// resource ids and their values are the actual resources
type Collection map[string]Resource

// CreateCollection creates a map from
func CreateCollection(resources []Resource) (Collection, error) {
	c := make(Collection)

	for _, r := range resources {
		id := r.ID()
		if _, ok := c[id]; ok {
			return c, fmt.Errorf("Duplicate resource declaration for %s", id)
		}
		c[id] = r
	}

	return c, nil
}

// DependencyGraph builds a dependency graph for the collection
func (c Collection) DependencyGraph() (*graph.Graph, error) {
	g := graph.New()

	// A map containing the resource ids and their nodes in the graph
	nodes := make(map[string]*graph.Node)
	for id := range c {
		node := graph.NewNode(id)
		nodes[id] = node
		g.AddNode(node)
	}

	// Connect the nodes in the graph
	for id, r := range c {
		for _, dep := range r.Dependencies() {
			if _, ok := c[dep]; !ok {
				return g, fmt.Errorf("%s wants %s, which does not exist", id, dep)
			}
			g.AddEdge(nodes[id], nodes[dep])
		}
	}

	return g, nil
}

// ReversedGraph builds a reverse dependency graph for the collection
func (c Collection) ReversedGraph() (*graph.Graph, error) {
	g := graph.New()

	// A map containing the resource ids and their nodes in the graph
	nodes := make(map[string]*graph.Node)
	for id := range c {
		node := graph.NewNode(id)
		nodes[id] = node
		g.AddNode(node)
	}

	// Connect the nodes in the graph
	for id, r := range c {
		for _, dep := range r.Dependencies() {
			if _, ok := c[dep]; !ok {
				return g, fmt.Errorf("%s wants %s, which does not exist", id, dep)
			}
			g.AddEdge(nodes[dep], nodes[id])
		}
	}

	return g, nil
}
<commit_msg>resource: remove Collection.ReversedGraph() method<commit_after>package resource

import (
	"fmt"

	"github.com/dnaeon/gru/graph"
)

// Collection type is a map which keys are the
// resource ids and their values are the actual resources
type Collection map[string]Resource

// CreateCollection creates a map from
func CreateCollection(resources []Resource) (Collection, error) {
	c := make(Collection)

	for _, r := range resources {
		id := r.ID()
		if _, ok := c[id]; ok {
			return c, fmt.Errorf("Duplicate resource declaration for %s", id)
		}
		c[id] = r
	}

	return c, nil
}

// DependencyGraph builds a dependency graph for the collection
func (c Collection) DependencyGraph() (*graph.Graph, error) {
	g := graph.New()

	// A map containing the resource ids and their nodes in the graph
	nodes := make(map[string]*graph.Node)
	for id := range c {
		node := graph.NewNode(id)
		nodes[id] = node
		g.AddNode(node)
	}

	// Connect the nodes in the graph
	for id, r := range c {
		for _, dep := range r.Dependencies() {
			if _, ok := c[dep]; !ok {
				return g, fmt.Errorf("%s wants %s, which does not exist", id, dep)
			}
			g.AddEdge(nodes[id], nodes[dep])
		}
	}

	return g, nil
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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
//
//     https://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 main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"

	// For triggering manual rebuilds.
	"golang.org/x/oauth2/google"
	"google.golang.org/api/cloudbuild/v1"
)

var redirects = map[string]string{
	"/change":    "https://github.com/google/gvisor",
	"/issue":     "https://github.com/google/gvisor/issues",
	"/issue/new": "https://github.com/google/gvisor/issues/new",
	"/pr":        "https://github.com/google/gvisor/pulls",

	// Redirects to compatibility docs.
	"/c":             "/docs/user_guide/compatibility",
	"/c/linux/amd64": "/docs/user_guide/compatibility/amd64",

	// Deprecated, but links continue to work.
	"/cl": "https://gvisor-review.googlesource.com",
}

var prefixHelpers = map[string]string{
	"change": "https://github.com/google/gvisor/commit/%s",
	"issue":  "https://github.com/google/gvisor/issues/%s",
	"pr":     "https://github.com/google/gvisor/pull/%s",

	// Redirects to compatibility docs.
	"c/linux/amd64": "/docs/user_guide/compatibility/amd64/#%s",

	// Deprecated, but links continue to work.
	"cl": "https://gvisor-review.googlesource.com/c/gvisor/+/%s",
}

var (
	validId     = regexp.MustCompile(`^[A-Za-z0-9-]*/?$`)
	goGetHeader = `<meta name="go-import" content="gvisor.dev/gvisor git https://github.com/google/gvisor">`
	goGetHTML5  = `<!doctype html><html><head><meta charset=utf-8>` + goGetHeader + `<title>Go-get</title></head><body></html>`
)

// wrappedHandler wraps an http.Handler.
//
// If the query parameters include go-get=1, then we redirect to a single
// static page that allows us to serve arbitrary Go packages.
func wrappedHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gg, ok := r.URL.Query()["go-get"]
		if ok && len(gg) == 1 && gg[0] == "1" {
			// Serve a trivial html page.
			w.Write([]byte(goGetHTML5))
			return
		}
		// Fallthrough.
		h.ServeHTTP(w, r)
	})
}

// redirectWithQuery redirects to the given target url preserving query parameters.
func redirectWithQuery(w http.ResponseWriter, r *http.Request, target string) {
	url := target
	if qs := r.URL.RawQuery; qs != "" {
		url += "?" + qs
	}
	http.Redirect(w, r, url, http.StatusFound)
}

// hostRedirectHandler redirects the www. domain to the naked domain.
func hostRedirectHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.Host, "www.") {
			// Redirect to the naked domain.
			r.URL.Scheme = "https"  // Assume https.
			r.URL.Host = r.Host[4:] // Remove the 'www.'
			http.Redirect(w, r, r.URL.String(), http.StatusMovedPermanently)
			return
		}
		h.ServeHTTP(w, r)
	})
}

// prefixRedirectHandler returns a handler that redirects to the given formated url.
func prefixRedirectHandler(prefix, baseURL string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if p := r.URL.Path; p == prefix {
			// Redirect /prefix/ to /prefix.
			http.Redirect(w, r, p[:len(p)-1], http.StatusFound)
			return
		}
		id := r.URL.Path[len(prefix):]
		if !validId.MatchString(id) {
			http.Error(w, "Not found", http.StatusNotFound)
			return
		}
		target := fmt.Sprintf(baseURL, id)
		redirectWithQuery(w, r, target)
	})
}

// redirectHandler returns a handler that redirects to the given url.
func redirectHandler(target string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		redirectWithQuery(w, r, target)
	})
}

// redirectRedirects registers redirect http handlers.
func registerRedirects(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}

	for prefix, baseURL := range prefixHelpers {
		p := "/" + prefix + "/"
		mux.Handle(p, hostRedirectHandler(wrappedHandler(prefixRedirectHandler(p, baseURL))))
	}

	for path, redirect := range redirects {
		mux.Handle(path, hostRedirectHandler(wrappedHandler(redirectHandler(redirect))))
	}
}

// registerStatic registers static file handlers
func registerStatic(mux *http.ServeMux, staticDir string) {
	if mux == nil {
		mux = http.DefaultServeMux
	}
	mux.Handle("/", hostRedirectHandler(wrappedHandler(http.FileServer(http.Dir(staticDir)))))
}

// registerRebuild registers the rebuild handler.
func registerRebuild(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}

	mux.Handle("/rebuild", wrappedHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := context.Background()
		credentials, err := google.FindDefaultCredentials(ctx, cloudbuild.CloudPlatformScope)
		if err != nil {
			http.Error(w, "credentials error: "+err.Error(), 500)
			return
		}
		cloudbuildService, err := cloudbuild.NewService(ctx)
		if err != nil {
			http.Error(w, "cloudbuild service error: "+err.Error(), 500)
			return
		}
		projectID := credentials.ProjectID
		if projectID == "" {
			// If running locally, then this project will not be
			// available. Use the default project here.
			projectID = "gvisor-website"
		}
		triggers, err := cloudbuildService.Projects.Triggers.List(projectID).Do()
		if err != nil {
			http.Error(w, "trigger list error: "+err.Error(), 500)
			return
		}
		if len(triggers.Triggers) < 1 {
			http.Error(w, "trigger list error: no triggers", 500)
			return
		}
		if _, err := cloudbuildService.Projects.Triggers.Run(
			projectID,
			triggers.Triggers[0].Id,
			&cloudbuild.RepoSource{
				// In the current project, require that a
				// github cloud source repository exists with
				// the given name, and build from master.
				BranchName: "master",
				RepoName:   "github_google_gvisor-website",
				ProjectId:  projectID,
			}).Do(); err != nil {
			http.Error(w, "run error: "+err.Error(), 500)
			return
		}
	})))
}

func envFlagString(name, def string) string {
	if val := os.Getenv(name); val != "" {
		return val
	}
	return def
}

var (
	addr      = flag.String("http", envFlagString("HTTP", ":8080"), "HTTP service address")
	staticDir = flag.String("static-dir", envFlagString("STATIC_DIR", "static"), "static files directory")
)

func main() {
	flag.Parse()

	registerRedirects(nil)
	registerRebuild(nil)
	registerStatic(nil, *staticDir)

	log.Fatal(http.ListenAndServe(*addr, nil))
}
<commit_msg>Redirect to the custom domain from appspot.com<commit_after>// Copyright 2019 Google LLC
//
// 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
//
//     https://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 main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"

	// For triggering manual rebuilds.
	"golang.org/x/oauth2/google"
	"google.golang.org/api/cloudbuild/v1"
)

var redirects = map[string]string{
	"/change":    "https://github.com/google/gvisor",
	"/issue":     "https://github.com/google/gvisor/issues",
	"/issue/new": "https://github.com/google/gvisor/issues/new",
	"/pr":        "https://github.com/google/gvisor/pulls",

	// Redirects to compatibility docs.
	"/c":             "/docs/user_guide/compatibility",
	"/c/linux/amd64": "/docs/user_guide/compatibility/amd64",

	// Deprecated, but links continue to work.
	"/cl": "https://gvisor-review.googlesource.com",
}

var prefixHelpers = map[string]string{
	"change": "https://github.com/google/gvisor/commit/%s",
	"issue":  "https://github.com/google/gvisor/issues/%s",
	"pr":     "https://github.com/google/gvisor/pull/%s",

	// Redirects to compatibility docs.
	"c/linux/amd64": "/docs/user_guide/compatibility/amd64/#%s",

	// Deprecated, but links continue to work.
	"cl": "https://gvisor-review.googlesource.com/c/gvisor/+/%s",
}

var (
	validId     = regexp.MustCompile(`^[A-Za-z0-9-]*/?$`)
	goGetHeader = `<meta name="go-import" content="gvisor.dev/gvisor git https://github.com/google/gvisor">`
	goGetHTML5  = `<!doctype html><html><head><meta charset=utf-8>` + goGetHeader + `<title>Go-get</title></head><body></html>`
)

// wrappedHandler wraps an http.Handler.
//
// If the query parameters include go-get=1, then we redirect to a single
// static page that allows us to serve arbitrary Go packages.
func wrappedHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gg, ok := r.URL.Query()["go-get"]
		if ok && len(gg) == 1 && gg[0] == "1" {
			// Serve a trivial html page.
			w.Write([]byte(goGetHTML5))
			return
		}
		// Fallthrough.
		h.ServeHTTP(w, r)
	})
}

// redirectWithQuery redirects to the given target url preserving query parameters.
func redirectWithQuery(w http.ResponseWriter, r *http.Request, target string) {
	url := target
	if qs := r.URL.RawQuery; qs != "" {
		url += "?" + qs
	}
	http.Redirect(w, r, url, http.StatusFound)
}

// hostRedirectHandler redirects the www. domain to the naked domain.
func hostRedirectHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.Host, "www.") {
			// Redirect to the naked domain.
			r.URL.Scheme = "https"  // Assume https.
			r.URL.Host = r.Host[4:] // Remove the 'www.'
			http.Redirect(w, r, r.URL.String(), http.StatusMovedPermanently)
			return
		}

		if *projectId != "" && r.Host == *projectId+".appspot.com" && *customHost != "" {
			// Redirect to the custom domain.
			r.URL.Scheme = "https" // Assume https.
			r.URL.Host = *customHost
			http.Redirect(w, r, r.URL.String(), http.StatusMovedPermanently)
			return
		}
		h.ServeHTTP(w, r)
	})
}

// prefixRedirectHandler returns a handler that redirects to the given formated url.
func prefixRedirectHandler(prefix, baseURL string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if p := r.URL.Path; p == prefix {
			// Redirect /prefix/ to /prefix.
			http.Redirect(w, r, p[:len(p)-1], http.StatusFound)
			return
		}
		id := r.URL.Path[len(prefix):]
		if !validId.MatchString(id) {
			http.Error(w, "Not found", http.StatusNotFound)
			return
		}
		target := fmt.Sprintf(baseURL, id)
		redirectWithQuery(w, r, target)
	})
}

// redirectHandler returns a handler that redirects to the given url.
func redirectHandler(target string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		redirectWithQuery(w, r, target)
	})
}

// redirectRedirects registers redirect http handlers.
func registerRedirects(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}

	for prefix, baseURL := range prefixHelpers {
		p := "/" + prefix + "/"
		mux.Handle(p, hostRedirectHandler(wrappedHandler(prefixRedirectHandler(p, baseURL))))
	}

	for path, redirect := range redirects {
		mux.Handle(path, hostRedirectHandler(wrappedHandler(redirectHandler(redirect))))
	}
}

// registerStatic registers static file handlers
func registerStatic(mux *http.ServeMux, staticDir string) {
	if mux == nil {
		mux = http.DefaultServeMux
	}
	mux.Handle("/", hostRedirectHandler(wrappedHandler(http.FileServer(http.Dir(staticDir)))))
}

// registerRebuild registers the rebuild handler.
func registerRebuild(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}

	mux.Handle("/rebuild", wrappedHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := context.Background()
		credentials, err := google.FindDefaultCredentials(ctx, cloudbuild.CloudPlatformScope)
		if err != nil {
			http.Error(w, "credentials error: "+err.Error(), 500)
			return
		}
		cloudbuildService, err := cloudbuild.NewService(ctx)
		if err != nil {
			http.Error(w, "cloudbuild service error: "+err.Error(), 500)
			return
		}
		projectID := credentials.ProjectID
		if projectID == "" {
			// If running locally, then this project will not be
			// available. Use the default project here.
			projectID = "gvisor-website"
		}
		triggers, err := cloudbuildService.Projects.Triggers.List(projectID).Do()
		if err != nil {
			http.Error(w, "trigger list error: "+err.Error(), 500)
			return
		}
		if len(triggers.Triggers) < 1 {
			http.Error(w, "trigger list error: no triggers", 500)
			return
		}
		if _, err := cloudbuildService.Projects.Triggers.Run(
			projectID,
			triggers.Triggers[0].Id,
			&cloudbuild.RepoSource{
				// In the current project, require that a
				// github cloud source repository exists with
				// the given name, and build from master.
				BranchName: "master",
				RepoName:   "github_google_gvisor-website",
				ProjectId:  projectID,
			}).Do(); err != nil {
			http.Error(w, "run error: "+err.Error(), 500)
			return
		}
	})))
}

func envFlagString(name, def string) string {
	if val := os.Getenv(name); val != "" {
		return val
	}
	return def
}

var (
	addr      = flag.String("http", envFlagString("HTTP", ":8080"), "HTTP service address")
	staticDir = flag.String("static-dir", envFlagString("STATIC_DIR", "static"), "static files directory")
	// Uses the standard GOOGLE_CLOUD_PROJECT environment variable set by App Engine.
	projectId  = flag.String("project-id", envFlagString("GOOGLE_CLOUD_PROJECT", ""), "The App Engine project ID.")
	customHost = flag.String("custom-domain", envFlagString("CUSTOM_DOMAIN", "gvisor.dev"), "The application's custom domain.")
)

func main() {
	flag.Parse()

	registerRedirects(nil)
	registerRebuild(nil)
	registerStatic(nil, *staticDir)

	log.Fatal(http.ListenAndServe(*addr, nil))
}
<|endoftext|>