<commit_before>/*
Copyright 2021 The Vitess Authors.

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

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

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

package consul

import (
	"context"
	"flag"
	"fmt"
	"os"
	"testing"
	"time"

	"vitess.io/vitess/go/vt/log"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance *cluster.LocalProcessCluster
	cell            = "zone1"
	hostname        = "localhost"
	KeyspaceName    = "customer"
	SchemaSQL       = `
CREATE TABLE t1 (
    c1 BIGINT NOT NULL,
    c2 BIGINT NOT NULL,
    c3 BIGINT,
    c4 varchar(100),
    PRIMARY KEY (c1),
    UNIQUE KEY (c2),
    UNIQUE KEY (c3),
    UNIQUE KEY (c4)
) ENGINE=Innodb;`
	VSchema = `
{
    "sharded": false,
    "tables": {
        "t1": {}
    }
}
`
)

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitCode := func() int {
		clusterInstance = cluster.NewCluster(cell, hostname)
		defer clusterInstance.Teardown()

		// Start topo server
		clusterInstance.TopoFlavor = "consul"
		if err := clusterInstance.StartTopo(); err != nil {
			return 1
		}

		// Start keyspace
		Keyspace := &cluster.Keyspace{
			Name:      KeyspaceName,
			SchemaSQL: SchemaSQL,
			VSchema:   VSchema,
		}
		if err := clusterInstance.StartUnshardedKeyspace(*Keyspace, 0, false); err != nil {
			log.Fatal(err.Error())
			return 1
		}

		// Start vtgate
		if err := clusterInstance.StartVtgate(); err != nil {
			log.Fatal(err.Error())
			return 1
		}

		return m.Run()
	}()
	os.Exit(exitCode)
}

func TestTopoRestart(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	vtParams := mysql.ConnParams{
		Host: "localhost",
		Port: clusterInstance.VtgateMySQLPort,
	}
	conn, err := mysql.Connect(ctx, &vtParams)
	require.Nil(t, err)
	defer conn.Close()

	execMulti(t, conn, `insert into t1(c1, c2, c3, c4) values (300,100,300,'abc'); ;; insert into t1(c1, c2, c3, c4) values (301,101,301,'abcd');;`)
	assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)] [INT64(301) INT64(101) INT64(301)]]`)

	defer execute(t, conn, `delete from t1`)

	ch := make(chan interface{})

	go func() {
		clusterInstance.TopoProcess.TearDown(clusterInstance.Cell, clusterInstance.OriginalVTDATAROOT, clusterInstance.CurrentVTDATAROOT, true, *clusterInstance.TopoFlavorString())

		// Some sleep to server few queries when topo is down.
		time.Sleep(400 * time.Millisecond)

		clusterInstance.TopoProcess.Setup(*clusterInstance.TopoFlavorString(), clusterInstance)

		// topo is up now.
		ch <- 1
	}()

	timeOut := time.After(5 * time.Second)

	for {
		select {
		case <-ch:
			return
		case <-timeOut:
			require.Fail(t, "timed out - topo process did not come up")
		case <-time.After(100 * time.Millisecond):
			assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)] [INT64(301) INT64(101) INT64(301)]]`)
		}
	}
}

func execute(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
	t.Helper()
	qr, err := conn.ExecuteFetch(query, 1000, true)
	require.NoError(t, err)
	return qr
}

func execMulti(t *testing.T, conn *mysql.Conn, query string) []*sqltypes.Result {
	t.Helper()
	var res []*sqltypes.Result
	qr, more, err := conn.ExecuteFetchMulti(query, 1000, true)
	res = append(res, qr)
	require.NoError(t, err)
	for more == true {
		qr, more, _, err = conn.ReadQueryResult(1000, true)
		require.NoError(t, err)
		res = append(res, qr)
	}
	return res
}

func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
	t.Helper()
	qr := execute(t, conn, query)
	got := fmt.Sprintf("%v", qr.Rows)
	diff := cmp.Diff(expected, got)
	if diff != "" {
		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
	}
}
<commit_msg>increase timeout<commit_after>/*
Copyright 2021 The Vitess Authors.

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

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

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

package consul

import (
	"context"
	"flag"
	"fmt"
	"os"
	"testing"
	"time"

	"vitess.io/vitess/go/vt/log"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance *cluster.LocalProcessCluster
	cell            = "zone1"
	hostname        = "localhost"
	KeyspaceName    = "customer"
	SchemaSQL       = `
CREATE TABLE t1 (
    c1 BIGINT NOT NULL,
    c2 BIGINT NOT NULL,
    c3 BIGINT,
    c4 varchar(100),
    PRIMARY KEY (c1),
    UNIQUE KEY (c2),
    UNIQUE KEY (c3),
    UNIQUE KEY (c4)
) ENGINE=Innodb;`
	VSchema = `
{
    "sharded": false,
    "tables": {
        "t1": {}
    }
}
`
)

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitCode := func() int {
		clusterInstance = cluster.NewCluster(cell, hostname)
		defer clusterInstance.Teardown()

		// Start topo server
		clusterInstance.TopoFlavor = "consul"
		if err := clusterInstance.StartTopo(); err != nil {
			return 1
		}

		// Start keyspace
		Keyspace := &cluster.Keyspace{
			Name:      KeyspaceName,
			SchemaSQL: SchemaSQL,
			VSchema:   VSchema,
		}
		if err := clusterInstance.StartUnshardedKeyspace(*Keyspace, 0, false); err != nil {
			log.Fatal(err.Error())
			return 1
		}

		// Start vtgate
		if err := clusterInstance.StartVtgate(); err != nil {
			log.Fatal(err.Error())
			return 1
		}

		return m.Run()
	}()
	os.Exit(exitCode)
}

func TestTopoRestart(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	vtParams := mysql.ConnParams{
		Host: "localhost",
		Port: clusterInstance.VtgateMySQLPort,
	}
	conn, err := mysql.Connect(ctx, &vtParams)
	require.Nil(t, err)
	defer conn.Close()

	execMulti(t, conn, `insert into t1(c1, c2, c3, c4) values (300,100,300,'abc'); ;; insert into t1(c1, c2, c3, c4) values (301,101,301,'abcd');;`)
	assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)] [INT64(301) INT64(101) INT64(301)]]`)

	defer execute(t, conn, `delete from t1`)

	ch := make(chan interface{})

	go func() {
		clusterInstance.TopoProcess.TearDown(clusterInstance.Cell, clusterInstance.OriginalVTDATAROOT, clusterInstance.CurrentVTDATAROOT, true, *clusterInstance.TopoFlavorString())

		// Some sleep to server few queries when topo is down.
		time.Sleep(400 * time.Millisecond)

		clusterInstance.TopoProcess.Setup(*clusterInstance.TopoFlavorString(), clusterInstance)

		// topo is up now.
		ch <- 1
	}()

	timeOut := time.After(15 * time.Second)

	for {
		select {
		case <-ch:
			return
		case <-timeOut:
			require.Fail(t, "timed out - topo process did not come up")
		case <-time.After(100 * time.Millisecond):
			assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)] [INT64(301) INT64(101) INT64(301)]]`)
		}
	}
}

func execute(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
	t.Helper()
	qr, err := conn.ExecuteFetch(query, 1000, true)
	require.NoError(t, err)
	return qr
}

func execMulti(t *testing.T, conn *mysql.Conn, query string) []*sqltypes.Result {
	t.Helper()
	var res []*sqltypes.Result
	qr, more, err := conn.ExecuteFetchMulti(query, 1000, true)
	res = append(res, qr)
	require.NoError(t, err)
	for more == true {
		qr, more, _, err = conn.ReadQueryResult(1000, true)
		require.NoError(t, err)
		res = append(res, qr)
	}
	return res
}

func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
	t.Helper()
	qr := execute(t, conn, query)
	got := fmt.Sprintf("%v", qr.Rows)
	diff := cmp.Diff(expected, got)
	if diff != "" {
		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
	}
}
<|endoftext|>


<commit_before>package merge

import (
	"runtime"
	"sort"
	"sync"
)

func sortBucket(comparators Comparators) {
	sort.Sort(comparators)
}

func copyChunk(chunk []Comparators) []Comparators {
	cp := make([]Comparators, len(chunk))
	copy(cp, chunk)
	return cp
}

// MultithreadedSortComparators will take a list of comparators
// and sort it using as many threads as are available.  The list
// is split into buckets for a bucket sort and then recursively
// merged using SymMerge.
func MultithreadedSortComparators(comparators Comparators) Comparators {
	toBeSorted := make(Comparators, len(comparators))
	copy(toBeSorted, comparators)

	var wg sync.WaitGroup
	chunks := chunk(toBeSorted, int64(runtime.NumCPU()))
	wg.Add(len(chunks))
	for i := 0; i < len(chunks); i++ {
		go func(i int) {
			sortBucket(chunks[i])
			wg.Done()
		}(i)
	}

	wg.Wait()
	todo := make([]Comparators, len(chunks)/2)
	for {
		todo = todo[:len(chunks)/2]
		wg.Add(len(chunks) / 2)
		for i := 0; i < len(chunks); i += 2 {
			go func(i int) {
				todo[i/2] = SymMerge(chunks[i], chunks[i+1])
				wg.Done()
			}(i)
		}

		wg.Wait()

		chunks = copyChunk(todo)
		if len(chunks) == 1 {
			break
		}
	}

	return chunks[0]
}

func chunk(comparators Comparators, numParts int64) []Comparators {
	parts := make([]Comparators, numParts)
	for i := int64(0); i < numParts; i++ {
		parts[i] = comparators[i*int64(len(comparators))/numParts : (i+1)*int64(len(comparators))/numParts]
	}
	return parts
}
<commit_msg>Fixed potential issue on single core machine.<commit_after>package merge

import (
	"runtime"
	"sort"
	"sync"
)

func sortBucket(comparators Comparators) {
	sort.Sort(comparators)
}

func copyChunk(chunk []Comparators) []Comparators {
	cp := make([]Comparators, len(chunk))
	copy(cp, chunk)
	return cp
}

// MultithreadedSortComparators will take a list of comparators
// and sort it using as many threads as are available.  The list
// is split into buckets for a bucket sort and then recursively
// merged using SymMerge.
func MultithreadedSortComparators(comparators Comparators) Comparators {
	toBeSorted := make(Comparators, len(comparators))
	copy(toBeSorted, comparators)

	var wg sync.WaitGroup

	numCPU := int64(runtime.NumCPU())
	if numCPU%2 == 1 { // single core machine
		numCPU++
	}

	chunks := chunk(toBeSorted, numCPU)
	wg.Add(len(chunks))
	for i := 0; i < len(chunks); i++ {
		go func(i int) {
			sortBucket(chunks[i])
			wg.Done()
		}(i)
	}

	wg.Wait()
	todo := make([]Comparators, len(chunks)/2)
	for {
		todo = todo[:len(chunks)/2]
		wg.Add(len(chunks) / 2)
		for i := 0; i < len(chunks); i += 2 {
			go func(i int) {
				todo[i/2] = SymMerge(chunks[i], chunks[i+1])
				wg.Done()
			}(i)
		}

		wg.Wait()

		chunks = copyChunk(todo)
		if len(chunks) == 1 {
			break
		}
	}

	return chunks[0]
}

func chunk(comparators Comparators, numParts int64) []Comparators {
	parts := make([]Comparators, numParts)
	for i := int64(0); i < numParts; i++ {
		parts[i] = comparators[i*int64(len(comparators))/numParts : (i+1)*int64(len(comparators))/numParts]
	}
	return parts
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"
)

// Should serve a known static error page if all backend servers are down
// and object isn't in cache/stale.
// NB: ideally this should be a page that we control that has a mechanism
//     to alert us that it has been served.
func TestFailoverErrorPageAllServersDown(t *testing.T) {
	ResetBackends(backendsByPriority)

	const expectedStatusCode = http.StatusServiceUnavailable
	const expectedBody = "Guru Meditation"

	originServer.Stop()
	backupServer1.Stop()
	backupServer2.Stop()

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatusCode {
		t.Errorf(
			"Invalid StatusCode received. Expected %d, got %d",
			expectedStatusCode,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	if bodyStr := string(body); !strings.Contains(bodyStr, expectedBody) {
		t.Errorf(
			"Received incorrect response body. Expected to contain %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should return the 5xx response from the last backup server if all
// preceeding servers also return a 5xx response.
func TestFailoverErrorPageAllServers5xx(t *testing.T) {
	ResetBackends(backendsByPriority)

	const expectedStatusCode = http.StatusServiceUnavailable
	const expectedBody = "lucky golden ticket"

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(originServer.Name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(backupServer1.Name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(expectedBody))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatusCode {
		t.Errorf(
			"Invalid StatusCode received. Expected %d, got %d",
			expectedStatusCode,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should back off requests against origin for a very short period of time
// if origin returns a 5xx response so as not to overwhelm it.
func TestFailoverOrigin5xxBackOff(t *testing.T) {
	t.Skip("Not implemented")
}

// Should serve stale object and not hit mirror(s) if origin is down and
// object is beyond TTL but still in cache.
func TestFailoverOriginDownServeStale(t *testing.T) {
	t.Skip("Not implemented")
}

// Should serve stale object and not hit mirror(s) if origin returns a 5xx
// response and object is beyond TTL but still in cache.
func TestFailoverOrigin5xxServeStale(t *testing.T) {
	ResetBackends(backendsByPriority)

	const expectedResponseStale = "going off like stilton"
	const expectedResponseFresh = "as fresh as daisies"

	const respTTL = time.Duration(2 * time.Second)
	const respTTLWithBuffer = 5 * respTTL
	// Allow varnish's beresp.saintmode to expire.
	const waitSaintMode = time.Duration(5 * time.Second)
	headerValue := fmt.Sprintf("max-age=%.0f", respTTL.Seconds())

	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	req := NewUniqueEdgeGET(t)

	var expectedBody string
	for requestCount := 1; requestCount < 6; requestCount++ {
		switch requestCount {
		case 1: // Request 1 populates cache.
			expectedBody = expectedResponseStale

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Cache-Control", headerValue)
				w.Write([]byte(expectedBody))
			})
		case 2: // Requests 2,3,4 come from stale.
			time.Sleep(respTTLWithBuffer)
			expectedBody = expectedResponseStale

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusServiceUnavailable)
				w.Write([]byte(originServer.Name))
			})
		case 5: // Last request comes directly from origin again.
			time.Sleep(waitSaintMode)
			expectedBody = expectedResponseFresh

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(expectedBody))
			})
		}

		resp := RoundTripCheckError(t, req)

		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Fatal(err)
		}
		if bodyStr := string(body); bodyStr != expectedBody {
			t.Errorf(
				"Request %d received incorrect response body. Expected %q, got %q",
				requestCount,
				expectedBody,
				bodyStr,
			)
		}
	}
}

// Should fallback to first mirror if origin is down and object is not in
// cache (active or stale).
func TestFailoverOriginDownUseFirstMirror(t *testing.T) {
	ResetBackends(backendsByPriority)

	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK

	originServer.Stop()
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(expectedBody))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received a request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should fallback to first mirror if origin returns 5xx response and object
// is not in cache (active or stale).
func TestFailoverOrigin5xxUseFirstMirror(t *testing.T) {
	ResetBackends(backendsByPriority)

	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK
	backendsSawRequest := map[string]bool{}

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := originServer.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		if !backendsSawRequest[name] {
			w.Write([]byte(expectedBody))
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
			w.Write([]byte(name))
		}
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received a request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should fallback to second mirror if both origin and first mirror are
// down.
func TestFailoverOriginDownFirstMirrorDownUseSecondMirror(t *testing.T) {
	ResetBackends(backendsByPriority)

	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK

	originServer.Stop()
	backupServer1.Stop()
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(expectedBody))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should fallback to second mirror if both origin and first mirror return
// 5xx responses.
func TestFailoverOrigin5xxFirstMirror5xxUseSecondMirror(t *testing.T) {
	ResetBackends(backendsByPriority)

	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK
	backendsSawRequest := map[string]bool{}

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := originServer.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		if !backendsSawRequest[name] {
			w.Write([]byte(expectedBody))
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
			w.Write([]byte(name))
		}
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should not fallback to mirror if origin returns a 5xx response with a
// No-Fallback header. In order to allow applications to present their own
// error pages.
func TestFailoverNoFallbackHeader(t *testing.T) {
	ResetBackends(backendsByPriority)

	const headerName = "No-Fallback"
	const expectedStatus = http.StatusServiceUnavailable
	const expectedBody = "custom error page"

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set(headerName, "")
		w.WriteHeader(expectedStatus)
		w.Write([]byte(expectedBody))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}
<commit_msg>Split origin down serve stale test<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"
)

// Should serve a known static error page if all backend servers are down
// and object isn't in cache/stale.
// NB: ideally this should be a page that we control that has a mechanism
//     to alert us that it has been served.
func TestFailoverErrorPageAllServersDown(t *testing.T) {
	ResetBackends(backendsByPriority)

	const expectedStatusCode = http.StatusServiceUnavailable
	const expectedBody = "Guru Meditation"

	originServer.Stop()
	backupServer1.Stop()
	backupServer2.Stop()

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatusCode {
		t.Errorf(
			"Invalid StatusCode received. Expected %d, got %d",
			expectedStatusCode,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	if bodyStr := string(body); !strings.Contains(bodyStr, expectedBody) {
		t.Errorf(
			"Received incorrect response body. Expected to contain %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should return the 5xx response from the last backup server if all
// preceeding servers also return a 5xx response.
func TestFailoverErrorPageAllServers5xx(t *testing.T) {
	ResetBackends(backendsByPriority)

	const expectedStatusCode = http.StatusServiceUnavailable
	const expectedBody = "lucky golden ticket"

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(originServer.Name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(backupServer1.Name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(expectedBody))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatusCode {
		t.Errorf(
			"Invalid StatusCode received. Expected %d, got %d",
			expectedStatusCode,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should back off requests against origin for a very short period of time
// if origin returns a 5xx response so as not to overwhelm it.
func TestFailoverOrigin5xxBackOff(t *testing.T) {
	t.Skip("Not implemented")
}

// Should serve reponse from first mirror and replace stale object if origin
// is down and health check has *not* expired.
// FIXME: This is not desired behaviour. We should serve from stale
//        immediately and not replace the stale object in cache.
func TestFailoverOriginDownHealthCheckNotExpiredReplaceStale(t *testing.T) {
	t.Skip("Not implemented")
}

// Should serve stale object and not hit mirror(s) if origin is down, health
// check has expired, and object is beyond TTL but still in cache.
// FIXME: This is not quite desired behaviour. We should not have to wait
//				for the stale object to become available.
func TestFailoverOriginDownHealthCheckHasExpiredServeStale(t *testing.T) {
	t.Skip("Not implemented")
}

// Should serve stale object and not hit mirror(s) if origin returns a 5xx
// response and object is beyond TTL but still in cache.
func TestFailoverOrigin5xxServeStale(t *testing.T) {
	ResetBackends(backendsByPriority)

	const expectedResponseStale = "going off like stilton"
	const expectedResponseFresh = "as fresh as daisies"

	const respTTL = time.Duration(2 * time.Second)
	const respTTLWithBuffer = 5 * respTTL
	// Allow varnish's beresp.saintmode to expire.
	const waitSaintMode = time.Duration(5 * time.Second)
	headerValue := fmt.Sprintf("max-age=%.0f", respTTL.Seconds())

	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	req := NewUniqueEdgeGET(t)

	var expectedBody string
	for requestCount := 1; requestCount < 6; requestCount++ {
		switch requestCount {
		case 1: // Request 1 populates cache.
			expectedBody = expectedResponseStale

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Cache-Control", headerValue)
				w.Write([]byte(expectedBody))
			})
		case 2: // Requests 2,3,4 come from stale.
			time.Sleep(respTTLWithBuffer)
			expectedBody = expectedResponseStale

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusServiceUnavailable)
				w.Write([]byte(originServer.Name))
			})
		case 5: // Last request comes directly from origin again.
			time.Sleep(waitSaintMode)
			expectedBody = expectedResponseFresh

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(expectedBody))
			})
		}

		resp := RoundTripCheckError(t, req)

		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Fatal(err)
		}
		if bodyStr := string(body); bodyStr != expectedBody {
			t.Errorf(
				"Request %d received incorrect response body. Expected %q, got %q",
				requestCount,
				expectedBody,
				bodyStr,
			)
		}
	}
}

// Should fallback to first mirror if origin is down and object is not in
// cache (active or stale).
func TestFailoverOriginDownUseFirstMirror(t *testing.T) {
	ResetBackends(backendsByPriority)

	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK

	originServer.Stop()
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(expectedBody))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received a request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should fallback to first mirror if origin returns 5xx response and object
// is not in cache (active or stale).
func TestFailoverOrigin5xxUseFirstMirror(t *testing.T) {
	ResetBackends(backendsByPriority)

	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK
	backendsSawRequest := map[string]bool{}

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := originServer.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		if !backendsSawRequest[name] {
			w.Write([]byte(expectedBody))
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
			w.Write([]byte(name))
		}
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received a request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should fallback to second mirror if both origin and first mirror are
// down.
func TestFailoverOriginDownFirstMirrorDownUseSecondMirror(t *testing.T) {
	ResetBackends(backendsByPriority)

	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK

	originServer.Stop()
	backupServer1.Stop()
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(expectedBody))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should fallback to second mirror if both origin and first mirror return
// 5xx responses.
func TestFailoverOrigin5xxFirstMirror5xxUseSecondMirror(t *testing.T) {
	ResetBackends(backendsByPriority)

	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK
	backendsSawRequest := map[string]bool{}

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := originServer.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		if !backendsSawRequest[name] {
			w.Write([]byte(expectedBody))
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
			w.Write([]byte(name))
		}
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should not fallback to mirror if origin returns a 5xx response with a
// No-Fallback header. In order to allow applications to present their own
// error pages.
func TestFailoverNoFallbackHeader(t *testing.T) {
	ResetBackends(backendsByPriority)

	const headerName = "No-Fallback"
	const expectedStatus = http.StatusServiceUnavailable
	const expectedBody = "custom error page"

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set(headerName, "")
		w.WriteHeader(expectedStatus)
		w.Write([]byte(expectedBody))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	req := NewUniqueEdgeGET(t)
	resp := RoundTripCheckError(t, req)

	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}
<|endoftext|>


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

package rest

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/couchbaselabs/go-couchbase"

	"github.com/couchbaselabs/sync_gateway/base"
	"github.com/couchbaselabs/sync_gateway/db"
)

// The URL that stats will be reported to if deployment_id is set in the config
const kStatsReportURL = "http://localhost:9999/stats"
const kStatsReportInterval = time.Hour
const kDefaultSlowServerCallWarningThreshold = 200 // ms

// Shared context of HTTP handlers: primarily a registry of databases by name. It also stores
// the configuration settings so handlers can refer to them.
// This struct is accessed from HTTP handlers running on multiple goroutines, so it needs to
// be thread-safe.
type ServerContext struct {
	config      *ServerConfig
	databases_  map[string]*db.DatabaseContext
	lock        sync.RWMutex
	statsTicker *time.Ticker
	HTTPClient  *http.Client
}

func NewServerContext(config *ServerConfig) *ServerContext {
	sc := &ServerContext{
		config:     config,
		databases_: map[string]*db.DatabaseContext{},
		HTTPClient: http.DefaultClient,
	}
	if config.Databases == nil {
		config.Databases = DbConfigMap{}
	}

	// Initialize the go-couchbase library's global configuration variables:
	couchbase.PoolSize = DefaultMaxCouchbaseConnections
	couchbase.PoolOverflow = DefaultMaxCouchbaseOverflowConnections
	if config.MaxCouchbaseConnections != nil {
		couchbase.PoolSize = *config.MaxCouchbaseConnections
	}
	if config.MaxCouchbaseOverflow != nil {
		couchbase.PoolOverflow = *config.MaxCouchbaseOverflow
	}

	slow := kDefaultSlowServerCallWarningThreshold
	if config.SlowServerCallWarningThreshold != nil {
		slow = *config.SlowServerCallWarningThreshold
	}
	couchbase.SlowServerCallWarningThreshold = time.Duration(slow) * time.Millisecond

	if config.DeploymentID != nil {
		sc.startStatsReporter()
	}
	return sc
}

func (sc *ServerContext) Close() {
	sc.lock.Lock()
	defer sc.lock.Unlock()

	sc.stopStatsReporter()
	for _, ctx := range sc.databases_ {
		ctx.Close()
	}
	sc.databases_ = nil
}

// Returns the DatabaseContext with the given name
func (sc *ServerContext) GetDatabase(name string) (*db.DatabaseContext, error) {
	sc.lock.RLock()
	dbc := sc.databases_[name]
	sc.lock.RUnlock()
	if dbc != nil {
		return dbc, nil
	} else if db.ValidateDatabaseName(name) != nil {
		return nil, base.HTTPErrorf(http.StatusBadRequest, "invalid database name %q", name)
	} else if sc.config.ConfigServer == nil {
		return nil, base.HTTPErrorf(http.StatusNotFound, "no such database %q", name)
	} else {
		// Let's ask the config server if it knows this database:
		base.Log("Asking config server %q about db %q...", *sc.config.ConfigServer, name)
		config, err := sc.getDbConfigFromServer(name)
		if err != nil {
			return nil, err
		}
		if dbc, err = sc.AddDatabaseFromConfig(config); err != nil {
			return nil, err
		}
		return dbc, nil
	}
}

func (sc *ServerContext) GetDatabaseConfig(name string) *DbConfig {
	sc.lock.RLock()
	config := sc.config.Databases[name]
	sc.lock.RUnlock()
	return config
}

func (sc *ServerContext) setDatabaseConfig(name string, config *DbConfig) {
	sc.lock.Lock()
	sc.config.Databases[name] = config
	sc.lock.Unlock()
}

func (sc *ServerContext) AllDatabaseNames() []string {
	sc.lock.Lock()
	defer sc.lock.Unlock()

	names := make([]string, 0, len(sc.databases_))
	for name, _ := range sc.databases_ {
		names = append(names, name)
	}
	return names
}

func (sc *ServerContext) registerDatabase(dbcontext *db.DatabaseContext) error {
	sc.lock.Lock()
	defer sc.lock.Unlock()

	name := dbcontext.Name
	if sc.databases_[name] != nil {
		return base.HTTPErrorf(http.StatusPreconditionFailed, // what CouchDB returns
			"Duplicate database name %q", name)
	}
	sc.databases_[name] = dbcontext
	return nil
}

// Adds a database to the ServerContext given its configuration.
func (sc *ServerContext) AddDatabaseFromConfig(config *DbConfig) (*db.DatabaseContext, error) {
	server := "http://localhost:8091"
	pool := "default"
	bucketName := config.name

	if config.Server != nil {
		server = *config.Server
	}
	if config.Pool != nil {
		pool = *config.Pool
	}
	if config.Bucket != nil {
		bucketName = *config.Bucket
	}
	dbName := config.name
	if dbName == "" {
		dbName = bucketName
	}
	base.Log("Opening db /%s as bucket %q, pool %q, server <%s>",
		dbName, bucketName, pool, server)

	if err := db.ValidateDatabaseName(dbName); err != nil {
		return nil, err
	}

	var importDocs, autoImport bool
	switch config.ImportDocs {
	case nil, false:
	case true:
		importDocs = true
	case "continuous":
		importDocs = true
		autoImport = true
	default:
		return nil, fmt.Errorf("Unrecognized value for ImportDocs: %#v", config.ImportDocs)
	}

	// Connect to the bucket and add the database:
	spec := base.BucketSpec{
		Server:     server,
		PoolName:   pool,
		BucketName: bucketName,
	}
	if config.Username != "" {
		spec.Auth = config
	}
	bucket, err := db.ConnectToBucket(spec)
	if err != nil {
		return nil, err
	}
	dbcontext, err := db.NewDatabaseContext(dbName, bucket, autoImport)
	if err != nil {
		return nil, err
	}

	syncFn := ""
	if config.Sync != nil {
		syncFn = *config.Sync
	}
	if err := dbcontext.ApplySyncFun(syncFn, importDocs); err != nil {
		return nil, err
	}

	if config.RevsLimit != nil && *config.RevsLimit > 0 {
		dbcontext.RevsLimit = *config.RevsLimit
	}

	if dbcontext.ChannelMapper == nil {
		base.Log("Using default sync function 'channel(doc.channels)' for database %q", dbName)
	}

	// Create default users & roles:
	if err := sc.installPrincipals(dbcontext, config.Roles, "role"); err != nil {
		return nil, err
	} else if err := sc.installPrincipals(dbcontext, config.Users, "user"); err != nil {
		return nil, err
	}

	// Install bucket-shadower if any:
	if shadow := config.Shadow; shadow != nil {
		if err := sc.startShadowing(dbcontext, shadow); err != nil {
			base.Warn("Database %q: unable to connect to external bucket for shadowing: %v",
				dbName, err)
		}
	}

	// Register it so HTTP handlers can find it:
	if err := sc.registerDatabase(dbcontext); err != nil {
		dbcontext.Close()
		return nil, err
	}
	sc.setDatabaseConfig(config.name, config)
	return dbcontext, nil
}

func (sc *ServerContext) startShadowing(dbcontext *db.DatabaseContext, shadow *ShadowConfig) error {
	var pattern *regexp.Regexp
	if shadow.Doc_id_regex != nil {
		var err error
		pattern, err = regexp.Compile(*shadow.Doc_id_regex)
		if err != nil {
			base.Warn("Invalid shadow doc_id_regex: %s", *shadow.Doc_id_regex)
			return err
		}
	}

	spec := base.BucketSpec{
		Server:     shadow.Server,
		PoolName:   "default",
		BucketName: shadow.Bucket,
	}
	if shadow.Pool != nil {
		spec.PoolName = *shadow.Pool
	}
	if shadow.Username != "" {
		spec.Auth = shadow
	}

	bucket, err := db.ConnectToBucket(spec)
	if err != nil {
		return err
	}
	shadower, err := db.NewShadower(dbcontext, bucket, pattern)
	if err != nil {
		bucket.Close()
		return err
	}
	dbcontext.Shadower = shadower

        //Remove credentials from server URL before logging
        url, err := couchbase.ParseURL(spec.Server)
        if err == nil {
	    base.Log("Database %q shadowing remote bucket %q, pool %q, server <%s:%s/%s>", dbcontext.Name, spec.BucketName, spec.PoolName, url.Scheme, url.Host, url.Path)
        }
	return nil
}

func (sc *ServerContext) RemoveDatabase(dbName string) bool {
	sc.lock.Lock()
	defer sc.lock.Unlock()

	context := sc.databases_[dbName]
	if context == nil {
		return false
	}
	base.Log("Closing db /%s (bucket %q)", context.Name, context.Bucket.GetName())
	context.Close()
	delete(sc.databases_, dbName)
	return true
}

func (sc *ServerContext) installPrincipals(context *db.DatabaseContext, spec map[string]*db.PrincipalConfig, what string) error {
	for name, princ := range spec {
		isGuest := name == "GUEST"
		if isGuest {
			internalName := ""
			princ.Name = &internalName
		} else {
			princ.Name = &name
		}
		_, err := context.UpdatePrincipal(*princ, (what == "user"), isGuest)
		if err != nil {
			// A conflict error just means updatePrincipal didn't overwrite an existing user.
			if status, _ := base.ErrorAsHTTPStatus(err); status != http.StatusConflict {
				return fmt.Errorf("Couldn't create %s %q: %v", what, name, err)
			}
		} else if isGuest {
			base.Log("    Reset guest user to config")
		} else {
			base.Log("    Created %s %q", what, name)
		}
	}
	return nil
}

// Fetch a configuration for a database from the ConfigServer
func (sc *ServerContext) getDbConfigFromServer(dbName string) (*DbConfig, error) {
	if sc.config.ConfigServer == nil {
		return nil, base.HTTPErrorf(http.StatusNotFound, "not_found")
	}

	urlStr := *sc.config.ConfigServer
	if !strings.HasSuffix(urlStr, "/") {
		urlStr += "/"
	}
	urlStr += url.QueryEscape(dbName)
	res, err := sc.HTTPClient.Get(urlStr)
	if err != nil {
		return nil, base.HTTPErrorf(http.StatusBadGateway,
			"Error contacting config server: %v", err)
	} else if res.StatusCode >= 300 {
		return nil, base.HTTPErrorf(res.StatusCode, res.Status)
	}

	var config DbConfig
	j := json.NewDecoder(res.Body)
	if err = j.Decode(&config); err != nil {
		return nil, base.HTTPErrorf(http.StatusBadGateway,
			"Bad response from config server: %v", err)
	}

	if err = config.setup(dbName); err != nil {
		return nil, err
	}
	return &config, nil
}

//////// STATISTICS REPORT:

func (sc *ServerContext) startStatsReporter() {
	interval := kStatsReportInterval
	if sc.config.StatsReportInterval != nil {
		if *sc.config.StatsReportInterval <= 0 {
			return
		}
		interval = time.Duration(*sc.config.StatsReportInterval) * time.Second
	}
	sc.statsTicker = time.NewTicker(interval)
	go func() {
		for _ = range sc.statsTicker.C {
			sc.reportStats()
		}
	}()
	base.Log("Will report server stats for %q every %v",
		*sc.config.DeploymentID, interval)
}

func (sc *ServerContext) stopStatsReporter() {
	if sc.statsTicker != nil {
		sc.statsTicker.Stop()
		sc.reportStats() // Report stuff since the last tick
	}
}

// POST a report of database statistics
func (sc *ServerContext) reportStats() {
	if sc.config.DeploymentID == nil {
		panic("Can't reportStats without DeploymentID")
	}
	stats := sc.Stats()
	if stats == nil {
		return // No activity
	}
	base.Log("Reporting server stats to %s ...", kStatsReportURL)
	body, _ := json.Marshal(stats)
	bodyReader := bytes.NewReader(body)
	_, err := sc.HTTPClient.Post(kStatsReportURL, "application/json", bodyReader)
	if err != nil {
		base.Warn("Error posting stats: %v", err)
	}
}

func (sc *ServerContext) Stats() map[string]interface{} {
	sc.lock.RLock()
	defer sc.lock.RUnlock()
	var stats []map[string]interface{}
	any := false
	for _, dbc := range sc.databases_ {
		max := dbc.ChangesClientStats.MaxCount()
		total := dbc.ChangesClientStats.TotalCount()
		dbc.ChangesClientStats.Reset()
		stats = append(stats, map[string]interface{}{
			"max_connections":   max,
			"total_connections": total,
		})
		any = any || total > 0
	}
	if !any {
		return nil
	}
	return map[string]interface{}{
		"deploymentID": *sc.config.DeploymentID,
		"databases":    stats,
	}
}
<commit_msg>ran file through Fmt<commit_after>//  Copyright (c) 2013 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package rest

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/couchbaselabs/go-couchbase"

	"github.com/couchbaselabs/sync_gateway/base"
	"github.com/couchbaselabs/sync_gateway/db"
)

// The URL that stats will be reported to if deployment_id is set in the config
const kStatsReportURL = "http://localhost:9999/stats"
const kStatsReportInterval = time.Hour
const kDefaultSlowServerCallWarningThreshold = 200 // ms

// Shared context of HTTP handlers: primarily a registry of databases by name. It also stores
// the configuration settings so handlers can refer to them.
// This struct is accessed from HTTP handlers running on multiple goroutines, so it needs to
// be thread-safe.
type ServerContext struct {
	config      *ServerConfig
	databases_  map[string]*db.DatabaseContext
	lock        sync.RWMutex
	statsTicker *time.Ticker
	HTTPClient  *http.Client
}

func NewServerContext(config *ServerConfig) *ServerContext {
	sc := &ServerContext{
		config:     config,
		databases_: map[string]*db.DatabaseContext{},
		HTTPClient: http.DefaultClient,
	}
	if config.Databases == nil {
		config.Databases = DbConfigMap{}
	}

	// Initialize the go-couchbase library's global configuration variables:
	couchbase.PoolSize = DefaultMaxCouchbaseConnections
	couchbase.PoolOverflow = DefaultMaxCouchbaseOverflowConnections
	if config.MaxCouchbaseConnections != nil {
		couchbase.PoolSize = *config.MaxCouchbaseConnections
	}
	if config.MaxCouchbaseOverflow != nil {
		couchbase.PoolOverflow = *config.MaxCouchbaseOverflow
	}

	slow := kDefaultSlowServerCallWarningThreshold
	if config.SlowServerCallWarningThreshold != nil {
		slow = *config.SlowServerCallWarningThreshold
	}
	couchbase.SlowServerCallWarningThreshold = time.Duration(slow) * time.Millisecond

	if config.DeploymentID != nil {
		sc.startStatsReporter()
	}
	return sc
}

func (sc *ServerContext) Close() {
	sc.lock.Lock()
	defer sc.lock.Unlock()

	sc.stopStatsReporter()
	for _, ctx := range sc.databases_ {
		ctx.Close()
	}
	sc.databases_ = nil
}

// Returns the DatabaseContext with the given name
func (sc *ServerContext) GetDatabase(name string) (*db.DatabaseContext, error) {
	sc.lock.RLock()
	dbc := sc.databases_[name]
	sc.lock.RUnlock()
	if dbc != nil {
		return dbc, nil
	} else if db.ValidateDatabaseName(name) != nil {
		return nil, base.HTTPErrorf(http.StatusBadRequest, "invalid database name %q", name)
	} else if sc.config.ConfigServer == nil {
		return nil, base.HTTPErrorf(http.StatusNotFound, "no such database %q", name)
	} else {
		// Let's ask the config server if it knows this database:
		base.Log("Asking config server %q about db %q...", *sc.config.ConfigServer, name)
		config, err := sc.getDbConfigFromServer(name)
		if err != nil {
			return nil, err
		}
		if dbc, err = sc.AddDatabaseFromConfig(config); err != nil {
			return nil, err
		}
		return dbc, nil
	}
}

func (sc *ServerContext) GetDatabaseConfig(name string) *DbConfig {
	sc.lock.RLock()
	config := sc.config.Databases[name]
	sc.lock.RUnlock()
	return config
}

func (sc *ServerContext) setDatabaseConfig(name string, config *DbConfig) {
	sc.lock.Lock()
	sc.config.Databases[name] = config
	sc.lock.Unlock()
}

func (sc *ServerContext) AllDatabaseNames() []string {
	sc.lock.Lock()
	defer sc.lock.Unlock()

	names := make([]string, 0, len(sc.databases_))
	for name, _ := range sc.databases_ {
		names = append(names, name)
	}
	return names
}

func (sc *ServerContext) registerDatabase(dbcontext *db.DatabaseContext) error {
	sc.lock.Lock()
	defer sc.lock.Unlock()

	name := dbcontext.Name
	if sc.databases_[name] != nil {
		return base.HTTPErrorf(http.StatusPreconditionFailed, // what CouchDB returns
			"Duplicate database name %q", name)
	}
	sc.databases_[name] = dbcontext
	return nil
}

// Adds a database to the ServerContext given its configuration.
func (sc *ServerContext) AddDatabaseFromConfig(config *DbConfig) (*db.DatabaseContext, error) {
	server := "http://localhost:8091"
	pool := "default"
	bucketName := config.name

	if config.Server != nil {
		server = *config.Server
	}
	if config.Pool != nil {
		pool = *config.Pool
	}
	if config.Bucket != nil {
		bucketName = *config.Bucket
	}
	dbName := config.name
	if dbName == "" {
		dbName = bucketName
	}
	base.Log("Opening db /%s as bucket %q, pool %q, server <%s>",
		dbName, bucketName, pool, server)

	if err := db.ValidateDatabaseName(dbName); err != nil {
		return nil, err
	}

	var importDocs, autoImport bool
	switch config.ImportDocs {
	case nil, false:
	case true:
		importDocs = true
	case "continuous":
		importDocs = true
		autoImport = true
	default:
		return nil, fmt.Errorf("Unrecognized value for ImportDocs: %#v", config.ImportDocs)
	}

	// Connect to the bucket and add the database:
	spec := base.BucketSpec{
		Server:     server,
		PoolName:   pool,
		BucketName: bucketName,
	}
	if config.Username != "" {
		spec.Auth = config
	}
	bucket, err := db.ConnectToBucket(spec)
	if err != nil {
		return nil, err
	}
	dbcontext, err := db.NewDatabaseContext(dbName, bucket, autoImport)
	if err != nil {
		return nil, err
	}

	syncFn := ""
	if config.Sync != nil {
		syncFn = *config.Sync
	}
	if err := dbcontext.ApplySyncFun(syncFn, importDocs); err != nil {
		return nil, err
	}

	if config.RevsLimit != nil && *config.RevsLimit > 0 {
		dbcontext.RevsLimit = *config.RevsLimit
	}

	if dbcontext.ChannelMapper == nil {
		base.Log("Using default sync function 'channel(doc.channels)' for database %q", dbName)
	}

	// Create default users & roles:
	if err := sc.installPrincipals(dbcontext, config.Roles, "role"); err != nil {
		return nil, err
	} else if err := sc.installPrincipals(dbcontext, config.Users, "user"); err != nil {
		return nil, err
	}

	// Install bucket-shadower if any:
	if shadow := config.Shadow; shadow != nil {
		if err := sc.startShadowing(dbcontext, shadow); err != nil {
			base.Warn("Database %q: unable to connect to external bucket for shadowing: %v",
				dbName, err)
		}
	}

	// Register it so HTTP handlers can find it:
	if err := sc.registerDatabase(dbcontext); err != nil {
		dbcontext.Close()
		return nil, err
	}
	sc.setDatabaseConfig(config.name, config)
	return dbcontext, nil
}

func (sc *ServerContext) startShadowing(dbcontext *db.DatabaseContext, shadow *ShadowConfig) error {
	var pattern *regexp.Regexp
	if shadow.Doc_id_regex != nil {
		var err error
		pattern, err = regexp.Compile(*shadow.Doc_id_regex)
		if err != nil {
			base.Warn("Invalid shadow doc_id_regex: %s", *shadow.Doc_id_regex)
			return err
		}
	}

	spec := base.BucketSpec{
		Server:     shadow.Server,
		PoolName:   "default",
		BucketName: shadow.Bucket,
	}
	if shadow.Pool != nil {
		spec.PoolName = *shadow.Pool
	}
	if shadow.Username != "" {
		spec.Auth = shadow
	}

	bucket, err := db.ConnectToBucket(spec)
	if err != nil {
		return err
	}
	shadower, err := db.NewShadower(dbcontext, bucket, pattern)
	if err != nil {
		bucket.Close()
		return err
	}
	dbcontext.Shadower = shadower

	//Remove credentials from server URL before logging
	url, err := couchbase.ParseURL(spec.Server)
	if err == nil {
		base.Log("Database %q shadowing remote bucket %q, pool %q, server <%s:%s/%s>", dbcontext.Name, spec.BucketName, spec.PoolName, url.Scheme, url.Host, url.Path)
	}
	return nil
}

func (sc *ServerContext) RemoveDatabase(dbName string) bool {
	sc.lock.Lock()
	defer sc.lock.Unlock()

	context := sc.databases_[dbName]
	if context == nil {
		return false
	}
	base.Log("Closing db /%s (bucket %q)", context.Name, context.Bucket.GetName())
	context.Close()
	delete(sc.databases_, dbName)
	return true
}

func (sc *ServerContext) installPrincipals(context *db.DatabaseContext, spec map[string]*db.PrincipalConfig, what string) error {
	for name, princ := range spec {
		isGuest := name == "GUEST"
		if isGuest {
			internalName := ""
			princ.Name = &internalName
		} else {
			princ.Name = &name
		}
		_, err := context.UpdatePrincipal(*princ, (what == "user"), isGuest)
		if err != nil {
			// A conflict error just means updatePrincipal didn't overwrite an existing user.
			if status, _ := base.ErrorAsHTTPStatus(err); status != http.StatusConflict {
				return fmt.Errorf("Couldn't create %s %q: %v", what, name, err)
			}
		} else if isGuest {
			base.Log("    Reset guest user to config")
		} else {
			base.Log("    Created %s %q", what, name)
		}
	}
	return nil
}

// Fetch a configuration for a database from the ConfigServer
func (sc *ServerContext) getDbConfigFromServer(dbName string) (*DbConfig, error) {
	if sc.config.ConfigServer == nil {
		return nil, base.HTTPErrorf(http.StatusNotFound, "not_found")
	}

	urlStr := *sc.config.ConfigServer
	if !strings.HasSuffix(urlStr, "/") {
		urlStr += "/"
	}
	urlStr += url.QueryEscape(dbName)
	res, err := sc.HTTPClient.Get(urlStr)
	if err != nil {
		return nil, base.HTTPErrorf(http.StatusBadGateway,
			"Error contacting config server: %v", err)
	} else if res.StatusCode >= 300 {
		return nil, base.HTTPErrorf(res.StatusCode, res.Status)
	}

	var config DbConfig
	j := json.NewDecoder(res.Body)
	if err = j.Decode(&config); err != nil {
		return nil, base.HTTPErrorf(http.StatusBadGateway,
			"Bad response from config server: %v", err)
	}

	if err = config.setup(dbName); err != nil {
		return nil, err
	}
	return &config, nil
}

//////// STATISTICS REPORT:

func (sc *ServerContext) startStatsReporter() {
	interval := kStatsReportInterval
	if sc.config.StatsReportInterval != nil {
		if *sc.config.StatsReportInterval <= 0 {
			return
		}
		interval = time.Duration(*sc.config.StatsReportInterval) * time.Second
	}
	sc.statsTicker = time.NewTicker(interval)
	go func() {
		for _ = range sc.statsTicker.C {
			sc.reportStats()
		}
	}()
	base.Log("Will report server stats for %q every %v",
		*sc.config.DeploymentID, interval)
}

func (sc *ServerContext) stopStatsReporter() {
	if sc.statsTicker != nil {
		sc.statsTicker.Stop()
		sc.reportStats() // Report stuff since the last tick
	}
}

// POST a report of database statistics
func (sc *ServerContext) reportStats() {
	if sc.config.DeploymentID == nil {
		panic("Can't reportStats without DeploymentID")
	}
	stats := sc.Stats()
	if stats == nil {
		return // No activity
	}
	base.Log("Reporting server stats to %s ...", kStatsReportURL)
	body, _ := json.Marshal(stats)
	bodyReader := bytes.NewReader(body)
	_, err := sc.HTTPClient.Post(kStatsReportURL, "application/json", bodyReader)
	if err != nil {
		base.Warn("Error posting stats: %v", err)
	}
}

func (sc *ServerContext) Stats() map[string]interface{} {
	sc.lock.RLock()
	defer sc.lock.RUnlock()
	var stats []map[string]interface{}
	any := false
	for _, dbc := range sc.databases_ {
		max := dbc.ChangesClientStats.MaxCount()
		total := dbc.ChangesClientStats.TotalCount()
		dbc.ChangesClientStats.Reset()
		stats = append(stats, map[string]interface{}{
			"max_connections":   max,
			"total_connections": total,
		})
		any = any || total > 0
	}
	if !any {
		return nil
	}
	return map[string]interface{}{
		"deploymentID": *sc.config.DeploymentID,
		"databases":    stats,
	}
}
<|endoftext|>


<commit_before>// Pin connection between FT232R and LCD (HD44780 compatible):
// TxD  (DBUS0) <--> D4
// RxD  (DBUS1) <--> D5
// RTS# (DBUS2) <--> D6
// CTS# (DBUS3) <--> D7
// DTR# (DBUS4) <--> E
// DSR# (DBUS5) <--> R/W#
// DCD# (DBUS6) <--> RS
package main

import (
	"flag"
	"fmt"
	"github.com/ziutek/ftdi"
	"github.com/ziutek/lcd/hdc"
	"os"
	"strings"
)

func checkErr(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

const (
	baudrate  = 1 << 17
	waitTicks = 6
)

var (
	x = flag.Int("x", 0, "set row")
	y = flag.Int("y", 0, "sec column")
	c = flag.Bool("c", false, "clear display")
)

func main() {
	flag.Parse()
	text := strings.Join(flag.Args(), " ")

	d, err := ftdi.OpenFirst(0x0403, 0x6001, ftdi.ChannelAny)
	checkErr(err)
	defer d.Close()
	checkErr(d.SetBitmode(0xff, ftdi.ModeBitbang))
	checkErr(d.SetBaudrate(baudrate / 16))

	lcd := hdc.NewDevice(hdc.NewBitbang(d, waitTicks), 4, 20)
	checkErr(lcd.Init())
	checkErr(lcd.SetDisplay(hdc.DisplayOn))

	if *c {
		checkErr(lcd.ClearDisplay())
	}

	checkErr(lcd.MoveCursor(*x, *y))
	_, err = fmt.Fprintf(lcd, text)
	checkErr(err)
}
<commit_msg>Fix arguments names<commit_after>// Pin connection between FT232R and LCD (HD44780 compatible):
// TxD  (DBUS0) <--> D4
// RxD  (DBUS1) <--> D5
// RTS# (DBUS2) <--> D6
// CTS# (DBUS3) <--> D7
// DTR# (DBUS4) <--> E
// DSR# (DBUS5) <--> R/W#
// DCD# (DBUS6) <--> RS
package main

import (
	"flag"
	"fmt"
	"github.com/ziutek/ftdi"
	"github.com/ziutek/lcd/hdc"
	"os"
	"strings"
)

func checkErr(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

const (
	baudrate  = 1 << 17
	waitTicks = 6
)

var (
	x = flag.Int("x", 0, "column number")
	y = flag.Int("y", 0, "row number")
	c = flag.Bool("c", false, "clear display")
)

func main() {
	flag.Parse()
	text := strings.Join(flag.Args(), " ")

	d, err := ftdi.OpenFirst(0x0403, 0x6001, ftdi.ChannelAny)
	checkErr(err)
	defer d.Close()
	checkErr(d.SetBitmode(0xff, ftdi.ModeBitbang))
	checkErr(d.SetBaudrate(baudrate / 16))

	lcd := hdc.NewDevice(hdc.NewBitbang(d, waitTicks), 4, 20)
	checkErr(lcd.Init())
	checkErr(lcd.SetDisplay(hdc.DisplayOn))

	if *c {
		checkErr(lcd.ClearDisplay())
	}

	checkErr(lcd.MoveCursor(*x, *y))
	_, err = fmt.Fprintf(lcd, text)
	checkErr(err)
}
<|endoftext|>


<commit_before>package cmd

import (
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/models/app"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/models/credential"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/models/info"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/options"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/sigwinch"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/terminal"
	"github.com/cloudfoundry-incubator/diego-ssh/helpers"
	"github.com/docker/docker/pkg/term"
)

type SecureShell interface {
	Connect(opts *options.SSHOptions) error
	InteractiveSession() error
	LocalPortForward() error
	Wait() error
	Close() error
}

//go:generate counterfeiter -o fakes/fake_secure_dialer.go . SecureDialer
type SecureDialer interface {
	Dial(network, address string, config *ssh.ClientConfig) (SecureClient, error)
}

//go:generate counterfeiter -o fakes/fake_secure_client.go . SecureClient
type SecureClient interface {
	NewSession() (SecureSession, error)
	Conn() ssh.Conn
	Dial(network, address string) (net.Conn, error)
	Wait() error
	Close() error
}

//go:generate counterfeiter -o fakes/fake_listener_factory.go . ListenerFactory
type ListenerFactory interface {
	Listen(network, address string) (net.Listener, error)
}

//go:generate counterfeiter -o fakes/fake_secure_session.go . SecureSession
type SecureSession interface {
	RequestPty(term string, height, width int, termModes ssh.TerminalModes) error
	SendRequest(name string, wantReply bool, payload []byte) (bool, error)
	StdinPipe() (io.WriteCloser, error)
	StdoutPipe() (io.Reader, error)
	StderrPipe() (io.Reader, error)
	Start(command string) error
	Shell() error
	Wait() error
	Close() error
}

type secureShell struct {
	secureDialer      SecureDialer
	terminalHelper    terminal.TerminalHelper
	listenerFactory   ListenerFactory
	keepAliveInterval time.Duration
	appFactory        app.AppFactory
	infoFactory       info.InfoFactory
	credFactory       credential.CredentialFactory
	secureClient      SecureClient
	opts              *options.SSHOptions

	localListeners []net.Listener
}

func NewSecureShell(
	secureDialer SecureDialer,
	terminalHelper terminal.TerminalHelper,
	listenerFactory ListenerFactory,
	keepAliveInterval time.Duration,
	appFactory app.AppFactory,
	infoFactory info.InfoFactory,
	credFactory credential.CredentialFactory,
) SecureShell {
	return &secureShell{
		secureDialer:      secureDialer,
		terminalHelper:    terminalHelper,
		listenerFactory:   listenerFactory,
		keepAliveInterval: keepAliveInterval,
		appFactory:        appFactory,
		infoFactory:       infoFactory,
		credFactory:       credFactory,
		localListeners:    []net.Listener{},
	}
}

func (c *secureShell) Connect(opts *options.SSHOptions) error {
	app, err := c.appFactory.Get(opts.AppName)
	if err != nil {
		return err
	}

	info, err := c.infoFactory.Get()
	if err != nil {
		return err
	}

	cred, err := c.credFactory.Get()
	if err != nil {
		return err
	}

	err = c.validateTarget(app, opts)
	if err != nil {
		return err
	}

	clientConfig := &ssh.ClientConfig{
		User: fmt.Sprintf("cf:%s/%d", app.Guid, opts.Instance),
		Auth: []ssh.AuthMethod{
			ssh.Password(cred.Token),
		},
		HostKeyCallback: fingerprintCallback(opts, info.SSHEndpointFingerprint),
	}

	secureClient, err := c.secureDialer.Dial("tcp", info.SSHEndpoint, clientConfig)
	if err != nil {
		return err
	}

	c.secureClient = secureClient
	c.opts = opts
	return nil
}

func (c *secureShell) Close() error {
	for _, listener := range c.localListeners {
		listener.Close()
	}
	return c.secureClient.Close()
}

func (c *secureShell) LocalPortForward() error {
	for _, forwardSpec := range c.opts.ForwardSpecs {
		listener, err := c.listenerFactory.Listen("tcp", forwardSpec.ListenAddress)
		if err != nil {
			return err
		}
		c.localListeners = append(c.localListeners, listener)

		go c.localForwardAcceptLoop(listener, forwardSpec.ConnectAddress)
	}

	return nil
}

func (c *secureShell) localForwardAcceptLoop(listener net.Listener, addr string) {
	defer listener.Close()

	for {
		conn, err := listener.Accept()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				time.Sleep(100 * time.Millisecond)
				continue
			}
			return
		}

		go c.handleForwardConnection(conn, addr)
	}
}

func (c *secureShell) handleForwardConnection(conn net.Conn, targetAddr string) {
	defer conn.Close()

	target, err := c.secureClient.Dial("tcp", targetAddr)
	if err != nil {
		return
	}
	defer target.Close()

	wg := &sync.WaitGroup{}
	wg.Add(2)

	go copyAndClose(wg, conn, target)
	go copyAndClose(wg, target, conn)
	wg.Wait()
}

func copyAndClose(wg *sync.WaitGroup, dest io.WriteCloser, src io.Reader) {
	io.Copy(dest, src)
	dest.Close()
	wg.Done()
}

func (c *secureShell) InteractiveSession() error {
	secureClient := c.secureClient
	opts := c.opts

	session, err := secureClient.NewSession()
	if err != nil {
		return fmt.Errorf("SSH session allocation failed: %s", err.Error())
	}
	defer session.Close()

	stdin, stdout, stderr := c.terminalHelper.StdStreams()

	inPipe, err := session.StdinPipe()
	if err != nil {
		return err
	}

	outPipe, err := session.StdoutPipe()
	if err != nil {
		return err
	}

	errPipe, err := session.StderrPipe()
	if err != nil {
		return err
	}

	stdinFd, stdinIsTerminal := c.terminalHelper.GetFdInfo(stdin)
	stdoutFd, stdoutIsTerminal := c.terminalHelper.GetFdInfo(stdout)

	if c.shouldAllocateTerminal(opts, stdinIsTerminal) {
		modes := ssh.TerminalModes{
			ssh.ECHO:          1,
			ssh.TTY_OP_ISPEED: 115200,
			ssh.TTY_OP_OSPEED: 115200,
		}

		width, height := c.getWindowDimensions(stdoutFd)

		err = session.RequestPty(c.terminalType(), height, width, modes)
		if err != nil {
			return err
		}

		state, err := c.terminalHelper.SetRawTerminal(stdinFd)
		if err == nil {
			defer c.terminalHelper.RestoreTerminal(stdinFd, state)
		}
	}

	if len(opts.Command) != 0 {
		cmd := strings.Join(opts.Command, " ")
		err = session.Start(cmd)
		if err != nil {
			return err
		}
	} else {
		err = session.Shell()
		if err != nil {
			return err
		}
	}

	go io.Copy(inPipe, stdin)
	go io.Copy(stdout, outPipe)
	go io.Copy(stderr, errPipe)

	if stdoutIsTerminal {
		resized := make(chan os.Signal, 16)

		if runtime.GOOS == "windows" {
			ticker := time.NewTicker(250 * time.Millisecond)
			defer ticker.Stop()

			go func() {
				for _ = range ticker.C {
					resized <- syscall.Signal(-1)
				}
				close(resized)
			}()
		} else {
			signal.Notify(resized, sigwinch.SIGWINCH())
			defer func() { signal.Stop(resized); close(resized) }()
		}

		go c.resize(resized, session, stdoutFd)
	}

	keepaliveStopCh := make(chan struct{})
	defer close(keepaliveStopCh)

	go keepalive(secureClient.Conn(), time.NewTicker(c.keepAliveInterval), keepaliveStopCh)

	return session.Wait()
}

func (c *secureShell) Wait() error {
	return c.secureClient.Wait()
}

func (c *secureShell) validateTarget(app app.App, opts *options.SSHOptions) error {
	if app.State != "STARTED" {
		return fmt.Errorf("Application %q is not in the STARTED state", opts.AppName)
	}

	if !app.Diego {
		return fmt.Errorf("Application %q is not running on Diego", opts.AppName)
	}

	return nil
}

type hostKeyCallback func(hostname string, remote net.Addr, key ssh.PublicKey) error

func fingerprintCallback(opts *options.SSHOptions, expectedFingerprint string) hostKeyCallback {
	if opts.SkipHostValidation {
		return nil
	}

	return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
		switch len(expectedFingerprint) {
		case helpers.SHA1_FINGERPRINT_LENGTH:
			fingerprint := helpers.SHA1Fingerprint(key)
			if fingerprint != expectedFingerprint {
				return fmt.Errorf("Host key verification failed.\n\nThe fingerprint of the received key was %q.", fingerprint)
			}
		case helpers.MD5_FINGERPRINT_LENGTH:
			fingerprint := helpers.MD5Fingerprint(key)
			if fingerprint != expectedFingerprint {
				return fmt.Errorf("Host key verification failed.\n\nThe fingerprint of the received key was %q.", fingerprint)
			}
		case 0:
			fingerprint := helpers.MD5Fingerprint(key)
			return fmt.Errorf("Unable to verify identity of host.\n\nThe fingerprint of the received key was %q.", fingerprint)
		default:
			return errors.New("Unsupported host key fingerprint format")
		}
		return nil
	}
}

func (c *secureShell) shouldAllocateTerminal(opts *options.SSHOptions, stdinIsTerminal bool) bool {
	switch opts.TerminalRequest {
	case options.REQUEST_TTY_FORCE:
		return true
	case options.REQUEST_TTY_NO:
		return false
	case options.REQUEST_TTY_YES:
		return stdinIsTerminal
	case options.REQUEST_TTY_AUTO:
		return len(opts.Command) == 0 && stdinIsTerminal
	default:
		return false
	}
}

func (c *secureShell) resize(resized <-chan os.Signal, session SecureSession, terminalFd uintptr) {
	type resizeMessage struct {
		Width       uint32
		Height      uint32
		PixelWidth  uint32
		PixelHeight uint32
	}

	var previousWidth, previousHeight int

	for _ = range resized {
		width, height := c.getWindowDimensions(terminalFd)

		if width == previousWidth && height == previousHeight {
			continue
		}

		message := resizeMessage{
			Width:  uint32(width),
			Height: uint32(height),
		}

		session.SendRequest("window-change", false, ssh.Marshal(message))

		previousWidth = width
		previousHeight = height
	}
}

func keepalive(conn ssh.Conn, ticker *time.Ticker, stopCh chan struct{}) {
	for {
		select {
		case <-ticker.C:
			conn.SendRequest("keepalive@cloudfoundry.org", true, nil)
		case <-stopCh:
			ticker.Stop()
			return
		}
	}
}

func (c *secureShell) terminalType() string {
	term := os.Getenv("TERM")
	if term == "" {
		term = "xterm"
	}
	return term
}

func (c *secureShell) getWindowDimensions(terminalFd uintptr) (width int, height int) {
	winSize, err := c.terminalHelper.GetWinsize(terminalFd)
	if err != nil {
		winSize = &term.Winsize{
			Width:  80,
			Height: 43,
		}
	}

	return int(winSize.Width), int(winSize.Height)
}

type secureDialer struct{}

func (d *secureDialer) Dial(network string, address string, config *ssh.ClientConfig) (SecureClient, error) {
	client, err := ssh.Dial(network, address, config)
	if err != nil {
		return nil, err
	}

	return &secureClient{client: client}, nil
}

func DefaultSecureDialer() SecureDialer {
	return &secureDialer{}
}

type secureClient struct{ client *ssh.Client }

func (sc *secureClient) Close() error   { return sc.client.Close() }
func (sc *secureClient) Conn() ssh.Conn { return sc.client.Conn }
func (sc *secureClient) Wait() error    { return sc.client.Wait() }
func (sc *secureClient) Dial(n, addr string) (net.Conn, error) {
	return sc.client.Dial(n, addr)
}
func (sc *secureClient) NewSession() (SecureSession, error) {
	return sc.client.NewSession()
}

type listenerFactory struct{}

func (lf *listenerFactory) Listen(network, address string) (net.Listener, error) {
	return net.Listen(network, address)
}

func DefaultListenerFactory() ListenerFactory {
	return &listenerFactory{}
}
<commit_msg>Forward close of stdin across ssh connection<commit_after>package cmd

import (
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/models/app"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/models/credential"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/models/info"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/options"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/sigwinch"
	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/terminal"
	"github.com/cloudfoundry-incubator/diego-ssh/helpers"
	"github.com/docker/docker/pkg/term"
)

type SecureShell interface {
	Connect(opts *options.SSHOptions) error
	InteractiveSession() error
	LocalPortForward() error
	Wait() error
	Close() error
}

//go:generate counterfeiter -o fakes/fake_secure_dialer.go . SecureDialer
type SecureDialer interface {
	Dial(network, address string, config *ssh.ClientConfig) (SecureClient, error)
}

//go:generate counterfeiter -o fakes/fake_secure_client.go . SecureClient
type SecureClient interface {
	NewSession() (SecureSession, error)
	Conn() ssh.Conn
	Dial(network, address string) (net.Conn, error)
	Wait() error
	Close() error
}

//go:generate counterfeiter -o fakes/fake_listener_factory.go . ListenerFactory
type ListenerFactory interface {
	Listen(network, address string) (net.Listener, error)
}

//go:generate counterfeiter -o fakes/fake_secure_session.go . SecureSession
type SecureSession interface {
	RequestPty(term string, height, width int, termModes ssh.TerminalModes) error
	SendRequest(name string, wantReply bool, payload []byte) (bool, error)
	StdinPipe() (io.WriteCloser, error)
	StdoutPipe() (io.Reader, error)
	StderrPipe() (io.Reader, error)
	Start(command string) error
	Shell() error
	Wait() error
	Close() error
}

type secureShell struct {
	secureDialer      SecureDialer
	terminalHelper    terminal.TerminalHelper
	listenerFactory   ListenerFactory
	keepAliveInterval time.Duration
	appFactory        app.AppFactory
	infoFactory       info.InfoFactory
	credFactory       credential.CredentialFactory
	secureClient      SecureClient
	opts              *options.SSHOptions

	localListeners []net.Listener
}

func NewSecureShell(
	secureDialer SecureDialer,
	terminalHelper terminal.TerminalHelper,
	listenerFactory ListenerFactory,
	keepAliveInterval time.Duration,
	appFactory app.AppFactory,
	infoFactory info.InfoFactory,
	credFactory credential.CredentialFactory,
) SecureShell {
	return &secureShell{
		secureDialer:      secureDialer,
		terminalHelper:    terminalHelper,
		listenerFactory:   listenerFactory,
		keepAliveInterval: keepAliveInterval,
		appFactory:        appFactory,
		infoFactory:       infoFactory,
		credFactory:       credFactory,
		localListeners:    []net.Listener{},
	}
}

func (c *secureShell) Connect(opts *options.SSHOptions) error {
	app, err := c.appFactory.Get(opts.AppName)
	if err != nil {
		return err
	}

	info, err := c.infoFactory.Get()
	if err != nil {
		return err
	}

	cred, err := c.credFactory.Get()
	if err != nil {
		return err
	}

	err = c.validateTarget(app, opts)
	if err != nil {
		return err
	}

	clientConfig := &ssh.ClientConfig{
		User: fmt.Sprintf("cf:%s/%d", app.Guid, opts.Instance),
		Auth: []ssh.AuthMethod{
			ssh.Password(cred.Token),
		},
		HostKeyCallback: fingerprintCallback(opts, info.SSHEndpointFingerprint),
	}

	secureClient, err := c.secureDialer.Dial("tcp", info.SSHEndpoint, clientConfig)
	if err != nil {
		return err
	}

	c.secureClient = secureClient
	c.opts = opts
	return nil
}

func (c *secureShell) Close() error {
	for _, listener := range c.localListeners {
		listener.Close()
	}
	return c.secureClient.Close()
}

func (c *secureShell) LocalPortForward() error {
	for _, forwardSpec := range c.opts.ForwardSpecs {
		listener, err := c.listenerFactory.Listen("tcp", forwardSpec.ListenAddress)
		if err != nil {
			return err
		}
		c.localListeners = append(c.localListeners, listener)

		go c.localForwardAcceptLoop(listener, forwardSpec.ConnectAddress)
	}

	return nil
}

func (c *secureShell) localForwardAcceptLoop(listener net.Listener, addr string) {
	defer listener.Close()

	for {
		conn, err := listener.Accept()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				time.Sleep(100 * time.Millisecond)
				continue
			}
			return
		}

		go c.handleForwardConnection(conn, addr)
	}
}

func (c *secureShell) handleForwardConnection(conn net.Conn, targetAddr string) {
	defer conn.Close()

	target, err := c.secureClient.Dial("tcp", targetAddr)
	if err != nil {
		return
	}
	defer target.Close()

	wg := &sync.WaitGroup{}
	wg.Add(2)

	go copyAndClose(wg, conn, target)
	go copyAndClose(wg, target, conn)
	wg.Wait()
}

func copyAndClose(wg *sync.WaitGroup, dest io.WriteCloser, src io.Reader) {
	io.Copy(dest, src)
	dest.Close()
	if wg != nil {
		wg.Done()
	}
}

func (c *secureShell) InteractiveSession() error {
	secureClient := c.secureClient
	opts := c.opts

	session, err := secureClient.NewSession()
	if err != nil {
		return fmt.Errorf("SSH session allocation failed: %s", err.Error())
	}
	defer session.Close()

	stdin, stdout, stderr := c.terminalHelper.StdStreams()

	inPipe, err := session.StdinPipe()
	if err != nil {
		return err
	}

	outPipe, err := session.StdoutPipe()
	if err != nil {
		return err
	}

	errPipe, err := session.StderrPipe()
	if err != nil {
		return err
	}

	stdinFd, stdinIsTerminal := c.terminalHelper.GetFdInfo(stdin)
	stdoutFd, stdoutIsTerminal := c.terminalHelper.GetFdInfo(stdout)

	if c.shouldAllocateTerminal(opts, stdinIsTerminal) {
		modes := ssh.TerminalModes{
			ssh.ECHO:          1,
			ssh.TTY_OP_ISPEED: 115200,
			ssh.TTY_OP_OSPEED: 115200,
		}

		width, height := c.getWindowDimensions(stdoutFd)

		err = session.RequestPty(c.terminalType(), height, width, modes)
		if err != nil {
			return err
		}

		state, err := c.terminalHelper.SetRawTerminal(stdinFd)
		if err == nil {
			defer c.terminalHelper.RestoreTerminal(stdinFd, state)
		}
	}

	if len(opts.Command) != 0 {
		cmd := strings.Join(opts.Command, " ")
		err = session.Start(cmd)
		if err != nil {
			return err
		}
	} else {
		err = session.Shell()
		if err != nil {
			return err
		}
	}

	go copyAndClose(nil, inPipe, stdin)
	go io.Copy(stdout, outPipe)
	go io.Copy(stderr, errPipe)

	if stdoutIsTerminal {
		resized := make(chan os.Signal, 16)

		if runtime.GOOS == "windows" {
			ticker := time.NewTicker(250 * time.Millisecond)
			defer ticker.Stop()

			go func() {
				for _ = range ticker.C {
					resized <- syscall.Signal(-1)
				}
				close(resized)
			}()
		} else {
			signal.Notify(resized, sigwinch.SIGWINCH())
			defer func() { signal.Stop(resized); close(resized) }()
		}

		go c.resize(resized, session, stdoutFd)
	}

	keepaliveStopCh := make(chan struct{})
	defer close(keepaliveStopCh)

	go keepalive(secureClient.Conn(), time.NewTicker(c.keepAliveInterval), keepaliveStopCh)

	return session.Wait()
}

func (c *secureShell) Wait() error {
	return c.secureClient.Wait()
}

func (c *secureShell) validateTarget(app app.App, opts *options.SSHOptions) error {
	if app.State != "STARTED" {
		return fmt.Errorf("Application %q is not in the STARTED state", opts.AppName)
	}

	if !app.Diego {
		return fmt.Errorf("Application %q is not running on Diego", opts.AppName)
	}

	return nil
}

type hostKeyCallback func(hostname string, remote net.Addr, key ssh.PublicKey) error

func fingerprintCallback(opts *options.SSHOptions, expectedFingerprint string) hostKeyCallback {
	if opts.SkipHostValidation {
		return nil
	}

	return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
		switch len(expectedFingerprint) {
		case helpers.SHA1_FINGERPRINT_LENGTH:
			fingerprint := helpers.SHA1Fingerprint(key)
			if fingerprint != expectedFingerprint {
				return fmt.Errorf("Host key verification failed.\n\nThe fingerprint of the received key was %q.", fingerprint)
			}
		case helpers.MD5_FINGERPRINT_LENGTH:
			fingerprint := helpers.MD5Fingerprint(key)
			if fingerprint != expectedFingerprint {
				return fmt.Errorf("Host key verification failed.\n\nThe fingerprint of the received key was %q.", fingerprint)
			}
		case 0:
			fingerprint := helpers.MD5Fingerprint(key)
			return fmt.Errorf("Unable to verify identity of host.\n\nThe fingerprint of the received key was %q.", fingerprint)
		default:
			return errors.New("Unsupported host key fingerprint format")
		}
		return nil
	}
}

func (c *secureShell) shouldAllocateTerminal(opts *options.SSHOptions, stdinIsTerminal bool) bool {
	switch opts.TerminalRequest {
	case options.REQUEST_TTY_FORCE:
		return true
	case options.REQUEST_TTY_NO:
		return false
	case options.REQUEST_TTY_YES:
		return stdinIsTerminal
	case options.REQUEST_TTY_AUTO:
		return len(opts.Command) == 0 && stdinIsTerminal
	default:
		return false
	}
}

func (c *secureShell) resize(resized <-chan os.Signal, session SecureSession, terminalFd uintptr) {
	type resizeMessage struct {
		Width       uint32
		Height      uint32
		PixelWidth  uint32
		PixelHeight uint32
	}

	var previousWidth, previousHeight int

	for _ = range resized {
		width, height := c.getWindowDimensions(terminalFd)

		if width == previousWidth && height == previousHeight {
			continue
		}

		message := resizeMessage{
			Width:  uint32(width),
			Height: uint32(height),
		}

		session.SendRequest("window-change", false, ssh.Marshal(message))

		previousWidth = width
		previousHeight = height
	}
}

func keepalive(conn ssh.Conn, ticker *time.Ticker, stopCh chan struct{}) {
	for {
		select {
		case <-ticker.C:
			conn.SendRequest("keepalive@cloudfoundry.org", true, nil)
		case <-stopCh:
			ticker.Stop()
			return
		}
	}
}

func (c *secureShell) terminalType() string {
	term := os.Getenv("TERM")
	if term == "" {
		term = "xterm"
	}
	return term
}

func (c *secureShell) getWindowDimensions(terminalFd uintptr) (width int, height int) {
	winSize, err := c.terminalHelper.GetWinsize(terminalFd)
	if err != nil {
		winSize = &term.Winsize{
			Width:  80,
			Height: 43,
		}
	}

	return int(winSize.Width), int(winSize.Height)
}

type secureDialer struct{}

func (d *secureDialer) Dial(network string, address string, config *ssh.ClientConfig) (SecureClient, error) {
	client, err := ssh.Dial(network, address, config)
	if err != nil {
		return nil, err
	}

	return &secureClient{client: client}, nil
}

func DefaultSecureDialer() SecureDialer {
	return &secureDialer{}
}

type secureClient struct{ client *ssh.Client }

func (sc *secureClient) Close() error   { return sc.client.Close() }
func (sc *secureClient) Conn() ssh.Conn { return sc.client.Conn }
func (sc *secureClient) Wait() error    { return sc.client.Wait() }
func (sc *secureClient) Dial(n, addr string) (net.Conn, error) {
	return sc.client.Dial(n, addr)
}
func (sc *secureClient) NewSession() (SecureSession, error) {
	return sc.client.NewSession()
}

type listenerFactory struct{}

func (lf *listenerFactory) Listen(network, address string) (net.Listener, error) {
	return net.Listen(network, address)
}

func DefaultListenerFactory() ListenerFactory {
	return &listenerFactory{}
}
<|endoftext|>


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

package channels

import (
	"fmt"
	"regexp"
)

var kValidChannelRegexp *regexp.Regexp

func init() {
	var err error
	kValidChannelRegexp, err = regexp.Compile(`^([-_.@\p{L}\p{Nd}]+|\*)$`)
	if err != nil {
		panic("Bad IsValidChannel regexp")
	}
}

func IsValidChannel(channel string) bool {
	return kValidChannelRegexp.MatchString(channel)
}

// Removes duplicates and invalid channel names.
// If 'starPower' is false, channel "*" is ignored.
// If it's true, channel "*" causes the output to be simply ["*"].
func SimplifyChannels(channels []string, starPower bool) ([]string, error) {
	if len(channels) == 0 {
		return channels, nil
	}
	result := make([]string, 0, len(channels))
	found := map[string]bool{}
	for _, ch := range channels {
		if !IsValidChannel(ch) {
			return nil, fmt.Errorf("Illegal channel name %q", ch)
		} else if ch == "*" {
			if starPower {
				return []string{"*"}, nil
			}
		} else if !found[ch] {
			found[ch] = true
			result = append(result, ch)
		}
	}
	return result, nil
}

// Returns true if the list contains the given channel name.
// A nil list is allowed and treated as an empty array.
func ContainsChannel(list []string, channelName string) bool {
	if list != nil {
		for _, item := range list {
			if item == channelName {
				return true
			}
		}
	}
	return false
}
<commit_msg>Make "+" and "=" legal in channel names<commit_after>//  Copyright (c) 2012-2013 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package channels

import (
	"fmt"
	"regexp"
)

var kValidChannelRegexp *regexp.Regexp

func init() {
	var err error
	kValidChannelRegexp, err = regexp.Compile(`^([-+=_.@\p{L}\p{Nd}]+|\*)$`)
	if err != nil {
		panic("Bad IsValidChannel regexp")
	}
}

func IsValidChannel(channel string) bool {
	return kValidChannelRegexp.MatchString(channel)
}

// Removes duplicates and invalid channel names.
// If 'starPower' is false, channel "*" is ignored.
// If it's true, channel "*" causes the output to be simply ["*"].
func SimplifyChannels(channels []string, starPower bool) ([]string, error) {
	if len(channels) == 0 {
		return channels, nil
	}
	result := make([]string, 0, len(channels))
	found := map[string]bool{}
	for _, ch := range channels {
		if !IsValidChannel(ch) {
			return nil, fmt.Errorf("Illegal channel name %q", ch)
		} else if ch == "*" {
			if starPower {
				return []string{"*"}, nil
			}
		} else if !found[ch] {
			found[ch] = true
			result = append(result, ch)
		}
	}
	return result, nil
}

// Returns true if the list contains the given channel name.
// A nil list is allowed and treated as an empty array.
func ContainsChannel(list []string, channelName string) bool {
	if list != nil {
		for _, item := range list {
			if item == channelName {
				return true
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>package dockerfile // import "github.com/docker/docker/builder/dockerfile"

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

	winio "github.com/Microsoft/go-winio"
	"github.com/docker/docker/pkg/idtools"
	"github.com/docker/docker/pkg/reexec"
	"github.com/docker/docker/pkg/system"
	"github.com/pkg/errors"
	"golang.org/x/sys/windows"
)

var pathBlacklist = map[string]bool{
	"c:\\":        true,
	"c:\\windows": true,
}

func init() {
	reexec.Register("windows-fix-permissions", fixPermissionsReexec)
}

func fixPermissions(source, destination string, identity idtools.Identity, _ bool) error {
	if identity.SID == "" {
		return nil
	}

	cmd := reexec.Command("windows-fix-permissions", source, destination, identity.SID)
	output, err := cmd.CombinedOutput()

	return errors.Wrapf(err, "failed to exec windows-fix-permissions: %s", output)
}

func fixPermissionsReexec() {
	err := fixPermissionsWindows(os.Args[1], os.Args[2], os.Args[3])
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		os.Exit(1)
	}
}

func fixPermissionsWindows(source, destination, SID string) error {

	privileges := []string{winio.SeRestorePrivilege, system.SeTakeOwnershipPrivilege}

	err := winio.EnableProcessPrivileges(privileges)
	if err != nil {
		return err
	}

	defer winio.DisableProcessPrivileges(privileges)

	sid, err := windows.StringToSid(SID)
	if err != nil {
		return err
	}

	// Owners on *nix have read/write/delete/read control and write DAC.
	// Add an ACE that grants this to the user/group specified with the
	// chown option. Currently Windows is not honoring the owner change,
	// however, they are aware of this and it should be fixed at some
	// point.

	sddlString := system.SddlAdministratorsLocalSystem
	sddlString += "(A;OICI;GRGWGXRCWDSD;;;" + SID + ")"

	securityDescriptor, err := winio.SddlToSecurityDescriptor(sddlString)
	if err != nil {
		return err
	}

	var daclPresent uint32
	var daclDefaulted uint32
	var dacl *byte

	err = system.GetSecurityDescriptorDacl(&securityDescriptor[0], &daclPresent, &dacl, &daclDefaulted)
	if err != nil {
		return err
	}

	return system.SetNamedSecurityInfo(windows.StringToUTF16Ptr(destination), system.SE_FILE_OBJECT, system.OWNER_SECURITY_INFORMATION|system.DACL_SECURITY_INFORMATION, sid, nil, dacl, nil)
}

func validateCopySourcePath(imageSource *imageMount, origPath, platform string) error {
	// validate windows paths from other images + LCOW
	if imageSource == nil || platform != "windows" {
		return nil
	}

	origPath = filepath.FromSlash(origPath)
	p := strings.ToLower(filepath.Clean(origPath))
	if !filepath.IsAbs(p) {
		if filepath.VolumeName(p) != "" {
			if p[len(p)-2:] == ":." { // case where clean returns weird c:. paths
				p = p[:len(p)-1]
			}
			p += "\\"
		} else {
			p = filepath.Join("c:\\", p)
		}
	}
	if _, blacklisted := pathBlacklist[p]; blacklisted {
		return errors.New("copy from c:\\ or c:\\windows is not allowed on windows")
	}
	return nil
}
<commit_msg>builder: remove use of deprecated pkg/system constants<commit_after>package dockerfile // import "github.com/docker/docker/builder/dockerfile"

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

	winio "github.com/Microsoft/go-winio"
	"github.com/docker/docker/pkg/idtools"
	"github.com/docker/docker/pkg/reexec"
	"github.com/docker/docker/pkg/system"
	"github.com/pkg/errors"
	"golang.org/x/sys/windows"
)

var pathBlacklist = map[string]bool{
	"c:\\":        true,
	"c:\\windows": true,
}

func init() {
	reexec.Register("windows-fix-permissions", fixPermissionsReexec)
}

func fixPermissions(source, destination string, identity idtools.Identity, _ bool) error {
	if identity.SID == "" {
		return nil
	}

	cmd := reexec.Command("windows-fix-permissions", source, destination, identity.SID)
	output, err := cmd.CombinedOutput()

	return errors.Wrapf(err, "failed to exec windows-fix-permissions: %s", output)
}

func fixPermissionsReexec() {
	err := fixPermissionsWindows(os.Args[1], os.Args[2], os.Args[3])
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		os.Exit(1)
	}
}

func fixPermissionsWindows(source, destination, SID string) error {

	privileges := []string{winio.SeRestorePrivilege, system.SeTakeOwnershipPrivilege}

	err := winio.EnableProcessPrivileges(privileges)
	if err != nil {
		return err
	}

	defer winio.DisableProcessPrivileges(privileges)

	sid, err := windows.StringToSid(SID)
	if err != nil {
		return err
	}

	// Owners on *nix have read/write/delete/read control and write DAC.
	// Add an ACE that grants this to the user/group specified with the
	// chown option. Currently Windows is not honoring the owner change,
	// however, they are aware of this and it should be fixed at some
	// point.

	sddlString := system.SddlAdministratorsLocalSystem
	sddlString += "(A;OICI;GRGWGXRCWDSD;;;" + SID + ")"

	securityDescriptor, err := winio.SddlToSecurityDescriptor(sddlString)
	if err != nil {
		return err
	}

	var daclPresent uint32
	var daclDefaulted uint32
	var dacl *byte

	err = system.GetSecurityDescriptorDacl(&securityDescriptor[0], &daclPresent, &dacl, &daclDefaulted)
	if err != nil {
		return err
	}

	return system.SetNamedSecurityInfo(windows.StringToUTF16Ptr(destination), windows.SE_FILE_OBJECT, windows.OWNER_SECURITY_INFORMATION|windows.DACL_SECURITY_INFORMATION, sid, nil, dacl, nil)
}

func validateCopySourcePath(imageSource *imageMount, origPath, platform string) error {
	// validate windows paths from other images + LCOW
	if imageSource == nil || platform != "windows" {
		return nil
	}

	origPath = filepath.FromSlash(origPath)
	p := strings.ToLower(filepath.Clean(origPath))
	if !filepath.IsAbs(p) {
		if filepath.VolumeName(p) != "" {
			if p[len(p)-2:] == ":." { // case where clean returns weird c:. paths
				p = p[:len(p)-1]
			}
			p += "\\"
		} else {
			p = filepath.Join("c:\\", p)
		}
	}
	if _, blacklisted := pathBlacklist[p]; blacklisted {
		return errors.New("copy from c:\\ or c:\\windows is not allowed on windows")
	}
	return nil
}
<|endoftext|>