<commit_before>package msp

import (
	"container/heap"
	"crypto/rand"
	"errors"
	"strings"
)

// A UserDatabase is an abstraction over the name -> share map returned by the
// secret splitter that allows an application to only decrypt or request shares
// when needed, rather than re-build a partial map of known data.
type UserDatabase interface {
	ValidUser(name string) bool
	CanGetShare(name string) bool
	GetShare(name string) ([][]byte, error)
}

type Condition interface { // Represents one condition in a predicate
	Ok(*UserDatabase) bool
}

type Name struct { // Type of condition
	string
	index int
}

func (n Name) Ok(db *UserDatabase) bool {
	return (*db).CanGetShare(n.string)
}

type TraceElem struct {
	loc   int
	names []string
	trace []string
}

type TraceSlice []TraceElem

func (ts TraceSlice) Len() int      { return len(ts) }
func (ts TraceSlice) Swap(i, j int) { ts[i], ts[j] = ts[j], ts[i] }

func (ts TraceSlice) Less(i, j int) bool {
	return len(ts[i].trace) < len(ts[j].trace)
}

func (ts *TraceSlice) Push(te interface{}) { *ts = append(*ts, te.(TraceElem)) }
func (ts *TraceSlice) Pop() interface{} {
	out := (*ts)[0]
	*ts = (*ts)[1 : len(*ts)-1]

	return out
}

func (ts TraceSlice) Compact() (index []int, names []string, trace []string) {
	for _, te := range ts {
		index = append(index, te.loc)
		names = append(names, te.names...)
		trace = append(trace, te.trace...)
	}

	ptr, cutoff := 0, len(trace)

TopLoop:
	for ptr < cutoff {
		for i := 0; i < ptr; i++ {
			if trace[i] == trace[ptr] {
				trace[ptr], trace[cutoff-1] = trace[cutoff-1], trace[ptr]
				cutoff--

				continue TopLoop
			}
		}

		ptr++
	}
	trace = trace[0:cutoff]

	return
}

type MSP Formatted

func StringToMSP(pred string) (m MSP, err error) {
	var f Formatted

	if -1 == strings.Index(pred, ",") {
		var r Raw
		r, err = StringToRaw(pred)
		if err != nil {
			return
		}

		f = r.Formatted()
	} else {
		f, err = StringToFormatted(pred)
		if err != nil {
			return
		}
	}

	return MSP(f), nil
}

func (m MSP) DerivePath(db *UserDatabase) (ok bool, names []string, locs []int, trace []string) {
	ts := &TraceSlice{}

	for i, cond := range m.Conds {
		switch cond.(type) {
		case Name:
			if (*db).CanGetShare(cond.(Name).string) {
				heap.Push(ts, TraceElem{
					i,
					[]string{cond.(Name).string},
					[]string{cond.(Name).string},
				})
			}

		case Formatted:
			sok, _, _, strace := MSP(cond.(Formatted)).DerivePath(db)
			if !sok {
				continue
			}

			heap.Push(ts, TraceElem{i, []string{}, strace})
		}

		if (*ts).Len() > m.Min {
			*ts = (*ts)[0:m.Min]
		}
	}

	ok = (*ts).Len() >= m.Min
	locs, names, trace = ts.Compact()
	return
}

// DistributeShares takes as input a secret and a user database and returns secret shares according to access structure
// described by the MSP.
func (m MSP) DistributeShares(sec []byte, db *UserDatabase) (map[string][][]byte, error) {
	out := make(map[string][][]byte)

	// Generate a Vandermonde matrix.
	height, width := len(m.Conds), m.Min
	M := Matrix(make([]Row, height))

	for i := 0; i < height; i++ {
		M[i] = NewRow(width)

		for j := 0; j < width; j++ {
			M[i][j][0] = byte(i + 1)
			M[i][j] = M[i][j].Exp(j)
		}
	}

	// Convert secret vector.
	s := NewRow(width)
	s[0] = FieldElem(sec)

	for i := 1; i < width; i++ {
		r := NewFieldElem()
		rand.Read(r)

		s[i] = FieldElem(r)
	}

	// Calculate shares.
	shares := M.Mul(s)

	// Distribute the shares.
	for i, cond := range m.Conds {
		share := shares[i]

		switch cond.(type) {
		case Name:
			name := cond.(Name).string
			if _, ok := out[name]; ok {
				out[name] = append(out[name], share)
			} else if (*db).ValidUser(name) {
				out[name] = [][]byte{share}
			} else {
				return out, errors.New("Unknown user in predicate.")
			}

		default:
			below := MSP(cond.(Formatted))
			subOut, err := below.DistributeShares(share, db)
			if err != nil {
				return out, err
			}

			for name, shares := range subOut {
				if _, ok := out[name]; ok {
					out[name] = append(out[name], shares...)
				} else {
					out[name] = shares
				}

			}
		}
	}

	return out, nil
}

// RecoverSecret takes a user database storing secret shares as input and returns the original secret.
func (m MSP) RecoverSecret(db *UserDatabase) ([]byte, error) {
	cache := make(map[string][][]byte, 0) // Caches un-used shares for a user.
	return m.recoverSecret(db, cache)
}

func (m MSP) recoverSecret(db *UserDatabase, cache map[string][][]byte) ([]byte, error) {
	var (
		index  = []int{}       // Indexes where given shares were in the matrix.
		shares = []FieldElem{} // Contains shares that will be used in reconstruction.
	)

	ok, names, locs, _ := m.DerivePath(db)
	if !ok {
		return nil, errors.New("Not enough shares to recover.")
	}

	for _, name := range names {
		if _, cached := cache[name]; !cached {
			out, err := (*db).GetShare(name)
			if err != nil {
				return nil, err
			}

			cache[name] = out
		}
	}

	for _, loc := range locs {
		gate := m.Conds[loc]
		index = append(index, loc+1)

		switch gate.(type) {
		case Name:
			if len(cache[gate.(Name).string]) <= gate.(Name).index {
				return nil, errors.New("Predicate / database mismatch!")
			}

			shares = append(shares, FieldElem(cache[gate.(Name).string][gate.(Name).index]))

		case Formatted:
			share, err := MSP(gate.(Formatted)).recoverSecret(db, cache)
			if err != nil {
				return nil, err
			}

			shares = append(shares, FieldElem(share))
		}
	}

	// Generate the Vandermonde matrix specific to whichever users' shares we're using.
	MSub := Matrix(make([]Row, m.Min))

	for i := 0; i < m.Min; i++ {
		MSub[i] = NewRow(m.Min)

		for j := 0; j < m.Min; j++ {
			MSub[i][j][0] = byte(index[i])
			MSub[i][j] = MSub[i][j].Exp(j)
		}
	}

	// Calculate the reconstruction vector and use it to recover the secret.
	r, ok := MSub.Recovery()
	if !ok {
		return nil, errors.New("Unable to find a reconstruction vector!")
	}

	// Compute dot product of the shares vector and the reconstruction vector to
	// recover the secret.
	s := Row(shares).DotProduct(r)

	return []byte(s), nil
}
<commit_msg>Comment the DerivePath logic.<commit_after>package msp

import (
	"container/heap"
	"crypto/rand"
	"errors"
	"strings"
)

// A UserDatabase is an abstraction over the name -> share map returned by the
// secret splitter that allows an application to only decrypt or request shares
// when needed, rather than re-build a partial map of known data.
type UserDatabase interface {
	ValidUser(name string) bool
	CanGetShare(name string) bool
	GetShare(name string) ([][]byte, error)
}

type Condition interface { // Represents one condition in a predicate
	Ok(*UserDatabase) bool
}

type Name struct { // Type of condition
	string
	index int
}

func (n Name) Ok(db *UserDatabase) bool {
	return (*db).CanGetShare(n.string)
}

type TraceElem struct {
	loc   int
	names []string
	trace []string
}

type TraceSlice []TraceElem

func (ts TraceSlice) Len() int      { return len(ts) }
func (ts TraceSlice) Swap(i, j int) { ts[i], ts[j] = ts[j], ts[i] }

func (ts TraceSlice) Less(i, j int) bool {
	return len(ts[i].trace) > len(ts[j].trace)
}

func (ts *TraceSlice) Push(te interface{}) { *ts = append(*ts, te.(TraceElem)) }
func (ts *TraceSlice) Pop() interface{} {
	old := *ts
	n := len(old)

	*ts = old[0 : n-1]
	out := old[n-1]

	return out
}

// Compact takes a trace slice and merges all of its fields.
//
// index: Union of all locations in the slice.
// names: Union of all names in the slice.
// trace: Union of all the traces in the slice.
func (ts TraceSlice) Compact() (index []int, names []string, trace []string) {
	for _, te := range ts {
		index = append(index, te.loc)
		names = append(names, te.names...)
		trace = append(trace, te.trace...)
	}

	// This is a QuickSort related algorithm.  It makes all the names in the trace unique so we don't double-count people.
	//
	// Invariant:  There are no duplicates in trace[0:ptr]
	// Algorithm:  Advance ptr by 1 and enforce the invariant.
	ptr, cutoff := 0, len(trace)

TopLoop:
	for ptr < cutoff { // Choose the next un-checked element of the slice.
		for i := 0; i < ptr; i++ { // Compare it to all elements before it.
			if trace[i] == trace[ptr] { // If we find a duplicate...
				trace[ptr], trace[cutoff-1] = trace[cutoff-1], trace[ptr] // Push the dup to the end of the surviving slice.
				cutoff--                                                  // Mark it for removal.

				continue TopLoop // Because trace[ptr] has been mutated, try to verify the invariant again w/o advancing ptr.
			}
		}

		ptr++ // There are no duplicates; move the ptr forward and start again.
	}
	trace = trace[0:cutoff]

	return
}

type MSP Formatted

func StringToMSP(pred string) (m MSP, err error) {
	var f Formatted

	if -1 == strings.Index(pred, ",") {
		var r Raw
		r, err = StringToRaw(pred)
		if err != nil {
			return
		}

		f = r.Formatted()
	} else {
		f, err = StringToFormatted(pred)
		if err != nil {
			return
		}
	}

	return MSP(f), nil
}

// DerivePath returns the cheapest way to satisfy the MSP (the one with the minimal number of delegations).
//
// ok:    True if the MSP can be satisfied with current delegations; false if not.
// names: The names in the top-level threshold gate that need to be delegated.
// locs:  The index in the treshold gate for each name.
// trace: All names that must be delegated for for this gate to be satisfied.
func (m MSP) DerivePath(db *UserDatabase) (ok bool, names []string, locs []int, trace []string) {
	ts := &TraceSlice{}

	for i, cond := range m.Conds {
		switch cond.(type) {
		case Name:
			if (*db).CanGetShare(cond.(Name).string) {
				heap.Push(ts, TraceElem{
					i,
					[]string{cond.(Name).string},
					[]string{cond.(Name).string},
				})
			}

		case Formatted:
			sok, _, _, strace := MSP(cond.(Formatted)).DerivePath(db)
			if sok {
				heap.Push(ts, TraceElem{i, []string{}, strace})
			}
		}

		if (*ts).Len() > m.Min { // If we can otherwise satisfy the threshold gate
			heap.Pop(ts) // Drop the TraceElem with the heaviest trace (the one that requires the most delegations).
		}
	}

	ok = (*ts).Len() >= m.Min
	locs, names, trace = ts.Compact()
	return
}

// DistributeShares takes as input a secret and a user database and returns secret shares according to access structure
// described by the MSP.
func (m MSP) DistributeShares(sec []byte, db *UserDatabase) (map[string][][]byte, error) {
	out := make(map[string][][]byte)

	// Generate a Vandermonde matrix.
	height, width := len(m.Conds), m.Min
	M := Matrix(make([]Row, height))

	for i := 0; i < height; i++ {
		M[i] = NewRow(width)

		for j := 0; j < width; j++ {
			M[i][j][0] = byte(i + 1)
			M[i][j] = M[i][j].Exp(j)
		}
	}

	// Convert secret vector.
	s := NewRow(width)
	s[0] = FieldElem(sec)

	for i := 1; i < width; i++ {
		r := NewFieldElem()
		rand.Read(r)

		s[i] = FieldElem(r)
	}

	// Calculate shares.
	shares := M.Mul(s)

	// Distribute the shares.
	for i, cond := range m.Conds {
		share := shares[i]

		switch cond.(type) {
		case Name:
			name := cond.(Name).string
			if _, ok := out[name]; ok {
				out[name] = append(out[name], share)
			} else if (*db).ValidUser(name) {
				out[name] = [][]byte{share}
			} else {
				return out, errors.New("Unknown user in predicate.")
			}

		default:
			below := MSP(cond.(Formatted))
			subOut, err := below.DistributeShares(share, db)
			if err != nil {
				return out, err
			}

			for name, shares := range subOut {
				if _, ok := out[name]; ok {
					out[name] = append(out[name], shares...)
				} else {
					out[name] = shares
				}

			}
		}
	}

	return out, nil
}

// RecoverSecret takes a user database storing secret shares as input and returns the original secret.
func (m MSP) RecoverSecret(db *UserDatabase) ([]byte, error) {
	cache := make(map[string][][]byte, 0) // Caches un-used shares for a user.
	return m.recoverSecret(db, cache)
}

func (m MSP) recoverSecret(db *UserDatabase, cache map[string][][]byte) ([]byte, error) {
	var (
		index  = []int{}       // Indexes where given shares were in the matrix.
		shares = []FieldElem{} // Contains shares that will be used in reconstruction.
	)

	ok, names, locs, _ := m.DerivePath(db)
	if !ok {
		return nil, errors.New("Not enough shares to recover.")
	}

	for _, name := range names {
		if _, cached := cache[name]; !cached {
			out, err := (*db).GetShare(name)
			if err != nil {
				return nil, err
			}

			cache[name] = out
		}
	}

	for _, loc := range locs {
		gate := m.Conds[loc]
		index = append(index, loc+1)

		switch gate.(type) {
		case Name:
			if len(cache[gate.(Name).string]) <= gate.(Name).index {
				return nil, errors.New("Predicate / database mismatch!")
			}

			shares = append(shares, FieldElem(cache[gate.(Name).string][gate.(Name).index]))

		case Formatted:
			share, err := MSP(gate.(Formatted)).recoverSecret(db, cache)
			if err != nil {
				return nil, err
			}

			shares = append(shares, FieldElem(share))
		}
	}

	// Generate the Vandermonde matrix specific to whichever users' shares we're using.
	MSub := Matrix(make([]Row, m.Min))

	for i := 0; i < m.Min; i++ {
		MSub[i] = NewRow(m.Min)

		for j := 0; j < m.Min; j++ {
			MSub[i][j][0] = byte(index[i])
			MSub[i][j] = MSub[i][j].Exp(j)
		}
	}

	// Calculate the reconstruction vector and use it to recover the secret.
	r, ok := MSub.Recovery()
	if !ok {
		return nil, errors.New("Unable to find a reconstruction vector!")
	}

	// Compute dot product of the shares vector and the reconstruction vector to
	// recover the secret.
	s := Row(shares).DotProduct(r)

	return []byte(s), nil
}
<|endoftext|>


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

// Common test setup code
package tests

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/GoogleCloudPlatform/cloudsql-proxy/logging"
	"github.com/GoogleCloudPlatform/cloudsql-proxy/proxy/proxy"

	"golang.org/x/crypto/ssh"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"
	"strings"
)

var (
	project      = flag.String("project", "", "Project to create the GCE test VM in")
	zone         = flag.String("zone", "us-central1-f", "Zone in which to create the VM")
	osImage      = flag.String("os", defaultOS, "OS image to use when creating a VM")
	vmName       = flag.String("vm_name", "proxy-test-gce", "Name of VM to create")
	databaseName = flag.String("db_name", "", "Fully-qualified Cloud SQL Instance (in the form of 'project:region:instance-name')")
)

const (
	defaultOS       = "https://www.googleapis.com/compute/v1/projects/debian-cloud/global/images/debian-8-jessie-v20160329"
	testTimeout     = 3 * time.Minute
	buildShLocation = "cmd/cloud_sql_proxy/build.sh"
)

func setupGCEProxy(t *testing.T, proxyArgs []string) (error, *ssh.Client) {
	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
	defer cancel()

	proxyBinary, err := compileProxy()
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("Built new cloud_sql_proxy binary")
	defer os.Remove(proxyBinary)

	cl, err := google.DefaultClient(ctx, "https://www.googleapis.com/auth/cloud-platform")
	if err != nil {
		t.Fatal(err)
	}

	ssh, err := newOrReuseVM(t.Logf, cl)
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("SSH to %s:%s succeeded", *project, *vmName)

	log.Printf("apt-get update...")
	var sout, serr bytes.Buffer
	if err := sshRun(ssh, "sudo apt-get update", nil, &sout, &serr); err != nil {
		t.Fatal("Failed 'sudo apt-get update' on remote machine: %v\n\nstandard out:\n%s\nstandard err:\n%s", err, &sout, &serr)
	}
	log.Printf("Install mysql client...")
	if err := sshRun(ssh, "sudo apt-get install -y mysql-client", nil, &sout, &serr); err != nil {
		t.Fatal("Failed 'sudo apt-get install -y mysql-client' on remote machine: %v\n\nstandard out:\n%s\nstandard err:\n%s", err, &sout, &serr)
	}
	if err = sshRun(ssh, "pkill cloud_sql_proxy", nil, &sout, &serr); err != nil {
		t.Logf("Failed to kill any cloud_sql_proxy process.")
	} else {
		t.Logf("Killed already running cloud_sql_proxy process.")
	}
	log.Printf("Copy binary to %s:%s...", *project, *vmName)
	this, err := os.Open(proxyBinary)
	if err != nil {
		t.Fatalf("Couldn't open %v for reading: %v", proxyBinary, err)
	}
	err = sshRun(ssh, "bash -c 'cat >cloud_sql_proxy; chmod +x cloud_sql_proxy; mkdir -p cloudsql'", this, &sout, &serr)
	this.Close()
	if err != nil {
		t.Fatalf("Couldn't scp to remote machine: %v\n\nstandard out:\n%s\nstandard err:\n%s", err, &sout, &serr)
	}
	logs, err := startProxy(ssh, "./cloud_sql_proxy -dir cloudsql -instances "+strings.Join(append([]string{*databaseName}, proxyArgs...), " "))
	if err != nil {
		t.Fatal(err)
	}
	defer logs.Close()
	// TODO: Instead of discarding all of the logs, verify that certain logs
	// happen during connects/disconnects.
	go io.Copy(ioutil.Discard, logs)
	t.Logf("Cloud SQL Proxy started on remote host")
	return err, ssh
}

var _ io.ReadCloser = (*process)(nil)

// process wraps a remotely executing process, turning it into an
// io.ReadCloser.
type process struct {
	io.Reader
	sess *ssh.Session
}

// TODO: Return the value of 'Wait'ing on the process. ssh.Session.Signal
// doesn't seem to have an effect, so calling it and then doing Wait doesn't do
// anything. Closing the session is the only way to clean up until I figure out
// what's wrong.
func (p *process) Close() error {
	return p.sess.Close()
}

func compileProxy() (binaryPath string, _ error) {
	// Find the 'build.sh' script by looking for it in cwd, cwd/.., and cwd/../..
	var buildSh string

	var parentPath []string
	for parents := 0; parents < 2; parents++ {
		cur := filepath.Join(append(parentPath, buildShLocation)...)
		if _, err := os.Stat(cur); err == nil {
			buildSh = cur
			break
		}
		parentPath = append(parentPath, "..")
	}
	if buildSh == "" {
		return "", fmt.Errorf("couldn't find %q; please cd into the local repository", buildShLocation)
	}

	cmd := exec.Command(buildSh)
	if out, err := cmd.CombinedOutput(); err != nil {
		return "", fmt.Errorf("error during build.sh execution: %v;\n%s", err, out)
	}

	return "cloud_sql_proxy", nil
}

// startProxy executes the cloud_sql_proxy via ssh. The returned ReadCloser
// must be serviced and closed when finished, otherwise the SSH connection may
// block.
func startProxy(ssh *ssh.Client, args string) (io.ReadCloser, error) {
	sess, err := ssh.NewSession()
	if err != nil {
		return nil, fmt.Errorf("couldn't open new session: %v", err)
	}
	pr, err := sess.StderrPipe()
	if err != nil {
		return nil, err
	}
	log.Printf("Running proxy...")
	if err := sess.Start(args); err != nil {
		return nil, err
	}

	// The proxy prints "Ready for new connections" after it starts up
	// correctly. Start a new goroutine looking for that value so that we can
	// time-out appropriately (in case something weird is going on).
	in := bufio.NewReader(pr)
	buf := new(bytes.Buffer)
	errCh := make(chan error, 1)
	go func() {
		for {
			bs, err := in.ReadBytes('\n')
			if err != nil {
				if err == io.EOF {
					log.Print("reading stderr gave EOF (remote process closed)")
					err = sess.Wait()
				}
				errCh <- fmt.Errorf("failed to run `%s`: %v", args, err)
				return
			}
			buf.Write(bs)
			if bytes.Contains(bs, []byte("Ready for new connections")) {
				errCh <- nil
				return
			}
		}
	}()

	select {
	case err := <-errCh:
		if err != nil {
			return nil, err
		}

		// Proxy process startup succeeded.
		return &process{
			io.MultiReader(buf, in),
			sess,
		}, nil
	case <-time.After(3 * time.Second):
		log.Printf("Timeout starting up `%v`", args)
	}

	// Starting the proxy timed out, so we should close the SSH session and
	// return an error after the process exits.
	// TODO: the sess.Signal method doesn't seem to work... that's what we
	// really want to do.
	err = sess.Close()
	select {
	case waitErr := <-errCh:
		if err == nil {
			err = waitErr
		}
	case <-time.After(2 * time.Second):
		log.Printf("Timeout while waiting for process after closing SSH session.")
		if err == nil {
			err = errors.New("timeout waiting for SSH connection to close")
		}
	}
	return nil, fmt.Errorf("timeout waiting for `%v`: error from close: %v; output was:\n\n%s", args, err, buf)
}

func sshRun(ssh *ssh.Client, cmd string, stdin io.Reader, stdout, stderr io.Writer) error {
	sess, err := ssh.NewSession()
	if err != nil {
		return err
	}

	sess.Stdin = stdin
	if stderr == nil && stdout == nil {
		if out, err := sess.CombinedOutput(cmd); err != nil {
			return fmt.Errorf("`%v`: %v; combined output was:\n%s", cmd, err, out)
		}
		return nil
	}
	sess.Stdout = stdout
	sess.Stderr = stderr

	return sess.Run(cmd)
}

func newOrReuseVM(logf func(string, ...interface{}), cl *http.Client) (*ssh.Client, error) {
	c, err := compute.New(cl)
	if err != nil {
		return nil, err
	}

	user := "test-user"
	pub, auth, err := sshKey()
	if err != nil {
		return nil, err
	}
	sshPubKey := user + ":" + pub

	var op *compute.Operation

	if inst, err := c.Instances.Get(*project, *zone, *vmName).Do(); err != nil {
		logf("Creating new instance (getting instance %v in project %v and zone %v failed: %v)", *vmName, *project, *zone, err)
		instProto := &compute.Instance{
			Name:        *vmName,
			MachineType: "zones/" + *zone + "/machineTypes/g1-small",
			Disks: []*compute.AttachedDisk{{
				AutoDelete: true,
				Boot:       true,
				InitializeParams: &compute.AttachedDiskInitializeParams{
					SourceImage: *osImage,
					DiskSizeGb:  10,
				}},
			},
			NetworkInterfaces: []*compute.NetworkInterface{{
				Network:       "projects/" + *project + "/global/networks/default",
				AccessConfigs: []*compute.AccessConfig{{Name: "External NAT", Type: "ONE_TO_ONE_NAT"}},
			}},
			Metadata: &compute.Metadata{
				Items: []*compute.MetadataItems{{
					Key: "sshKeys", Value: &sshPubKey,
				}},
			},
			Tags: &compute.Tags{Items: []string{"ssh"}},
			ServiceAccounts: []*compute.ServiceAccount{{
				Email:  "default",
				Scopes: []string{proxy.SQLScope},
			}},
		}
		op, err = c.Instances.Insert(*project, *zone, instProto).Do()
		if err != nil {
			return nil, err
		}
	} else {
		logf("attempting to reuse instance %v (in project %v and zone %v)...", *vmName, *project, *zone)
		set := false
		md := inst.Metadata
		for _, v := range md.Items {
			if v.Key == "sshKeys" {
				v.Value = &sshPubKey
				set = true
				break
			}
		}
		if !set {
			md.Items = append(md.Items, &compute.MetadataItems{Key: "sshKeys", Value: &sshPubKey})
		}
		op, err = c.Instances.SetMetadata(*project, *zone, *vmName, md).Do()
		if err != nil {
			return nil, err
		}
	}

	for {
		if op.Error != nil && len(op.Error.Errors) > 0 {
			return nil, fmt.Errorf("errors: %v", op.Error.Errors)
		}

		log.Printf("%v %v (%v)", op.OperationType, op.TargetLink, op.Status)
		if op.Status == "DONE" {
			break
		}
		time.Sleep(5 * time.Second)

		op, err = c.ZoneOperations.Get(*project, *zone, op.Name).Do()
		if err != nil {
			return nil, err
		}
	}

	inst, err := c.Instances.Get(*project, *zone, *vmName).Do()
	if err != nil {
		return nil, fmt.Errorf("error getting instance after it was created: %v", err)
	}
	ip := inst.NetworkInterfaces[0].AccessConfigs[0].NatIP

	var lastErr error
	for try := 0; try < 10; try++ {
		if lastErr != nil {
			const sleepTime = 10 * time.Second
			logging.Errorf("%v; sleeping for %v then retrying", lastErr, sleepTime)
			time.Sleep(sleepTime)
		}
		ssh, err := ssh.Dial("tcp", ip+":22", &ssh.ClientConfig{
			User:            user,
			Auth:            []ssh.AuthMethod{auth},
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		})
		if err == nil {
			return ssh, nil
		}
		lastErr = fmt.Errorf("couldn't ssh to %v (IP=%v): %v", *vmName, ip, err)
	}
	return nil, lastErr
}

func sshKey() (pubKey string, auth ssh.AuthMethod, err error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", nil, err
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return "", nil, err
	}
	pub, err := ssh.NewPublicKey(&key.PublicKey)
	if err != nil {
		return "", nil, err
	}
	return string(ssh.MarshalAuthorizedKey(pub)), ssh.PublicKeys(signer), nil
}

func TestMain(m *testing.M) {
	flag.Parse()
	switch "" {
	case *project:
		log.Fatal("Must set -project")
	case *databaseName:
		log.Fatal("Must set -db_name")
	}

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

// Common test setup code
package tests

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/GoogleCloudPlatform/cloudsql-proxy/logging"
	"github.com/GoogleCloudPlatform/cloudsql-proxy/proxy/proxy"

	"golang.org/x/crypto/ssh"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"
	"strings"
)

var (
	project      = flag.String("project", "", "Project to create the GCE test VM in")
	zone         = flag.String("zone", "us-central1-f", "Zone in which to create the VM")
	osImage      = flag.String("os", defaultOS, "OS image to use when creating a VM")
	vmName       = flag.String("vm_name", "proxy-test-gce", "Name of VM to create")
	databaseName = flag.String("db_name", "", "Fully-qualified Cloud SQL Instance (in the form of 'project:region:instance-name')")
)

const (
	defaultOS       = "https://www.googleapis.com/compute/v1/projects/debian-cloud/global/images/debian-8-jessie-v20160329"
	testTimeout     = 3 * time.Minute
	buildShLocation = "cmd/cloud_sql_proxy/build.sh"
)

func setupGCEProxy(t *testing.T, proxyArgs []string) (error, *ssh.Client) {
	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
	defer cancel()

	proxyBinary, err := compileProxy()
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("Built new cloud_sql_proxy binary")
	defer os.Remove(proxyBinary)

	cl, err := google.DefaultClient(ctx, "https://www.googleapis.com/auth/cloud-platform")
	if err != nil {
		t.Fatal(err)
	}

	ssh, err := newOrReuseVM(t.Logf, cl)
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("SSH to %s:%s succeeded", *project, *vmName)

	log.Printf("apt-get update...")
	var sout, serr bytes.Buffer
	if err := sshRun(ssh, "sudo apt-get update", nil, &sout, &serr); err != nil {
		t.Fatal("Failed 'sudo apt-get update' on remote machine: %v\n\nstandard out:\n%s\nstandard err:\n%s", err, &sout, &serr)
	}
	log.Printf("Install mysql client...")
	if err := sshRun(ssh, "sudo apt-get install -y mysql-client", nil, &sout, &serr); err != nil {
		t.Fatal("Failed 'sudo apt-get install -y mysql-client' on remote machine: %v\n\nstandard out:\n%s\nstandard err:\n%s", err, &sout, &serr)
	}
	if err = sshRun(ssh, "pkill cloud_sql_proxy", nil, &sout, &serr); err != nil {
		t.Logf("Failed to kill any cloud_sql_proxy process.")
	} else {
		t.Logf("Killed already running cloud_sql_proxy process.")
	}
	log.Printf("Copy binary to %s:%s...", *project, *vmName)
	this, err := os.Open(proxyBinary)
	if err != nil {
		t.Fatalf("Couldn't open %v for reading: %v", proxyBinary, err)
	}
	err = sshRun(ssh, "bash -c 'cat >cloud_sql_proxy; chmod +x cloud_sql_proxy; mkdir -p cloudsql'", this, &sout, &serr)
	this.Close()
	if err != nil {
		t.Fatalf("Couldn't scp to remote machine: %v\n\nstandard out:\n%s\nstandard err:\n%s", err, &sout, &serr)
	}
	logs, err := startProxy(ssh, "./cloud_sql_proxy -dir cloudsql -instances "+strings.Join(append([]string{*databaseName}, proxyArgs...), " "))
	if err != nil {
		t.Fatal(err)
	}
	defer logs.Close()
	// TODO: Instead of discarding all of the logs, verify that certain logs
	// happen during connects/disconnects.
	go io.Copy(ioutil.Discard, logs)
	t.Logf("Cloud SQL Proxy started on remote host")
	return err, ssh
}

var _ io.ReadCloser = (*process)(nil)

// process wraps a remotely executing process, turning it into an
// io.ReadCloser.
type process struct {
	io.Reader
	sess *ssh.Session
}

// TODO: Return the value of 'Wait'ing on the process. ssh.Session.Signal
// doesn't seem to have an effect, so calling it and then doing Wait doesn't do
// anything. Closing the session is the only way to clean up until I figure out
// what's wrong.
func (p *process) Close() error {
	return p.sess.Close()
}

func compileProxy() (binaryPath string, _ error) {
	// Find the 'build.sh' script by looking for it in cwd, cwd/.., and cwd/../..
	var buildSh string

	var parentPath []string
	for parents := 0; parents < 2; parents++ {
		cur := filepath.Join(append(parentPath, buildShLocation)...)
		if _, err := os.Stat(cur); err == nil {
			buildSh = cur
			break
		}
		parentPath = append(parentPath, "..")
	}
	if buildSh == "" {
		return "", fmt.Errorf("couldn't find %q; please cd into the local repository", buildShLocation)
	}

	cmd := exec.Command(buildSh)
	if out, err := cmd.CombinedOutput(); err != nil {
		return "", fmt.Errorf("error during build.sh execution: %v;\n%s", err, out)
	}

	return "cloud_sql_proxy", nil
}

// startProxy executes the cloud_sql_proxy via ssh. The returned ReadCloser
// must be serviced and closed when finished, otherwise the SSH connection may
// block.
func startProxy(ssh *ssh.Client, args string) (io.ReadCloser, error) {
	sess, err := ssh.NewSession()
	if err != nil {
		return nil, fmt.Errorf("couldn't open new session: %v", err)
	}
	pr, err := sess.StderrPipe()
	if err != nil {
		return nil, err
	}
	log.Printf("Running proxy...")
	if err := sess.Start(args); err != nil {
		return nil, err
	}

	// The proxy prints "Ready for new connections" after it starts up
	// correctly. Start a new goroutine looking for that value so that we can
	// time-out appropriately (in case something weird is going on).
	in := bufio.NewReader(pr)
	buf := new(bytes.Buffer)
	errCh := make(chan error, 1)
	go func() {
		for {
			bs, err := in.ReadBytes('\n')
			if err != nil {
				if err == io.EOF {
					log.Print("reading stderr gave EOF (remote process closed)")
					err = sess.Wait()
				}
				errCh <- fmt.Errorf("failed to run `%s`: %v", args, err)
				return
			}
			buf.Write(bs)
			if bytes.Contains(bs, []byte("Ready for new connections")) {
				errCh <- nil
				return
			}
		}
	}()

	select {
	case err := <-errCh:
		if err != nil {
			return nil, err
		}

		// Proxy process startup succeeded.
		return &process{
			io.MultiReader(buf, in),
			sess,
		}, nil
	case <-time.After(3 * time.Second):
		log.Printf("Timeout starting up `%v`", args)
	}

	// Starting the proxy timed out, so we should close the SSH session and
	// return an error after the process exits.
	// TODO: the sess.Signal method doesn't seem to work... that's what we
	// really want to do.
	err = sess.Close()
	select {
	case waitErr := <-errCh:
		if err == nil {
			err = waitErr
		}
	case <-time.After(2 * time.Second):
		log.Printf("Timeout while waiting for process after closing SSH session.")
		if err == nil {
			err = errors.New("timeout waiting for SSH connection to close")
		}
	}
	return nil, fmt.Errorf("timeout waiting for `%v`: error from close: %v; output was:\n\n%s", args, err, buf)
}

func sshRun(ssh *ssh.Client, cmd string, stdin io.Reader, stdout, stderr io.Writer) error {
	sess, err := ssh.NewSession()
	if err != nil {
		return err
	}

	sess.Stdin = stdin
	if stderr == nil && stdout == nil {
		if out, err := sess.CombinedOutput(cmd); err != nil {
			return fmt.Errorf("`%v`: %v; combined output was:\n%s", cmd, err, out)
		}
		return nil
	}
	sess.Stdout = stdout
	sess.Stderr = stderr

	return sess.Run(cmd)
}

func newOrReuseVM(logf func(string, ...interface{}), cl *http.Client) (*ssh.Client, error) {
	c, err := compute.New(cl)
	if err != nil {
		return nil, err
	}

	user := "test-user"
	pub, auth, err := sshKey()
	if err != nil {
		return nil, err
	}
	sshPubKey := user + ":" + pub

	var op *compute.Operation

	if inst, err := c.Instances.Get(*project, *zone, *vmName).Do(); err != nil {
		logf("Creating new instance (getting instance %v in project %v and zone %v failed: %v)", *vmName, *project, *zone, err)
		instProto := &compute.Instance{
			Name:        *vmName,
			MachineType: "zones/" + *zone + "/machineTypes/g1-small",
			Disks: []*compute.AttachedDisk{{
				AutoDelete: true,
				Boot:       true,
				InitializeParams: &compute.AttachedDiskInitializeParams{
					SourceImage: *osImage,
					DiskSizeGb:  10,
				}},
			},
			NetworkInterfaces: []*compute.NetworkInterface{{
				Network:       "projects/" + *project + "/global/networks/default",
				AccessConfigs: []*compute.AccessConfig{{Name: "External NAT", Type: "ONE_TO_ONE_NAT"}},
			}},
			Metadata: &compute.Metadata{
				Items: []*compute.MetadataItems{{
					Key: "sshKeys", Value: &sshPubKey,
				}},
			},
			Tags: &compute.Tags{Items: []string{"ssh"}},
			ServiceAccounts: []*compute.ServiceAccount{{
				Email:  "default",
				Scopes: []string{proxy.SQLScope},
			}},
		}
		op, err = c.Instances.Insert(*project, *zone, instProto).Do()
		if err != nil {
			return nil, err
		}
	} else {
		logf("attempting to reuse instance %v (in project %v and zone %v)...", *vmName, *project, *zone)
		set := false
		md := inst.Metadata
		for _, v := range md.Items {
			if v.Key == "sshKeys" {
				v.Value = &sshPubKey
				set = true
				break
			}
		}
		if !set {
			md.Items = append(md.Items, &compute.MetadataItems{Key: "sshKeys", Value: &sshPubKey})
		}
		op, err = c.Instances.SetMetadata(*project, *zone, *vmName, md).Do()
		if err != nil {
			return nil, err
		}
	}

	for {
		if op.Error != nil && len(op.Error.Errors) > 0 {
			return nil, fmt.Errorf("errors: %v", op.Error.Errors)
		}

		log.Printf("%v %v (%v)", op.OperationType, op.TargetLink, op.Status)
		if op.Status == "DONE" {
			break
		}
		time.Sleep(5 * time.Second)

		op, err = c.ZoneOperations.Get(*project, *zone, op.Name).Do()
		if err != nil {
			return nil, err
		}
	}

	inst, err := c.Instances.Get(*project, *zone, *vmName).Do()
	if err != nil {
		return nil, fmt.Errorf("error getting instance after it was created: %v", err)
	}
	ip := inst.NetworkInterfaces[0].NetworkIP

	var lastErr error
	for try := 0; try < 10; try++ {
		if lastErr != nil {
			const sleepTime = 10 * time.Second
			logging.Errorf("%v; sleeping for %v then retrying", lastErr, sleepTime)
			time.Sleep(sleepTime)
		}
		ssh, err := ssh.Dial("tcp", ip+":22", &ssh.ClientConfig{
			User:            user,
			Auth:            []ssh.AuthMethod{auth},
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		})
		if err == nil {
			return ssh, nil
		}
		lastErr = fmt.Errorf("couldn't ssh to %v (IP=%v): %v", *vmName, ip, err)
	}
	return nil, lastErr
}

func sshKey() (pubKey string, auth ssh.AuthMethod, err error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", nil, err
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return "", nil, err
	}
	pub, err := ssh.NewPublicKey(&key.PublicKey)
	if err != nil {
		return "", nil, err
	}
	return string(ssh.MarshalAuthorizedKey(pub)), ssh.PublicKeys(signer), nil
}

func TestMain(m *testing.M) {
	flag.Parse()
	switch "" {
	case *project:
		log.Fatal("Must set -project")
	case *databaseName:
		log.Fatal("Must set -db_name")
	}

	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>package mux

import (
	"github.com/freehaha/go-r3"
	"github.com/gorilla/context"
	"log"
	"net/http"
	"runtime"
)

type Router struct {
	tree            *r3.Tree
	NotFoundHandler http.Handler
}

var methods map[string]r3.Method = map[string]r3.Method{
	"GET":     r3.MethodGet,
	"PUT":     r3.MethodPut,
	"POST":    r3.MethodPost,
	"DELETE":  r3.MethodDelete,
	"OPTIONS": r3.MethodOptions,
	"PATCH":   r3.MethodPatch,
}

/* instanciate a new Router */
func NewRouter() *Router {
	r := &Router{}
	r.tree = r3.NewTree(10)
	runtime.SetFinalizer(r, finalizeRouter)
	return r
}

func finalizeRouter(r *Router) {
	log.Print("finalizing router")
	r.tree.Free()
}

func (r *Router) Compile() error {
	return r.tree.Compile()
}

func (r *Router) Free() {
	finalizeRouter(r)
}

func Vars(r *http.Request) []string {
	return context.Get(r, "vars").([]string)
}

/* insert a path to the router with specific handlerfunc */
func (r *Router) HandleFunc(method r3.Method, path string, handler http.HandlerFunc) {
	r.Handle(method, path, http.HandlerFunc(handler))
}

/* insert a path to the router with specific handler */
func (r *Router) Handle(method r3.Method, path string, handler http.Handler) {
	r.tree.InsertRoute(method, path, handler)
	/* keep a reference on it so it doens't get GCed */
}

/* implement Mux */
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	ent := r3.NewMatchEntry(req.URL.Path)
	ent.SetRequestMethod(methods[req.Method])
	log.Print(req.URL.Path)
	log.Print(req.Method)
	defer ent.Free()
	route := r.tree.MatchRoute(ent)
	if route != nil {
		context.Set(req, "vars", ent.Vars())
		handler := route.Data().(http.Handler)
		handler.ServeHTTP(w, req)
	} else {
		handler := r.NotFoundHandler
		if handler == nil {
			handler = http.NotFoundHandler()
		}
		handler.ServeHTTP(w, req)
	}
}
<commit_msg>add Negroni middleware implementation and remove debugging messages<commit_after>package mux

import (
	"github.com/freehaha/go-r3"
	"github.com/gorilla/context"
	"net/http"
	"runtime"
)

type Router struct {
	tree            *r3.Tree
	NotFoundHandler http.Handler
}

type NegroniRouter struct {
	Router
}

var methods map[string]r3.Method = map[string]r3.Method{
	"GET":     r3.MethodGet,
	"PUT":     r3.MethodPut,
	"POST":    r3.MethodPost,
	"DELETE":  r3.MethodDelete,
	"OPTIONS": r3.MethodOptions,
	"PATCH":   r3.MethodPatch,
}

/* instanciate a new Router */
func NewRouter() *Router {
	r := &Router{}
	r.tree = r3.NewTree(10)
	runtime.SetFinalizer(r, finalizeRouter)
	return r
}

func finalizeRouter(r *Router) {
	r.tree.Free()
}

func finalizeNRouter(r *NegroniRouter) {
	r.tree.Free()
}

func (r *Router) Compile() error {
	return r.tree.Compile()
}

func (r *Router) Free() {
	finalizeRouter(r)
}

func Vars(r *http.Request) []string {
	return context.Get(r, "vars").([]string)
}

/* insert a path to the router with specific handlerfunc */
func (r *Router) HandleFunc(method r3.Method, path string, handler http.HandlerFunc) {
	r.Handle(method, path, http.HandlerFunc(handler))
}

/* insert a path to the router with specific handler */
func (r *Router) Handle(method r3.Method, path string, handler http.Handler) {
	r.tree.InsertRoute(method, path, handler)
	/* keep a reference on it so it doens't get GCed */
}

/* implement Mux */
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	ent := r3.NewMatchEntry(req.URL.Path)
	ent.SetRequestMethod(methods[req.Method])
	defer ent.Free()
	route := r.tree.MatchRoute(ent)
	if route != nil {
		context.Set(req, "vars", ent.Vars())
		handler := route.Data().(http.Handler)
		handler.ServeHTTP(w, req)
	} else {
		handler := r.NotFoundHandler
		if handler == nil {
			handler = http.NotFoundHandler()
		}
		handler.ServeHTTP(w, req)
	}
}

func NewNegroniRouter() *NegroniRouter {
	r := &NegroniRouter{
		Router{},
	}
	r.tree = r3.NewTree(10)
	runtime.SetFinalizer(r, finalizeNRouter)
	return r
}

/* implement negroni middleware */
func (r *NegroniRouter) ServeHTTP(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
	ent := r3.NewMatchEntry(req.URL.Path)
	ent.SetRequestMethod(methods[req.Method])
	defer ent.Free()
	route := r.tree.MatchRoute(ent)
	if route != nil {
		context.Set(req, "vars", ent.Vars())
		handler := route.Data().(http.Handler)
		handler.ServeHTTP(w, req)
	} else {
		/* skip to next middleware */
		next(w, req)
	}
}
<|endoftext|>


<commit_before>package toscaviewer

// This is a basic example
// Thanks http://thenewstack.io/make-a-restful-json-api-go/ for the tutorial
import (
	"fmt"
	"io"

	"github.com/owulveryck/toscalib"
)

// PrintDot convert the ToscaDefinitionure in dot format
// in order to generate a graph with graphviz
// This function is mostly used for debugging purpose and may change a lot in the future
func PrintDot(w io.Writer, topology toscalib.ToscaDefinition) {
	dotCode := fmt.Sprintf("digraph G {\n")
	dotCode = fmt.Sprintf("%v\tgraph [ rankdir = \"LR\" ];\n", dotCode)
	for nodeName, nodeDetail := range topology.TopologyTemplate.NodeTemplates {
		// For each node, create a record
		dotCode = fmt.Sprintf("%v\t%v [label=\"<nodeName> %v|<nodeType> %v", dotCode, nodeName, nodeName, nodeDetail.Type)
		//Display the properties
		if nodeDetail.Properties != nil {
			dotCode = fmt.Sprintf("%v|{Properties|{", dotCode)
			for property, _ := range nodeDetail.Properties {
				dotCode = fmt.Sprintf("%v%v|", dotCode, property)
			}
			dotCode = fmt.Sprintf("%v}}", dotCode)
		}
		// Display the requirements
		if nodeDetail.Requirements != nil {
			dotCode = fmt.Sprintf("%v|{Requirements|{", dotCode)
			i := 0
			pipe := "|"
			for _, requirementAssignements := range nodeDetail.Requirements {
				for requirement, _ := range requirementAssignements {
					if i == len(requirementAssignements) {
						pipe = ""
					}
					i = i + 1
					dotCode = fmt.Sprintf("%v<%v>%v%v", dotCode, requirement, requirement, pipe)
				}
			}
			dotCode = fmt.Sprintf("%v}}", dotCode)
		}
		// Display the capabilities
		dotCode = fmt.Sprintf("%v|{<capabilities>Capabilities|{", dotCode)
		if nodeDetail.Capabilities != nil {
			i := 1
			pipe := "|"
			for capability, _ := range nodeDetail.Capabilities {
				if i == len(nodeDetail.Capabilities) {
					pipe = ""
				}
				i = i + 1
				dotCode = fmt.Sprintf("%v<%v>%v%v", dotCode, capability, capability, pipe)
			}
		}
		dotCode = fmt.Sprintf("%v}}", dotCode)

		dotCode = fmt.Sprintf("%v\" shape=record style=rounded color=blue]\n", dotCode)
		// If requirements are found
		//		dotCode = fmt.Sprintf( "\t\"%v\" [ label = \"%v\" shape = circle color=blue]\n", nodeName, nodeName)
		//		dotCode = fmt.Sprintf( "\t\"%v\" [ label = \"%v\" shape = record color=red]\n", nodeDetail.Type, nodeDetail.Type)
		//dotCode = fmt.Sprintf( "\t\"%v\" -> \"%v\" [ color = red ]\n", nodeDetail.Type, nodeName)
		if nodeDetail.Requirements != nil {
			for _, requirementAssignements := range nodeDetail.Requirements {
				for requirementName, requirementAssignement := range requirementAssignements {

					dotCode = fmt.Sprintf("%v\t%v:%v -> %v:capabilities [label = %v color=red];\n", dotCode, nodeName, requirementName, requirementAssignement.Node, requirementName)
				}
			}
		}
	}
	dotCode = fmt.Sprintf("%v}\n", dotCode)
	//fmt.Printf("DEBUG: \n%v\n", dotCode)
	fmt.Fprintf(w, "%v", dotCode)
}
<commit_msg>moved to the toscalib<commit_after><|endoftext|>


<commit_before>package rpc

import (
	"encoding/gob"
	"errors"
	"github.com/mitchellh/packer/packer"
	"io"
	"log"
	"net"
	"net/rpc"
	"time"
)

// An implementation of packer.Communicator where the communicator is actually
// executed over an RPC connection.
type communicator struct {
	client *rpc.Client
}

// CommunicatorServer wraps a packer.Communicator implementation and makes
// it exportable as part of a Golang RPC server.
type CommunicatorServer struct {
	c packer.Communicator
}

type CommandFinished struct {
	ExitStatus int
}

type CommunicatorStartArgs struct {
	Command         string
	StdinAddress    string
	StdoutAddress   string
	StderrAddress   string
	ResponseAddress string
}

type CommunicatorDownloadArgs struct {
	Path          string
	WriterAddress string
}

type CommunicatorUploadArgs struct {
	Path          string
	ReaderAddress string
}

func Communicator(client *rpc.Client) *communicator {
	return &communicator{client}
}

func (c *communicator) Start(cmd *packer.RemoteCmd) (err error) {
	var args CommunicatorStartArgs
	args.Command = cmd.Command

	if cmd.Stdin != nil {
		stdinL := netListenerInRange(portRangeMin, portRangeMax)
		args.StdinAddress = stdinL.Addr().String()
		go serveSingleCopy("stdin", stdinL, nil, cmd.Stdin)
	}

	if cmd.Stdout != nil {
		stdoutL := netListenerInRange(portRangeMin, portRangeMax)
		args.StdoutAddress = stdoutL.Addr().String()
		go serveSingleCopy("stdout", stdoutL, cmd.Stdout, nil)
	}

	if cmd.Stderr != nil {
		stderrL := netListenerInRange(portRangeMin, portRangeMax)
		args.StderrAddress = stderrL.Addr().String()
		go serveSingleCopy("stderr", stderrL, cmd.Stderr, nil)
	}

	responseL := netListenerInRange(portRangeMin, portRangeMax)
	args.ResponseAddress = responseL.Addr().String()

	go func() {
		defer responseL.Close()

		conn, err := responseL.Accept()
		if err != nil {
			log.Panic(err)
		}

		defer conn.Close()

		decoder := gob.NewDecoder(conn)

		var finished CommandFinished
		if err := decoder.Decode(&finished); err != nil {
			log.Panic(err)
		}

		cmd.ExitStatus = finished.ExitStatus
		cmd.Exited = true
	}()

	err = c.client.Call("Communicator.Start", &args, new(interface{}))
	return
}

func (c *communicator) Upload(path string, r io.Reader) (err error) {
	// We need to create a server that can proxy the reader data
	// over because we can't simply gob encode an io.Reader
	readerL := netListenerInRange(portRangeMin, portRangeMax)
	if readerL == nil {
		err = errors.New("couldn't allocate listener for upload reader")
		return
	}

	// Make sure at the end of this call, we close the listener
	defer readerL.Close()

	// Pipe the reader through to the connection
	go serveSingleCopy("uploadReader", readerL, nil, r)

	args := CommunicatorUploadArgs{
		path,
		readerL.Addr().String(),
	}

	err = c.client.Call("Communicator.Upload", &args, new(interface{}))
	return
}

func (c *communicator) Download(path string, w io.Writer) (err error) {
	// We need to create a server that can proxy that data downloaded
	// into the writer because we can't gob encode a writer directly.
	writerL := netListenerInRange(portRangeMin, portRangeMax)
	if writerL == nil {
		err = errors.New("couldn't allocate listener for download writer")
		return
	}

	// Make sure we close the listener once we're done because we'll be done
	defer writerL.Close()

	// Serve a single connection and a single copy
	go serveSingleCopy("downloadWriter", writerL, w, nil)

	args := CommunicatorDownloadArgs{
		path,
		writerL.Addr().String(),
	}

	err = c.client.Call("Communicator.Download", &args, new(interface{}))
	return
}

func (c *CommunicatorServer) Start(args *CommunicatorStartArgs, reply *interface{}) (err error) {
	// Build the RemoteCmd on this side so that it all pipes over
	// to the remote side.
	var cmd packer.RemoteCmd
	cmd.Command = args.Command

	toClose := make([]net.Conn, 0)
	if args.StdinAddress != "" {
		stdinC, err := net.Dial("tcp", args.StdinAddress)
		if err != nil {
			return err
		}

		toClose = append(toClose, stdinC)
		cmd.Stdin = stdinC
	}

	if args.StdoutAddress != "" {
		stdoutC, err := net.Dial("tcp", args.StdoutAddress)
		if err != nil {
			return err
		}

		toClose = append(toClose, stdoutC)
		cmd.Stdout = stdoutC
	}

	if args.StderrAddress != "" {
		stderrC, err := net.Dial("tcp", args.StderrAddress)
		if err != nil {
			return err
		}

		toClose = append(toClose, stderrC)
		cmd.Stderr = stderrC
	}

	// Connect to the response address so we can write our result to it
	// when ready.
	responseC, err := net.Dial("tcp", args.ResponseAddress)
	if err != nil {
		return err
	}

	responseWriter := gob.NewEncoder(responseC)

	// Start the actual command
	err = c.c.Start(&cmd)

	// Start a goroutine to spin and wait for the process to actual
	// exit. When it does, report it back to caller...
	go func() {
		defer responseC.Close()
		for _, conn := range toClose {
			defer conn.Close()
		}

		for !cmd.Exited {
			time.Sleep(50 * time.Millisecond)
		}

		responseWriter.Encode(&CommandFinished{cmd.ExitStatus})
	}()

	return
}

func (c *CommunicatorServer) Upload(args *CommunicatorUploadArgs, reply *interface{}) (err error) {
	readerC, err := net.Dial("tcp", args.ReaderAddress)
	if err != nil {
		return
	}

	defer readerC.Close()

	err = c.c.Upload(args.Path, readerC)
	return
}

func (c *CommunicatorServer) Download(args *CommunicatorDownloadArgs, reply *interface{}) (err error) {
	writerC, err := net.Dial("tcp", args.WriterAddress)
	if err != nil {
		return
	}

	defer writerC.Close()

	err = c.c.Download(args.Path, writerC)
	return
}

func serveSingleCopy(name string, l net.Listener, dst io.Writer, src io.Reader) {
	defer l.Close()

	conn, err := l.Accept()
	if err != nil {
		log.Printf("'%s' accept error: %s", name, err)
		return
	}

	// Be sure to close the connection after we're done copying so
	// that an EOF will successfully be sent to the remote side
	defer conn.Close()

	// The connection is the destination/source that is nil
	if dst == nil {
		dst = conn
	} else {
		src = conn
	}

	written, err := io.Copy(dst, src)
	log.Printf("%d bytes written for '%s'", written, name)
	if err != nil {
		log.Printf("'%s' copy error: %s", name, err)
	}
}
<commit_msg>packer/rpc: Use the SetExited for remote communicators<commit_after>package rpc

import (
	"encoding/gob"
	"errors"
	"github.com/mitchellh/packer/packer"
	"io"
	"log"
	"net"
	"net/rpc"
	"time"
)

// An implementation of packer.Communicator where the communicator is actually
// executed over an RPC connection.
type communicator struct {
	client *rpc.Client
}

// CommunicatorServer wraps a packer.Communicator implementation and makes
// it exportable as part of a Golang RPC server.
type CommunicatorServer struct {
	c packer.Communicator
}

type CommandFinished struct {
	ExitStatus int
}

type CommunicatorStartArgs struct {
	Command         string
	StdinAddress    string
	StdoutAddress   string
	StderrAddress   string
	ResponseAddress string
}

type CommunicatorDownloadArgs struct {
	Path          string
	WriterAddress string
}

type CommunicatorUploadArgs struct {
	Path          string
	ReaderAddress string
}

func Communicator(client *rpc.Client) *communicator {
	return &communicator{client}
}

func (c *communicator) Start(cmd *packer.RemoteCmd) (err error) {
	var args CommunicatorStartArgs
	args.Command = cmd.Command

	if cmd.Stdin != nil {
		stdinL := netListenerInRange(portRangeMin, portRangeMax)
		args.StdinAddress = stdinL.Addr().String()
		go serveSingleCopy("stdin", stdinL, nil, cmd.Stdin)
	}

	if cmd.Stdout != nil {
		stdoutL := netListenerInRange(portRangeMin, portRangeMax)
		args.StdoutAddress = stdoutL.Addr().String()
		go serveSingleCopy("stdout", stdoutL, cmd.Stdout, nil)
	}

	if cmd.Stderr != nil {
		stderrL := netListenerInRange(portRangeMin, portRangeMax)
		args.StderrAddress = stderrL.Addr().String()
		go serveSingleCopy("stderr", stderrL, cmd.Stderr, nil)
	}

	responseL := netListenerInRange(portRangeMin, portRangeMax)
	args.ResponseAddress = responseL.Addr().String()

	go func() {
		defer responseL.Close()

		conn, err := responseL.Accept()
		if err != nil {
			log.Panic(err)
		}

		defer conn.Close()

		decoder := gob.NewDecoder(conn)

		var finished CommandFinished
		if err := decoder.Decode(&finished); err != nil {
			log.Panic(err)
		}

		cmd.SetExited(finished.ExitStatus)
	}()

	err = c.client.Call("Communicator.Start", &args, new(interface{}))
	return
}

func (c *communicator) Upload(path string, r io.Reader) (err error) {
	// We need to create a server that can proxy the reader data
	// over because we can't simply gob encode an io.Reader
	readerL := netListenerInRange(portRangeMin, portRangeMax)
	if readerL == nil {
		err = errors.New("couldn't allocate listener for upload reader")
		return
	}

	// Make sure at the end of this call, we close the listener
	defer readerL.Close()

	// Pipe the reader through to the connection
	go serveSingleCopy("uploadReader", readerL, nil, r)

	args := CommunicatorUploadArgs{
		path,
		readerL.Addr().String(),
	}

	err = c.client.Call("Communicator.Upload", &args, new(interface{}))
	return
}

func (c *communicator) Download(path string, w io.Writer) (err error) {
	// We need to create a server that can proxy that data downloaded
	// into the writer because we can't gob encode a writer directly.
	writerL := netListenerInRange(portRangeMin, portRangeMax)
	if writerL == nil {
		err = errors.New("couldn't allocate listener for download writer")
		return
	}

	// Make sure we close the listener once we're done because we'll be done
	defer writerL.Close()

	// Serve a single connection and a single copy
	go serveSingleCopy("downloadWriter", writerL, w, nil)

	args := CommunicatorDownloadArgs{
		path,
		writerL.Addr().String(),
	}

	err = c.client.Call("Communicator.Download", &args, new(interface{}))
	return
}

func (c *CommunicatorServer) Start(args *CommunicatorStartArgs, reply *interface{}) (err error) {
	// Build the RemoteCmd on this side so that it all pipes over
	// to the remote side.
	var cmd packer.RemoteCmd
	cmd.Command = args.Command

	toClose := make([]net.Conn, 0)
	if args.StdinAddress != "" {
		stdinC, err := net.Dial("tcp", args.StdinAddress)
		if err != nil {
			return err
		}

		toClose = append(toClose, stdinC)
		cmd.Stdin = stdinC
	}

	if args.StdoutAddress != "" {
		stdoutC, err := net.Dial("tcp", args.StdoutAddress)
		if err != nil {
			return err
		}

		toClose = append(toClose, stdoutC)
		cmd.Stdout = stdoutC
	}

	if args.StderrAddress != "" {
		stderrC, err := net.Dial("tcp", args.StderrAddress)
		if err != nil {
			return err
		}

		toClose = append(toClose, stderrC)
		cmd.Stderr = stderrC
	}

	// Connect to the response address so we can write our result to it
	// when ready.
	responseC, err := net.Dial("tcp", args.ResponseAddress)
	if err != nil {
		return err
	}

	responseWriter := gob.NewEncoder(responseC)

	// Start the actual command
	err = c.c.Start(&cmd)

	// Start a goroutine to spin and wait for the process to actual
	// exit. When it does, report it back to caller...
	go func() {
		defer responseC.Close()
		for _, conn := range toClose {
			defer conn.Close()
		}

		for !cmd.Exited {
			time.Sleep(50 * time.Millisecond)
		}

		responseWriter.Encode(&CommandFinished{cmd.ExitStatus})
	}()

	return
}

func (c *CommunicatorServer) Upload(args *CommunicatorUploadArgs, reply *interface{}) (err error) {
	readerC, err := net.Dial("tcp", args.ReaderAddress)
	if err != nil {
		return
	}

	defer readerC.Close()

	err = c.c.Upload(args.Path, readerC)
	return
}

func (c *CommunicatorServer) Download(args *CommunicatorDownloadArgs, reply *interface{}) (err error) {
	writerC, err := net.Dial("tcp", args.WriterAddress)
	if err != nil {
		return
	}

	defer writerC.Close()

	err = c.c.Download(args.Path, writerC)
	return
}

func serveSingleCopy(name string, l net.Listener, dst io.Writer, src io.Reader) {
	defer l.Close()

	conn, err := l.Accept()
	if err != nil {
		log.Printf("'%s' accept error: %s", name, err)
		return
	}

	// Be sure to close the connection after we're done copying so
	// that an EOF will successfully be sent to the remote side
	defer conn.Close()

	// The connection is the destination/source that is nil
	if dst == nil {
		dst = conn
	} else {
		src = conn
	}

	written, err := io.Copy(dst, src)
	log.Printf("%d bytes written for '%s'", written, name)
	if err != nil {
		log.Printf("'%s' copy error: %s", name, err)
	}
}
<|endoftext|>


<commit_before>package rest

import (
	"compress/gzip"
	"io"
	"net/http"
	"strings"
	"sync"

	"github.com/couchbase/sync_gateway/base"
)

// An implementation of http.ResponseWriter that wraps another instance and transparently applies
// GZip compression when appropriate.
type EncodedResponseWriter struct {
	http.ResponseWriter
	gz        *gzip.Writer
	status    int
	sniffDone bool
}

// Creates a new EncodedResponseWriter, or returns nil if the request doesn't allow encoded responses.
func NewEncodedResponseWriter(response http.ResponseWriter, rq *http.Request) *EncodedResponseWriter {
	if !strings.Contains(rq.Header.Get("Accept-Encoding"), "gzip") ||
		rq.Method == "HEAD" || rq.Method == "PUT" || rq.Method == "DELETE" {
		return nil
	}

	// Workaround for https://github.com/couchbase/sync_gateway/issues/1419
	// if the user agent is empty or earlier than 1.2, then we never want to gzip
	// the entire response for requests to the _bulk_get endpoint, since the clients
	// are not equipped to handle that.
	if strings.Contains(rq.URL.Path, "_bulk_get") {
		userAgentVersion := NewUserAgentVersion(rq.Header.Get("User-Agent"))
		if userAgentVersion.IsBefore(1, 2) {
			return nil
		}
	}

	return &EncodedResponseWriter{ResponseWriter: response}
}

func (w *EncodedResponseWriter) WriteHeader(status int) {
	w.status = status
	w.sniff(nil) // Must do it now because headers can't be changed after WriteHeader call
	w.ResponseWriter.WriteHeader(status)
}

func (w *EncodedResponseWriter) Write(b []byte) (int, error) {
	w.sniff(b)
	if w.gz != nil {
		return w.gz.Write(b)
	} else {
		return w.ResponseWriter.Write(b)
	}
}

func (w *EncodedResponseWriter) disableCompression() {
	if w.sniffDone {
		base.Warn("EncodedResponseWriter: Too late to disableCompression!")
	}
	w.sniffDone = true
}

func (w *EncodedResponseWriter) sniff(bytes []byte) {
	if w.sniffDone {
		return
	}
	w.sniffDone = true
	// Check the content type, sniffing the initial data if necessary:
	respType := w.Header().Get("Content-Type")
	if respType == "" && bytes != nil {
		respType = http.DetectContentType(bytes)
		w.Header().Set("Content-Type", respType)
	}

	// Can/should we compress the response?
	if w.status >= 300 || w.Header().Get("Content-Encoding") != "" ||
		(!strings.HasPrefix(respType, "application/json") && !strings.HasPrefix(respType, "text/") && !strings.HasPrefix(respType, "multipart/mixed")) {
		return
	}

	// OK, we can compress the response:
	//base.LogTo("HTTP+", "GZip-compressing response")
	w.Header().Set("Content-Encoding", "gzip")
	w.Header().Del("Content-Length") // length is unknown due to compression

	w.gz = GetGZipWriter(w.ResponseWriter)
}

// Flushes the GZip encoder buffer, and if possible flushes output to the network.
func (w *EncodedResponseWriter) Flush() {
	if w.gz != nil {
		w.gz.Flush()
	}
	switch r := w.ResponseWriter.(type) {
	case http.Flusher:
		r.Flush()
	}
}

// The writer should be closed when output is complete, to flush the GZip encoder buffer.
func (w *EncodedResponseWriter) Close() {
	if w.gz != nil {
		ReturnGZipWriter(w.gz)
		w.gz = nil
	}
}

func (w *EncodedResponseWriter) CloseNotify() <-chan bool {
	var closeNotify <-chan bool
	cn, ok := w.ResponseWriter.(http.CloseNotifier)
	if ok {
		closeNotify = cn.CloseNotify()
	}
	return closeNotify
}

//////// GZIP WRITER CACHE:

var zipperCache sync.Pool

// Gets a gzip writer from the pool, or creates a new one if the pool is empty:
func GetGZipWriter(writer io.Writer) *gzip.Writer {
	if gz, ok := zipperCache.Get().(*gzip.Writer); ok {
		gz.Reset(writer)
		return gz
	} else {
		return gzip.NewWriter(writer)
	}
}

// Closes a gzip writer and returns it to the pool:
func ReturnGZipWriter(gz *gzip.Writer) {
	gz.Close()
	zipperCache.Put(gz)
}
<commit_msg>Do not use EncodedResponseWriter if the current HTTP request is asking for a WebSocket upgrade, WebSockets us per message compression not whole response compression. (#1945)<commit_after>package rest

import (
	"compress/gzip"
	"io"
	"net/http"
	"strings"
	"sync"

	"github.com/couchbase/sync_gateway/base"
)

// An implementation of http.ResponseWriter that wraps another instance and transparently applies
// GZip compression when appropriate.
type EncodedResponseWriter struct {
	http.ResponseWriter
	gz        *gzip.Writer
	status    int
	sniffDone bool
}

// Creates a new EncodedResponseWriter, or returns nil if the request doesn't allow encoded responses.
func NewEncodedResponseWriter(response http.ResponseWriter, rq *http.Request) *EncodedResponseWriter {
	isWebSocketRequest := strings.ToLower(rq.Header.Get("Upgrade")) == "websocket" &&
		strings.Contains(strings.ToLower(rq.Header.Get("Connection")), "upgrade")

	if isWebSocketRequest || !strings.Contains(rq.Header.Get("Accept-Encoding"), "gzip") ||
		rq.Method == "HEAD" || rq.Method == "PUT" || rq.Method == "DELETE" {
		return nil
	}

	// Workaround for https://github.com/couchbase/sync_gateway/issues/1419
	// if the user agent is empty or earlier than 1.2, then we never want to gzip
	// the entire response for requests to the _bulk_get endpoint, since the clients
	// are not equipped to handle that.
	if strings.Contains(rq.URL.Path, "_bulk_get") {
		userAgentVersion := NewUserAgentVersion(rq.Header.Get("User-Agent"))
		if userAgentVersion.IsBefore(1, 2) {
			return nil
		}
	}

	return &EncodedResponseWriter{ResponseWriter: response}
}

func (w *EncodedResponseWriter) WriteHeader(status int) {
	w.status = status
	w.sniff(nil) // Must do it now because headers can't be changed after WriteHeader call
	w.ResponseWriter.WriteHeader(status)
}

func (w *EncodedResponseWriter) Write(b []byte) (int, error) {
	w.sniff(b)
	if w.gz != nil {
		return w.gz.Write(b)
	} else {
		return w.ResponseWriter.Write(b)
	}
}

func (w *EncodedResponseWriter) disableCompression() {
	if w.sniffDone {
		base.Warn("EncodedResponseWriter: Too late to disableCompression!")
	}
	w.sniffDone = true
}

func (w *EncodedResponseWriter) sniff(bytes []byte) {
	if w.sniffDone {
		return
	}
	w.sniffDone = true
	// Check the content type, sniffing the initial data if necessary:
	respType := w.Header().Get("Content-Type")
	if respType == "" && bytes != nil {
		respType = http.DetectContentType(bytes)
		w.Header().Set("Content-Type", respType)
	}

	// Can/should we compress the response?
	if w.status >= 300 || w.Header().Get("Content-Encoding") != "" ||
		(!strings.HasPrefix(respType, "application/json") && !strings.HasPrefix(respType, "text/") && !strings.HasPrefix(respType, "multipart/mixed")) {
		return
	}

	// OK, we can compress the response:
	//base.LogTo("HTTP+", "GZip-compressing response")
	w.Header().Set("Content-Encoding", "gzip")
	w.Header().Del("Content-Length") // length is unknown due to compression

	w.gz = GetGZipWriter(w.ResponseWriter)
}

// Flushes the GZip encoder buffer, and if possible flushes output to the network.
func (w *EncodedResponseWriter) Flush() {
	if w.gz != nil {
		w.gz.Flush()
	}
	switch r := w.ResponseWriter.(type) {
	case http.Flusher:
		r.Flush()
	}
}

// The writer should be closed when output is complete, to flush the GZip encoder buffer.
func (w *EncodedResponseWriter) Close() {
	if w.gz != nil {
		ReturnGZipWriter(w.gz)
		w.gz = nil
	}
}

func (w *EncodedResponseWriter) CloseNotify() <-chan bool {
	var closeNotify <-chan bool
	cn, ok := w.ResponseWriter.(http.CloseNotifier)
	if ok {
		closeNotify = cn.CloseNotify()
	}
	return closeNotify
}

//////// GZIP WRITER CACHE:

var zipperCache sync.Pool

// Gets a gzip writer from the pool, or creates a new one if the pool is empty:
func GetGZipWriter(writer io.Writer) *gzip.Writer {
	if gz, ok := zipperCache.Get().(*gzip.Writer); ok {
		gz.Reset(writer)
		return gz
	} else {
		return gzip.NewWriter(writer)
	}
}

// Closes a gzip writer and returns it to the pool:
func ReturnGZipWriter(gz *gzip.Writer) {
	gz.Close()
	zipperCache.Put(gz)
}
<|endoftext|>


<commit_before>package trie

// The double array structure looks like this:
//
//          BASE                CHECK
// -+-------------------+-------------------
// 0| <unused>          | <unused>
// 1| free list pointer | free list pointer
// 2| 3                 | 0
// 3| ...               | ...
//
// There are no free slots initially, so the free list pointers are initialized to -1 (themselves). BASE(2) is
// the root node, and initially points to the next slot of the double array pool (i.e., beyond the end of the
// initially allocated double array).
const (
	daFreeListIndex = 1
	daRootIndex     = 2
)

// Translate a raw input byte to an index into the double array. The index cannot be 0. We're not using alpha
// maps as libdatrie does (for now), so the quick solution is to map 0 -> 256 and 1-255 map to themselves.
func byteToDAIndex(b byte) int {
	if b == '\x00' {
		return 256
	}
	return int(b)
}

type doubleArrayCell struct {
	base  int32
	check int32
}

type doubleArray struct {
	cells []doubleArrayCell
}

func newDoubleArray() *doubleArray {
	cells := []doubleArrayCell{
		doubleArrayCell{0, 0},
		doubleArrayCell{-1, -1},
		doubleArrayCell{daRootIndex + 1, 0},
	}
	return &doubleArray{cells}
}

// Get the value of BASE at some index i.
func (da *doubleArray) base(i int) int32 { return da.cells[i].base }

// Get the value of CHECK at some index i.
func (da *doubleArray) check(i int) int32 { return da.cells[i].check }

// Free all the cells in the range [start, end). Returns ok if successful.
func (da *doubleArray) free(start, end int) bool {
	if start <= daRootIndex || end > len(da.cells) {
		return false
	}
	/*previous := daFreeListIndex*/
	current := daFreeListIndex
	firstIteration := true
	// Walk the free list to find the nearest preceding free node.
	for current != daFreeListIndex && !firstIteration {
		next = -da.check(current)
		if next >= start {
			break
		}

	}

}

// A helper to walk to along a trie edge inside the double array. Returns the next index along the edge ch. ch
// must be converted using byteToDAIndex first; 0 is reserved for an end-of-string character here. If the DA
// cell at s points to the tail, then inTail will be true and next is the tailBlockList index. If the walk is
// not possible (either s is invalid or there is no such child) then ok is false. next and inTail should be
// ignored unless ok == true.
func (da *doubleArray) walk(s int, ch int) (next int, inTail, ok bool) {
	if s >= len(da.cells) || s < daRootIndex {
		return 0, false, false
	}
	sContents := da.base(s)
	if sContents < 0 {
		// A tail pointer
		return int(-sContents), true, true
	}
	next = int(sContents) + ch
	if next >= len(da.cells) || next < daRootIndex {
		return 0, false, false
	}
	if int(da.check(next)) == s {
		return next, false, true
	}
	return 0, false, false
}

// Adds a new transition from state s along an edge ch, possibly relocating an existing base or expanding the
// double array. Returns ok unless there was some problem and it failed.
func (da *doubleArray) addBase(s int, ch int) bool {
	target := da.check(s)
	if target >= len(da.cells) {
		// WIP
	}
}
<commit_msg>Write some more completely untested double array code.<commit_after>package trie

// The double array structure looks like this:
//
//          BASE                CHECK
// -+-------------------+-------------------
// 0| <unused>          | <unused>
// 1| free list pointer | free list pointer
// 2| 3                 | 0
// 3| ...               | ...
//
// There are no free slots initially, so the free list pointers are initialized to -1 (themselves). BASE(2) is
// the root node, and initially points to the next slot of the double array pool (i.e., beyond the end of the
// initially allocated double array).
const (
	daFreeListIndex = 1
	daRootIndex     = 2
)

// Translate a raw input byte to an index into the double array. The index cannot be 0. We're not using alpha
// maps as libdatrie does (for now), so the quick solution is to map 0 -> 256 and 1-255 map to themselves.
func byteToDAIndex(b byte) int {
	if b == '\x00' {
		return 256
	}
	return int(b)
}

type doubleArrayCell struct {
	base  int32
	check int32
}

type doubleArray struct {
	cells []doubleArrayCell
}

func newDoubleArray() *doubleArray {
	cells := []doubleArrayCell{
		doubleArrayCell{0, 0},
		doubleArrayCell{-1, -1},
		doubleArrayCell{daRootIndex + 1, 0},
	}
	return &doubleArray{cells}
}

// Get the value of BASE at some index i.
func (da *doubleArray) base(i int) int32 { return da.cells[i].base }

// Get the value of CHECK at some index i.
func (da *doubleArray) check(i int) int32 { return da.cells[i].check }

// Set the previous free pointer (-BASE) of a cell, so that the previous of s points to t.
func (da *doubleArray) setPrevious(s, t int) { da.cells[s].base = int32(-t) }

// Set the next free pointer (-CHECK) of a cell, so that the next of s points to t.
func (da *doubleArray) setNext(s, t int) { da.cells[s].check = int32(-t) }

// Remove a cell from the free list. Doesn't verify that the cell is free first, and doesn't modify the
// contents of s.
func (da *doubleArray) removeFromFreeList(s int) {
	nextFreeCell, previousFreeCell := int(-da.check(s)), int(-da.base(s))
	da.setNext(previousFreeCell, nextFreeCell)
	da.setPrevious(nextFreeCell, previousFreeCell)
}

// Free all the cells in the range [start, end). Returns ok if successful.
func (da *doubleArray) free(start, end int) bool {
	if start <= daRootIndex || end > len(da.cells) {
		return false
	}
	current := daFreeListIndex
	// Walk the free list to find the nearest preceding free node.
	var next int
	for {
		next = int(-da.check(current))
		if next >= start || next == daFreeListIndex {
			break
		}
		current = next
	}
	// Previous will be the preceding node; next will be the next one; we'll walk current up the range and set
	// the pointers as we go.
	previous := current
	for current = start; current < end; current++ {
		if n := da.check(current); n < 0 {
			// Node is already free.
			next = int(-n)
			// Previous already point to the current node, because this loop keeps the free list in a consistent
			// state.
			previous = current
		} else {
			// Node is currently occupied.
			da.setNext(previous, current)
			da.setPrevious(current, previous)
			da.setNext(current, next)
			da.setPrevious(next, current)
			previous = current
		}
	}
}

// Resize the double array to accomodate up to maxIndex (so the len(cells) == maxIndex + 1) and free the newly
// allocated cells. We could use append to take care of this but it will be less efficient to construct a new
// intermediate slice. Returns ok == true unless there was a problem.
func (da *doubleArray) resize(newSize int) bool {
	currentSize := len(da.cells)
	if newSize <= currentSize {
		return true
	}
	if newSize > cap(da.cells) {
		// Allocate a new slice of double the desired size, for room to grow.
		newDoubleArrayCells := make([]doubleArrayCell, newSize * 2)
		copy(newDoubleArray, da.cells)
		da.cells = newDoubleArrayCells
	}
	da.cells = da.cells[0:newSize]

	// Free the new cells.
	previous = int(-da.base(daFreeListIndex))
	var current int
	for current = currentSize; current < newSize; current++ {
		da.setNext(previous, current)
		da.setPrevious(current, previous)
	}
	// Fix the loop-around
	da.setNext(current, daRootIndex)
	da.setPrevious(daRootIndex, current)
	return true
}

// A helper to walk to along a trie edge inside the double array. Returns the next index along the edge ch. ch
// must be converted using byteToDAIndex first; 0 is reserved for an end-of-string character here. If the DA
// cell at s points to the tail, then inTail will be true and next is the tailBlockList index. If the walk is
// not possible (either s is invalid or there is no such child) then ok is false. next and inTail should be
// ignored unless ok == true.
func (da *doubleArray) walk(s int, ch int) (next int, inTail, ok bool) {
	if s >= len(da.cells) || s < daRootIndex {
		return 0, false, false
	}
	sContents := da.base(s)
	if sContents < 0 {
		// A tail pointer
		return int(-sContents), true, true
	}
	next = int(sContents) + ch
	if next >= len(da.cells) || next < daRootIndex {
		return 0, false, false
	}
	if int(da.check(next)) == s {
		return next, false, true
	}
	return 0, false, false
}

// Adds a new transition from state s along an edge ch, possibly relocating an existing base or expanding the
// double array. Returns ok == true unless there was some problem and it failed.
func (da *doubleArray) addBase(s int, ch int) bool {
	base := da.check(s)

	if base < 0 {
		// Cell is free. Use the first free cell as the target.
		target := int(-da.check(daFreeListIndex))
		if target == daFreeListIndex {
			// No free cells.
			if ok := da.resize(len(da.cells) +1); !ok {
				panic("Error resizing double array.")
			}
			target = int(-da.check(daFreeListIndex))
		}
		// Remove the cell from the free list.
		da.removeFromFreeList(target)
		// Use this cell as the base for our new edge.
		da.cells[target].check = int32(s)
		da.cells[target].base = 0 // TODO: Is this correct?
		// Now call ourselves recursively. The second time around we won't hit this if statement (we'll use the
		// base we just created).
		return da.addBase(s, ch)
	}

	// Cell is already in use.
	t := int(-base) + ch
	if da.check(t) == s {
		// The transition already exists.
		return true
	}
	if t >= len(da.cells) {
		// The transition overflows the double array. Add more cells.
		if ok := da.resize(t+1); !ok {
			panic("Error resizing double array.")
		}
	} else {
		// The transition conflicts with an existing transition. Relocate that base out of the way.
		panic("Base relocation unimplemented.")
	}

	// t is now free and ready to be used. Remove it from the free list and update its contents.
	da.removeFromFreeList(t)
	da.cells[t].check = s
	da.cells[t].base = 0 // TODO: Correct?
	// WIP
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var (
	command  string
	timeout  uint64
	username string
	hosts    []string
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s command timeout user hosts...\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "timeout should be in seconds (can be a float)\n")
	os.Exit(1)
}

func parseArgs() {
	if len(os.Args) < 5 {
		usage()
	}
	command = os.Args[1]
	timeoutf, err := strconv.ParseFloat(os.Args[2], 32)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error while parsing timeout: %v\n", err)
		usage()
	}
	timeout = uint64(timeoutf * 1000) // convert to ms
	username = os.Args[3]
	hosts = os.Args[4:]
}

type Req struct {
	Action  string
	Cmd     string
	Timeout uint64
	Hosts   []string
}

type InitializeComplete struct {
	InitializeComplete bool
}

type ConnectionProgress struct {
	ConnectedHost string
}

type UserError struct {
	IsCritical bool
	ErrorMsg   string
}

type Reply struct {
	Hostname string
	Stdout   string
	Stderr   string
	Success  bool
	ErrMsg   string
}

type FinalReply struct {
	TotalTime     float64
	TimedOutHosts map[string]bool
}

type Resp struct {
	Type               string
	InitializeComplete `json:",omitempty,inline"`
	ConnectionProgress `json:",omitempty,inline"`
	UserError          `json:",omitempty,inline"`
	Reply              `json:",omitempty,inline"`
	FinalReply         `json:",omitempty,inline"`
}

func main() {
	parseArgs()
	cmd := exec.Command("GoSSHa", "-l", username)
	send, err := cmd.StdinPipe()
	if err != nil {
		die("failed to create stdin pipe: %v\n", err)
	}
	recv, err := cmd.StdoutPipe()
	if err != nil {
		die("failed to create stdout pipe: %v\n", err)
	}
	if err := cmd.Start(); err != nil {
		die("failed to start parallel ssh agent: %v\n", err)
	}
	scanner := bufio.NewScanner(recv)
	if !scanner.Scan() {
		die("error occurred while initializing parallel ssh agent: %v\n", scanner.Err())
	}
	var resp Resp
	if err := json.Unmarshal(scanner.Bytes(), &resp); err != nil {
		die("error decoding response: %v\n", err)
	}
	if resp.Type != "InitializeComplete" || !resp.InitializeComplete.InitializeComplete {
		die("error occurred while initilizing parallel ssh agent\n")
	}
	enc := json.NewEncoder(send)
	req := Req{
		Action:  "ssh",
		Cmd:     command,
		Timeout: timeout,
		Hosts:   hosts,
	}
	if err := enc.Encode(&req); err != nil {
		die("error occurred while encoding request: %v\n", err)
	}
	cleanExit := true
	for scanner.Scan() {
		var resp Resp
		if err := json.Unmarshal(scanner.Bytes(), &resp); err != nil {
			fmt.Fprintf(os.Stderr, "error occured decoding response from parallel ssh agent: %v\n", err)
			cleanExit = false
			continue
		}
		switch resp.Type {
		case "ConnectionProgress":
			fmt.Printf("Connected to %s\n", resp.ConnectionProgress.ConnectedHost)
		case "UserError":
			if resp.UserError.IsCritical {
				fmt.Fprintf(os.Stderr, "Error: %s\n", resp.UserError.ErrorMsg)
				cmd.Process.Kill()
				cleanExit = false
			} else {
				fmt.Fprintf(os.Stderr, "Warning: %s\n", resp.UserError.ErrorMsg)
			}
		case "Reply":
			fmt.Printf("==== Response from %s ====\n", resp.Reply.Hostname)
			fmt.Printf(" - Success: %t\n", resp.Reply.Success)
			if resp.Reply.ErrMsg != "" {
				fmt.Printf(" - ErrMsg: %s\n", resp.Reply.ErrMsg)
			}
			if resp.Reply.Stdout != "" {
				fmt.Println(" - Stdout:")
				fmt.Println(resp.Reply.Stdout)
			}
			if resp.Reply.Stderr != "" {
				fmt.Println(" - Stderr:")
				fmt.Println(resp.Reply.Stderr)
			}
		case "FinalReply":
			fmt.Println("==== Summary ====")
			fmt.Printf("- Execution Time: %f\n", resp.FinalReply.TotalTime)
			var timedout []string
			for h, t := range resp.FinalReply.TimedOutHosts {
				if t {
					timedout = append(timedout, h)
				}
			}
			if len(timedout) > 0 {
				fmt.Printf("- Timed out: %s\n", strings.Join(timedout, " "))
			}
			cmd.Process.Kill()
		default:
			fmt.Printf("Unexpected response: %s\n", scanner.Text())
			cmd.Process.Kill()
		}
	}
	if !cleanExit {
		os.Exit(2)
	}
}

func die(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format, args)
	os.Exit(1)
}
<commit_msg>make this a little more usable<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/spf13/pflag"
)

var (
	command    []string
	timeoutDur time.Duration
	timeoutMS  uint64
	username   string
	hosts      []string
)

func init() {
	pflag.DurationVarP(&timeoutDur, "timeout", "t", 10*time.Second, "timeout for the command to complete")
	pflag.StringVarP(&username, "user", "u", "root", "user to login with")
	pflag.StringSliceVarP(&hosts, "hosts", "h", nil, "hosts to run command on")
}

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [options] -- command\n", os.Args[0])
	pflag.PrintDefaults()
}

func parseArgs() {
	pflag.Usage = usage
	pflag.Parse()
	if pflag.NArg() < 1 {
		log.Fatal("must specify a command")
	}
	command = pflag.Args()
	timeoutMS = uint64(timeoutDur / time.Millisecond)
}

type Req struct {
	Action  string
	Cmd     string
	Timeout uint64
	Hosts   []string
}

type InitializeComplete struct {
	InitializeComplete bool
}

type ConnectionProgress struct {
	ConnectedHost string
}

type UserError struct {
	IsCritical bool
	ErrorMsg   string
}

type Reply struct {
	Hostname string
	Stdout   string
	Stderr   string
	Success  bool
	ErrMsg   string
}

type FinalReply struct {
	TotalTime     float64
	TimedOutHosts map[string]bool
}

type Resp struct {
	Type               string
	InitializeComplete `json:",omitempty,inline"`
	ConnectionProgress `json:",omitempty,inline"`
	UserError          `json:",omitempty,inline"`
	Reply              `json:",omitempty,inline"`
	FinalReply         `json:",omitempty,inline"`
}

func run() error {
	cmd := exec.Command("GoSSHa", "-l", username)
	send, err := cmd.StdinPipe()
	if err != nil {
		return errors.Wrap(err, "failed to create stdin pipe")
	}
	recv, err := cmd.StdoutPipe()
	if err != nil {
		return errors.Wrap(err, "failed to create stdout pipe")
	}
	if err := cmd.Start(); err != nil {
		return errors.Wrap(err, "failed to start parallel ssh agent")
	}
	defer cmd.Process.Kill()
	scanner := bufio.NewScanner(recv)
	if !scanner.Scan() {
		return errors.Wrap(scanner.Err(), "error occurred during parallel ssh agent init")
	}
	var resp Resp
	if err := json.Unmarshal(scanner.Bytes(), &resp); err != nil {
		return errors.Wrap(err, "error decoding response")
	}
	if resp.Type != "InitializeComplete" || !resp.InitializeComplete.InitializeComplete {
		return errors.New("error occurred during parallel ssh agent init")
	}
	enc := json.NewEncoder(send)
	req := Req{
		Action:  "ssh",
		Cmd:     strings.Join(command, " "),
		Timeout: timeoutMS,
		Hosts:   hosts,
	}
	if err := enc.Encode(&req); err != nil {
		return errors.Wrap(err, "error occurred while encoding request: %v\n")
	}
	var failedHosts []string
	for scanner.Scan() {
		var resp Resp
		if err := json.Unmarshal(scanner.Bytes(), &resp); err != nil {
			errors.Wrap(err, "error occured decoding response from parallel ssh agent")
			continue
		}
		switch resp.Type {
		case "ConnectionProgress":
			log.Printf("[%s] connected\n", resp.ConnectionProgress.ConnectedHost)
		case "UserError":
			if resp.UserError.IsCritical {
				return errors.New("critical error: " + resp.UserError.ErrorMsg)
			} else {
				log.Printf("warning: %s\n", resp.UserError.ErrorMsg)
			}
		case "Reply":
			status := "success"
			if !resp.Reply.Success {
				status = "failure"
				failedHosts = append(failedHosts, resp.Reply.Hostname)
			}
			log.Printf("[%s] %s", resp.Reply.Hostname, status)
			if resp.Reply.ErrMsg != "" {
				log.Printf("[%s] error message: %s", resp.Reply.Hostname, resp.Reply.ErrMsg)
			}
			if resp.Reply.Stdout != "" {
				log.Printf("[%s] stdout:\n%s", resp.Reply.Stdout)
			}
			if resp.Reply.Stderr != "" {
				log.Printf("[%s] stderr:\n%s", resp.Reply.Stderr)
			}
		case "FinalReply":
			fmt.Println("summary:")
			fmt.Printf("\texecution time: %f s\n", resp.FinalReply.TotalTime)
			var timedout []string
			for h, t := range resp.FinalReply.TimedOutHosts {
				if t {
					timedout = append(timedout, h)
				}
			}
			if len(timedout) > 0 {
				fmt.Printf("\ttimed out: %s\n", strings.Join(timedout, " "))
				return errors.New("some hosts timed out")
			}
			if len(failedHosts) > 0 {
				return errors.Errorf("failed hosts: %s", strings.Join(failedHosts, " "))
			}
			return nil
		default:
			return errors.New("unexpected response: " + scanner.Text())
		}
	}
	return nil
}

func main() {
	log.SetPrefix("psshcmd: ")
	log.SetFlags(0)
	parseArgs()
	if err := run(); err != nil {
		errors.Fprint(os.Stderr, err)
		os.Exit(2)
	}
}
<|endoftext|>


<commit_before>package main
import (
	"github.com/gorilla/mux"
	"github.com/tylerb/graceful"
	"net/http"
	"fmt"
	"os"
	"github.com/mailgun/oxy/forward"
	"github.com/mailgun/oxy/roundrobin"
	"github.com/mailgun/oxy/testutils"
	"time"
	"net"
	"os/signal"
	"syscall"
	"github.com/BurntSushi/toml"
)

type Backend struct {
	Servers []string
}

type Server struct {
	Url string
}

type Rule struct {
	Category string
	Value string
}

type Route struct {
	Backends []string
	Rules map[string]Rule
}

type Config struct {
	Backends map[string]Backend
	Servers map[string]Server
	Routes map[string]Route
}

var srv *graceful.Server

func main() {

	var config Config
	if _, err := toml.DecodeFile("tortuous.toml", &config); err != nil {
		fmt.Println(err)
		return
	}
	fmt.Printf("%+v\n", config )

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	systemRouter := mux.NewRouter()
	systemRouter.Methods("POST").Path("/").HandlerFunc(ReloadHandler)
	systemRouter.Methods("GET").Path("/").HandlerFunc(GetPidHandler)
	go http.ListenAndServe(":8000", systemRouter)

	userRouter := mux.NewRouter()

	/*for i := range config.Routes {
		fmt.Printf("%+v\n", config.Routes[i] )
	}*/

	fwd, _ := forward.New()
	lb, _ := roundrobin.New(fwd)

	lb.UpsertServer(testutils.ParseURI("http://172.17.0.2:80"))
	lb.UpsertServer(testutils.ParseURI("http://172.17.0.3:80"))

	userRouter.Host("test.zenika.fr").Handler(lb)

	goAway := false
	go func() {
		sig := <-sigs
		fmt.Println("I have to go...", sig)
		goAway = true
		srv.Stop(10 * time.Second)
	}()

	for{
		if (goAway){
			break
		}
		fmt.Println("Started")
		srv = &graceful.Server{
			Timeout: 10 * time.Second,
			NoSignalHandling: true,

			ConnState: func(conn net.Conn, state http.ConnState) {
				fmt.Println( "Connection ", state)
			},

			Server: &http.Server{
				Addr: ":8001",
				Handler: userRouter,
			},
		}

		go srv.ListenAndServe()
		<- srv.StopChan()
		fmt.Println("Stopped")
	}
}

func ReloadHandler(rw http.ResponseWriter, r *http.Request) {
	srv.Stop(10 * time.Second)
}

func GetPidHandler(rw http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(rw, "%d", os.Getpid())
}
<commit_msg>Config loaded and working<commit_after>package main
import (
	"github.com/gorilla/mux"
	"github.com/tylerb/graceful"
	"net/http"
	"fmt"
	"os"
	"github.com/mailgun/oxy/forward"
	"github.com/mailgun/oxy/roundrobin"
	"github.com/mailgun/oxy/testutils"
	"time"
	"net"
	"os/signal"
	"syscall"
	"github.com/BurntSushi/toml"
	"encoding/json"
)

type Backend struct {
	Servers []string
}

type Server struct {
	Url string
}

type Rule struct {
	Category string
	Value string
}

type Route struct {
	Backends []string
	Rules map[string]Rule
}

type Config struct {
	Backends map[string]Backend
	Servers map[string]Server
	Routes map[string]Route
}

var srv *graceful.Server
var userRouter *mux.Router
var config = new(Config)

func main() {

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	systemRouter := mux.NewRouter()
	systemRouter.Methods("POST").Path("/restart").HandlerFunc(RestartHandler)
	systemRouter.Methods("POST").Path("/reload").HandlerFunc(ReloadConfigHandler)
	systemRouter.Methods("GET").Path("/").HandlerFunc(GetConfigHandler)
	go http.ListenAndServe(":8000", systemRouter)

	userRouter = LoadConfig()

	goAway := false
	go func() {
		sig := <-sigs
		fmt.Println("I have to go...", sig)
		goAway = true
		srv.Stop(10 * time.Second)
	}()

	for{
		if (goAway){
			break
		}
		fmt.Println("Started")
		srv = &graceful.Server{
			Timeout: 10 * time.Second,
			NoSignalHandling: true,

			ConnState: func(conn net.Conn, state http.ConnState) {
				fmt.Println( "Connection ", state)
			},

			Server: &http.Server{
				Addr: ":8001",
				Handler: userRouter,
			},
		}

		go srv.ListenAndServe()
		<- srv.StopChan()
		fmt.Println("Stopped")
	}
}

func LoadConfig() *mux.Router{
	if metadata, err := toml.DecodeFile("tortuous.toml", config); err != nil {
		fmt.Println(err)
		return nil
	}else{
		fmt.Printf("Loaded config: %+v\n", metadata )
	}

	router := mux.NewRouter()
	for routeName, route := range config.Routes {
		fmt.Println("Creating route", routeName)
		fwd, _ := forward.New()
		newRoute:= router.NewRoute()
		for ruleName, rule := range route.Rules{
			fmt.Println("Creating rule", ruleName)
			newRoute = newRoute.Host(rule.Value)
		}
		for _, backendName := range route.Backends {
			fmt.Println("Creating backend", backendName)
			lb, _ := roundrobin.New(fwd)
			for _, serverName := range config.Backends[backendName].Servers {
				fmt.Println("Creating server", serverName)
				lb.UpsertServer(testutils.ParseURI(config.Servers[serverName].Url))
			}
			newRoute.Handler(lb)
		}
	}
	return router
}

func ReloadConfigHandler(rw http.ResponseWriter, r *http.Request) {
	userRouter = LoadConfig()
}

func RestartHandler(rw http.ResponseWriter, r *http.Request) {
	srv.Stop(10 * time.Second)
}

func GetConfigHandler(rw http.ResponseWriter, r *http.Request) {
	if jsonRes, err := json.Marshal(config); err != nil {
		fmt.Println(err)
		return
	}else{
		fmt.Fprintf(rw, "%s", jsonRes)
	}
}
<|endoftext|>