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

// Package boomer provides commands to run load tests and display results.
package boomer

import (
	"crypto/tls"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/rakyll/pb"
)

// client is the http.Client that will be used to make all requests
// to the destination.
var client *http.Client

type result struct {
	err           error
	statusCode    int
	duration      time.Duration
	contentLength int64
}

type ReqOpts struct {
	Method   string
	URL      string
	Header   http.Header
	Body     string
	Username string
	Password string
	ReadAll  bool
}

// Creates a req object from req options
func (r *ReqOpts) Request() *http.Request {
	req, _ := http.NewRequest(r.Method, r.URL, strings.NewReader(r.Body))
	req.Header = r.Header
	if r.Username != "" && r.Password != "" {
		req.SetBasicAuth(r.Username, r.Password)
	}
	return req
}

type Boomer struct {
	// Req represents the options of the request to be made.
	// TODO(jbd): Make it work with an http.Request instead.
	Req *ReqOpts

	// N is the total number of requests to make.
	N int

	// C is the concurrency level, the number of concurrent workers to run.
	C int

	// Timeout in seconds.
	Timeout int

	// Qps is the rate limit.
	Qps int

	// AllowInsecure is an option to allow insecure TLS/SSL certificates.
	AllowInsecure bool

	// DisableCompression is an option to disable compression in response
	DisableCompression bool

	// DisableKeepAlives is an option to prevents re-use of TCP connections between different HTTP requests
	DisableKeepAlives bool

	// Output represents the output type. If "csv" is provided, the
	// output will be dumped as a csv stream.
	Output string

	// ProxyAddr is the address of HTTP proxy server in the format on "host:port".
	// Optional.
	ProxyAddr *url.URL

	bar     *pb.ProgressBar
	results chan *result
}

func (b *Boomer) startProgress() {
	if b.Output != "" {
		return
	}
	b.bar = pb.New(b.N)
	b.bar.Format("Bom !")
	b.bar.Start()
}

func (b *Boomer) finalizeProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Finish()
}

func (b *Boomer) incProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Increment()
}

// Run makes all the requests, prints the summary. It blocks until
// all work is done.
func (b *Boomer) Run() {
	b.results = make(chan *result, b.N)
	b.startProgress()

	start := time.Now()
	b.runWorkers()
	b.finalizeProgress()

	newReport(b.N, b.results, b.Output, time.Now().Sub(start)).finalize()
	close(b.results)
}

func (b *Boomer) runWorker(wg *sync.WaitGroup, ch chan *http.Request, readAll bool) {
	for req := range ch {
		s := time.Now()

		var code int
		var size int64

		resp, err := client.Do(req)
		if err == nil {
			size = resp.ContentLength
			code = resp.StatusCode
			if readAll {
				_, err = io.Copy(ioutil.Discard, resp.Body)
			}
			resp.Body.Close()
		}

		wg.Done()
		b.incProgress()
		b.results <- &result{
			statusCode:    code,
			duration:      time.Now().Sub(s),
			err:           err,
			contentLength: size,
		}
	}
}

func (b *Boomer) runWorkers() {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: b.AllowInsecure,
		},
		DisableCompression: b.DisableCompression,
		DisableKeepAlives:  b.DisableKeepAlives,
		// TODO(jbd): Add dial timeout.
		TLSHandshakeTimeout: time.Duration(b.Timeout) * time.Millisecond,
		Proxy:               http.ProxyURL(b.ProxyAddr),
	}
	client = &http.Client{Transport: tr}

	var wg sync.WaitGroup
	wg.Add(b.N)

	var throttle <-chan time.Time
	if b.Qps > 0 {
		throttle = time.Tick(time.Duration(1e6/(b.Qps)) * time.Microsecond)
	}
	jobs := make(chan *http.Request, b.N)
	for i := 0; i < b.C; i++ {
		go func() {
			b.runWorker(&wg, jobs, b.Req.ReadAll)
		}()
	}
	for i := 0; i < b.N; i++ {
		if b.Qps > 0 {
			<-throttle
		}
		jobs <- b.Req.Request()
	}
	close(jobs)
	wg.Wait()
}
<commit_msg>Readability improvements.<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package boomer provides commands to run load tests and display results.
package boomer

import (
	"crypto/tls"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/rakyll/pb"
)

// client is the http.Client that will be used to make all requests
// to the destination.
var client *http.Client

type result struct {
	err           error
	statusCode    int
	duration      time.Duration
	contentLength int64
}

type ReqOpts struct {
	Method   string
	URL      string
	Header   http.Header
	Body     string
	Username string
	Password string
	ReadAll  bool
}

// Creates a req object from req options
func (r *ReqOpts) Request() *http.Request {
	req, _ := http.NewRequest(r.Method, r.URL, strings.NewReader(r.Body))
	req.Header = r.Header
	if r.Username != "" && r.Password != "" {
		req.SetBasicAuth(r.Username, r.Password)
	}
	return req
}

type Boomer struct {
	// Req represents the options of the request to be made.
	// TODO(jbd): Make it work with an http.Request instead.
	Req *ReqOpts

	// N is the total number of requests to make.
	N int

	// C is the concurrency level, the number of concurrent workers to run.
	C int

	// Timeout in seconds.
	Timeout int

	// Qps is the rate limit.
	Qps int

	// AllowInsecure is an option to allow insecure TLS/SSL certificates.
	AllowInsecure bool

	// DisableCompression is an option to disable compression in response
	DisableCompression bool

	// DisableKeepAlives is an option to prevents re-use of TCP connections between different HTTP requests
	DisableKeepAlives bool

	// Output represents the output type. If "csv" is provided, the
	// output will be dumped as a csv stream.
	Output string

	// ProxyAddr is the address of HTTP proxy server in the format on "host:port".
	// Optional.
	ProxyAddr *url.URL

	bar     *pb.ProgressBar
	results chan *result
}

func (b *Boomer) startProgress() {
	if b.Output != "" {
		return
	}
	b.bar = pb.New(b.N)
	b.bar.Format("Bom !")
	b.bar.Start()
}

func (b *Boomer) finalizeProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Finish()
}

func (b *Boomer) incProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Increment()
}

// Run makes all the requests, prints the summary. It blocks until
// all work is done.
func (b *Boomer) Run() {
	b.results = make(chan *result, b.N)
	b.startProgress()

	start := time.Now()
	b.runWorkers()
	b.finalizeProgress()

	newReport(b.N, b.results, b.Output, time.Now().Sub(start)).finalize()
	close(b.results)
}

func (b *Boomer) runWorker(wg *sync.WaitGroup, ch chan *http.Request, readAll bool) {
	for req := range ch {
		s := time.Now()

		var code int
		var size int64

		resp, err := client.Do(req)
		if err == nil {
			size = resp.ContentLength
			code = resp.StatusCode
			if readAll {
				_, err = io.Copy(ioutil.Discard, resp.Body)
			}
			resp.Body.Close()
		}

		wg.Done()
		b.incProgress()
		b.results <- &result{
			statusCode:    code,
			duration:      time.Now().Sub(s),
			err:           err,
			contentLength: size,
		}
	}
}

func (b *Boomer) runWorkers() {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: b.AllowInsecure,
		},
		DisableCompression: b.DisableCompression,
		DisableKeepAlives:  b.DisableKeepAlives,
		// TODO(jbd): Add dial timeout.
		TLSHandshakeTimeout: time.Duration(b.Timeout) * time.Millisecond,
		Proxy:               http.ProxyURL(b.ProxyAddr),
	}
	client = &http.Client{Transport: tr}

	var wg sync.WaitGroup
	wg.Add(b.N)

	var throttle <-chan time.Time
	if b.Qps > 0 {
		throttle = time.Tick(time.Duration(1e6/(b.Qps)) * time.Microsecond)
	}

	jobsch := make(chan *http.Request, b.N)
	for i := 0; i < b.C; i++ {
		go b.runWorker(&wg, jobsch, b.Req.ReadAll)
	}

	for i := 0; i < b.N; i++ {
		if b.Qps > 0 {
			<-throttle
		}
		jobsch <- b.Req.Request()
	}
	close(jobsch)

	wg.Wait()
}
<|endoftext|>


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

// Verifies compiler-internal data structures.

package Verifier

import Utils "utils"
import Scanner "scanner"
import Globals "globals"
import Object "object"
import Type "type"
import Universe "universe"
import Import "import"
import AST "ast"

func Error(msg string) {
	panic("internal compiler error: ", msg, "\n");
}

type Verifier struct {
	comp *Globals.Compilation;

	// various sets for marking the graph (and thus avoid cycles)
	objs *map[*Globals.Object] bool;
	typs *map[*Globals.Type] bool;
	pkgs *map[*Globals.Package] bool;
}

func (V *Verifier) VerifyObject(obj *Globals.Object, pnolev int);

func (V *Verifier) VerifyType(typ *Globals.Type) {
	if V.typs[typ] {
		return;  // already verified
	}
	V.typs[typ] = true;

	if typ.obj != nil {
		V.VerifyObject(typ.obj, 0);
	}

	switch typ.form {
	case Type.VOID:
		break;  // TODO for now - remove eventually
	case Type.FORWARD:
		if typ.scope == nil {
			Error("forward types must have a scope");
		}
		break;
	case Type.NIL:
		break;
	case Type.BOOL:
		break;
	case Type.UINT:
		break;
	case Type.INT:
		break;
	case Type.FLOAT:
		break;
	case Type.STRING:
		break;
	case Type.ANY:
		break;
	case Type.ALIAS:
		break;
	case Type.ARRAY:
		break;
	case Type.STRUCT:
		break;
	case Type.INTERFACE:
		break;
	case Type.MAP:
		break;
	case Type.CHANNEL:
		break;
	case Type.FUNCTION:
		break;
	case Type.POINTER:
		break;
	case Type.REFERENCE:
		break;
	default:
		Error("illegal type form " + Type.FormStr(typ.form));
	}
}

func (V *Verifier) VerifyObject(obj *Globals.Object, pnolev int) {
	if V.objs[obj] {
		return;  // already verified
	}
	V.objs[obj] = true;

	// all objects have a non-nil type
	V.VerifyType(obj.typ);

	switch obj.kind {
	case Object.CONST:
		break;
	case Object.TYPE:
		break;
	case Object.VAR:
		break;
	case Object.FUNC:
		break;
	case Object.PACKAGE:
		break;
	case Object.LABEL:
		break;
	default:
		Error("illegal object kind " + Object.KindStr(obj.kind));
	}
}

func (V *Verifier) VerifyScope(scope *Globals.Scope) {
	for p := scope.entries.first; p != nil; p = p.next {
		V.VerifyObject(p.obj, 0);
	}
}

func (V *Verifier) VerifyPackage(pkg *Globals.Package, pno int) {
	if V.pkgs[pkg] {
		return;  // already verified
	}
	V.pkgs[pkg] = true;

	V.VerifyObject(pkg.obj, pno);
	V.VerifyScope(pkg.scope);
}

func (V *Verifier) Verify(comp *Globals.Compilation) {
	// initialize Verifier
	V.comp = comp;
	V.objs = new(map[*Globals.Object] bool);
	V.typs = new(map[*Globals.Type] bool);
	V.pkgs = new(map[*Globals.Package] bool);

	// verify all packages
	filenames := new(map[string] bool);
	for i := 0; i < comp.pkg_ref; i++ {
		pkg := comp.pkg_list[i];
		// each pkg filename must appear only once
		if filenames[pkg.file_name] {
			Error("package filename present more then once");
		}
		filenames[pkg.file_name] = true;
		V.VerifyPackage(pkg, i);
	}
}

export func Verify(comp *Globals.Compilation) {
	V := new(Verifier);
	V.Verify(comp);
}
<commit_msg>- updated code to work again with latest 6g version<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Verifies compiler-internal data structures.

package Verifier

import Utils "utils"
import Scanner "scanner"
import Globals "globals"
import Object "object"
import Type "type"
import Universe "universe"
import Import "import"
import AST "ast"

func Error(msg string) {
	panic("internal compiler error: ", msg, "\n");
}

type Verifier struct {
	comp *Globals.Compilation;

	// various sets for marking the graph (and thus avoid cycles)
	objs *map[*Globals.Object] bool;
	typs *map[*Globals.Type] bool;
	pkgs *map[*Globals.Package] bool;
}

func (V *Verifier) VerifyObject(obj *Globals.Object, pnolev int);

func (V *Verifier) VerifyType(typ *Globals.Type) {
	if present, ok := V.typs[typ]; present {
		return;  // already verified
	}
	V.typs[typ] = true;

	if typ.obj != nil {
		V.VerifyObject(typ.obj, 0);
	}

	switch typ.form {
	case Type.VOID:
		break;  // TODO for now - remove eventually
	case Type.FORWARD:
		if typ.scope == nil {
			Error("forward types must have a scope");
		}
		break;
	case Type.NIL:
		break;
	case Type.BOOL:
		break;
	case Type.UINT:
		break;
	case Type.INT:
		break;
	case Type.FLOAT:
		break;
	case Type.STRING:
		break;
	case Type.ANY:
		break;
	case Type.ALIAS:
		break;
	case Type.ARRAY:
		break;
	case Type.STRUCT:
		break;
	case Type.INTERFACE:
		break;
	case Type.MAP:
		break;
	case Type.CHANNEL:
		break;
	case Type.FUNCTION:
		break;
	case Type.POINTER:
		break;
	case Type.REFERENCE:
		break;
	default:
		Error("illegal type form " + Type.FormStr(typ.form));
	}
}

func (V *Verifier) VerifyObject(obj *Globals.Object, pnolev int) {
	if present, ok := V.objs[obj]; present {
		return;  // already verified
	}
	V.objs[obj] = true;

	// all objects have a non-nil type
	V.VerifyType(obj.typ);

	switch obj.kind {
	case Object.CONST:
		break;
	case Object.TYPE:
		break;
	case Object.VAR:
		break;
	case Object.FUNC:
		break;
	case Object.PACKAGE:
		break;
	case Object.LABEL:
		break;
	default:
		Error("illegal object kind " + Object.KindStr(obj.kind));
	}
}

func (V *Verifier) VerifyScope(scope *Globals.Scope) {
	for p := scope.entries.first; p != nil; p = p.next {
		V.VerifyObject(p.obj, 0);
	}
}

func (V *Verifier) VerifyPackage(pkg *Globals.Package, pno int) {
	if present, ok := V.pkgs[pkg]; present {
		return;  // already verified
	}
	V.pkgs[pkg] = true;

	V.VerifyObject(pkg.obj, pno);
	V.VerifyScope(pkg.scope);
}

func (V *Verifier) Verify(comp *Globals.Compilation) {
	// initialize Verifier
	V.comp = comp;
	V.objs = new(map[*Globals.Object] bool);
	V.typs = new(map[*Globals.Type] bool);
	V.pkgs = new(map[*Globals.Package] bool);

	// verify all packages
	filenames := new(map[string] bool);
	for i := 0; i < comp.pkg_ref; i++ {
		pkg := comp.pkg_list[i];
		// each pkg filename must appear only once
		if present, ok := filenames[pkg.file_name]; present {
			Error("package filename present more than once");
		}
		filenames[pkg.file_name] = true;
		V.VerifyPackage(pkg, i);
	}
}

export func Verify(comp *Globals.Compilation) {
	V := new(Verifier);
	V.Verify(comp);
}
<|endoftext|>


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

// Package leaktest provides tools to detect leaked goroutines in tests.
// To use it, call "defer leaktest.AfterTest(t)" at the beginning of each
// test that may use goroutines, and add a TestMain function for the package
// which calls leaktest.TestMainWithLeakCheck.
package leaktest

import (
	"fmt"
	"net/http"
	"os"
	"runtime"
	"sort"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

var hasFailed int32 // updated atomially (though not necessary in normal usage)

// TestMainWithLeakCheck is an implementation of TestMain which verifies that
// there are no leaked goroutines at the end of the run (except those created
// by the system which are on a whitelist). Usage:
//
//
// // Adjust the relative path as needed.
// //go:generate ../util/leaktest/add-leaktest.sh *_test.go
//
// func TestMain(m *testing.M) {
//   leaktest.TestMainWithLeakCheck(m)
// }
func TestMainWithLeakCheck(m *testing.M) {
	v := m.Run()
	if v == 0 && goroutineLeaked() {
		os.Exit(1)
	}
	os.Exit(v)
}

func interestingGoroutines() (gs []string) {
	buf := make([]byte, 2<<20)
	buf = buf[:runtime.Stack(buf, true)]
	for _, g := range strings.Split(string(buf), "\n\n") {
		sl := strings.SplitN(g, "\n", 2)
		if len(sl) != 2 {
			continue
		}
		stack := strings.TrimSpace(sl[1])
		if stack == "" ||
			strings.Contains(stack, "created by google.golang.org/grpc.(*Server).serveStreams") ||
			strings.Contains(stack, "created by golang.org/x/net/http2.(*serverConn).processHeaderBlockFragment") ||
			strings.Contains(stack, "created by net.startServer") ||
			strings.Contains(stack, "created by testing.RunTests") ||
			strings.Contains(stack, "closeWriteAndWait") ||
			strings.Contains(stack, "testing.Main(") ||
			// These only show up with GOTRACEBACK=2; Issue 5005 (comment 28)
			strings.Contains(stack, "runtime.goexit") ||
			strings.Contains(stack, "created by runtime.gc") ||
			strings.Contains(stack, "github.com/cockroachdb/cockroach/util/leaktest.interestingGoroutines") ||
			strings.Contains(stack, "runtime.MHeap_Scavenger") ||
			strings.Contains(stack, "github.com/cockroachdb/cockroach/util/log.init") ||
			strings.Contains(stack, "os/signal.loop") {
			continue
		}
		gs = append(gs, stack)
	}
	sort.Strings(gs)
	return
}

// Verify the other tests didn't leave any goroutines running.
func goroutineLeaked() bool {
	if testing.Short() {
		// not counting goroutines for leakage in -short mode
		return false
	}
	var stackCount map[string]int
	for i := 0; i < 8; i++ {
		gs := interestingGoroutines()

		n := 0
		stackCount = make(map[string]int)
		for _, g := range gs {
			stackCount[g]++
			n++
		}

		if n == 0 {
			return false
		}
		time.Sleep(10 * time.Millisecond)
	}
	fmt.Fprintf(os.Stderr, "Too many goroutines running after tests.\n")
	for stack, count := range stackCount {
		fmt.Fprintf(os.Stderr, "%d instances of:\n%s\n", count, stack)
	}
	return true
}

// AfterTest should be called (generally with "defer leaktest.AfterTest(t)")
// from each test which uses goroutines. This waits for all goroutines
// on a blacklist to terminate and provides more precise error reporting
// than TestMainWithLeakCheck alone.
// If a previous test's check has already failed, this is a noop (to avoid
// failing unrelated tests).
func AfterTest(t testing.TB) {
	if atomic.LoadInt32(&hasFailed) > 0 {
		t.Log("prior leak detected, leaktest disabled")
		return
	}
	if r := recover(); r != nil {
		// Don't bother with leaktest if we're recovering from a panic.
		panic(r)
	}
	http.DefaultTransport.(*http.Transport).CloseIdleConnections()
	if testing.Short() || t.Failed() {
		// If a test has failed, chances are it hasn't shut down cleanly, so
		// stop checking for leaks in this run altogether.
		atomic.StoreInt32(&hasFailed, 1)
		return
	}
	if r := recover(); r != nil {
		panic(r)
	}
	var bad string
	badSubstring := map[string]string{
		").readLoop(":                                  "a Transport",
		").writeLoop(":                                 "a Transport",
		"created by net/http/httptest.(*Server).Start": "an httptest.Server",
		"timeoutHandler":                               "a TimeoutHandler",
		"net.(*netFD).connect(":                        "a timing out dial",
		").noteClientGone(":                            "a closenotifier sender",
		"created by net/rpc.NewClientWithCodec":        "an rpc client",
		"created by net/rpc.(*Server.ServeCodec)":      "an rpc server connection",
		"(*Store).Start":                               "a store",
		"(*Range).Send":                                "a range command",
	}
	var stacks string
	for i := 0; i < 8; i++ {
		bad = ""
		stacks = strings.Join(interestingGoroutines(), "\n\n")
		for substr, what := range badSubstring {
			if strings.Contains(stacks, substr) {
				bad = what
			}
		}
		if bad == "" {
			return
		}
		// Bad stuff found, but goroutines might just still be
		// shutting down, so give it some time.
		time.Sleep(10 * time.Millisecond)
	}
	atomic.StoreInt32(&hasFailed, 1)
	t.Errorf("Test appears to have leaked %s:\n%s", bad, stacks)
}
<commit_msg>leaktest: remove grpc exemptions<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package leaktest provides tools to detect leaked goroutines in tests.
// To use it, call "defer leaktest.AfterTest(t)" at the beginning of each
// test that may use goroutines, and add a TestMain function for the package
// which calls leaktest.TestMainWithLeakCheck.
package leaktest

import (
	"fmt"
	"net/http"
	"os"
	"runtime"
	"sort"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

var hasFailed int32 // updated atomially (though not necessary in normal usage)

// TestMainWithLeakCheck is an implementation of TestMain which verifies that
// there are no leaked goroutines at the end of the run (except those created
// by the system which are on a whitelist). Usage:
//
//
// // Adjust the relative path as needed.
// //go:generate ../util/leaktest/add-leaktest.sh *_test.go
//
// func TestMain(m *testing.M) {
//   leaktest.TestMainWithLeakCheck(m)
// }
func TestMainWithLeakCheck(m *testing.M) {
	v := m.Run()
	if v == 0 && goroutineLeaked() {
		os.Exit(1)
	}
	os.Exit(v)
}

func interestingGoroutines() (gs []string) {
	buf := make([]byte, 2<<20)
	buf = buf[:runtime.Stack(buf, true)]
	for _, g := range strings.Split(string(buf), "\n\n") {
		sl := strings.SplitN(g, "\n", 2)
		if len(sl) != 2 {
			continue
		}
		stack := strings.TrimSpace(sl[1])
		if stack == "" ||
			strings.Contains(stack, "created by net.startServer") ||
			strings.Contains(stack, "created by testing.RunTests") ||
			strings.Contains(stack, "closeWriteAndWait") ||
			strings.Contains(stack, "testing.Main(") ||
			// These only show up with GOTRACEBACK=2; Issue 5005 (comment 28)
			strings.Contains(stack, "runtime.goexit") ||
			strings.Contains(stack, "created by runtime.gc") ||
			strings.Contains(stack, "github.com/cockroachdb/cockroach/util/leaktest.interestingGoroutines") ||
			strings.Contains(stack, "runtime.MHeap_Scavenger") ||
			strings.Contains(stack, "github.com/cockroachdb/cockroach/util/log.init") ||
			strings.Contains(stack, "os/signal.loop") {
			continue
		}
		gs = append(gs, stack)
	}
	sort.Strings(gs)
	return
}

// Verify the other tests didn't leave any goroutines running.
func goroutineLeaked() bool {
	if testing.Short() {
		// not counting goroutines for leakage in -short mode
		return false
	}
	var stackCount map[string]int
	for i := 0; i < 8; i++ {
		gs := interestingGoroutines()

		n := 0
		stackCount = make(map[string]int)
		for _, g := range gs {
			stackCount[g]++
			n++
		}

		if n == 0 {
			return false
		}
		time.Sleep(10 * time.Millisecond)
	}
	fmt.Fprintf(os.Stderr, "Too many goroutines running after tests.\n")
	for stack, count := range stackCount {
		fmt.Fprintf(os.Stderr, "%d instances of:\n%s\n", count, stack)
	}
	return true
}

// AfterTest should be called (generally with "defer leaktest.AfterTest(t)")
// from each test which uses goroutines. This waits for all goroutines
// on a blacklist to terminate and provides more precise error reporting
// than TestMainWithLeakCheck alone.
// If a previous test's check has already failed, this is a noop (to avoid
// failing unrelated tests).
func AfterTest(t testing.TB) {
	if atomic.LoadInt32(&hasFailed) > 0 {
		t.Log("prior leak detected, leaktest disabled")
		return
	}
	if r := recover(); r != nil {
		// Don't bother with leaktest if we're recovering from a panic.
		panic(r)
	}
	http.DefaultTransport.(*http.Transport).CloseIdleConnections()
	if testing.Short() || t.Failed() {
		// If a test has failed, chances are it hasn't shut down cleanly, so
		// stop checking for leaks in this run altogether.
		atomic.StoreInt32(&hasFailed, 1)
		return
	}
	if r := recover(); r != nil {
		panic(r)
	}
	var bad string
	badSubstring := map[string]string{
		").readLoop(":                                  "a Transport",
		").writeLoop(":                                 "a Transport",
		"created by net/http/httptest.(*Server).Start": "an httptest.Server",
		"timeoutHandler":                               "a TimeoutHandler",
		"net.(*netFD).connect(":                        "a timing out dial",
		").noteClientGone(":                            "a closenotifier sender",
		"created by net/rpc.NewClientWithCodec":        "an rpc client",
		"created by net/rpc.(*Server.ServeCodec)":      "an rpc server connection",
		"(*Store).Start":                               "a store",
		"(*Range).Send":                                "a range command",
	}
	var stacks string
	for i := 0; i < 8; i++ {
		bad = ""
		stacks = strings.Join(interestingGoroutines(), "\n\n")
		for substr, what := range badSubstring {
			if strings.Contains(stacks, substr) {
				bad = what
			}
		}
		if bad == "" {
			return
		}
		// Bad stuff found, but goroutines might just still be
		// shutting down, so give it some time.
		time.Sleep(10 * time.Millisecond)
	}
	atomic.StoreInt32(&hasFailed, 1)
	t.Errorf("Test appears to have leaked %s:\n%s", bad, stacks)
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/smancke/guble/server/cluster"

	log "github.com/Sirupsen/logrus"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server/auth"
	"github.com/smancke/guble/store"
)

const (
	overloadedHandleChannelRatio = 0.9
	handleChannelCapacity        = 500
	subscribeChannelCapacity     = 10
	unsubscribeChannelCapacity   = 10
)

// Router interface provides a mechanism for PubSub messaging
type Router interface {
	AccessManager() (auth.AccessManager, error)
	MessageStore() (store.MessageStore, error)
	KVStore() (store.KVStore, error)
	Cluster() *cluster.Cluster

	Fetch(store.FetchRequest) error

	Subscribe(r *Route) (*Route, error)
	Unsubscribe(r *Route)
	HandleMessage(message *protocol.Message) error
}

// Helper struct to pass `Route` to subscription channel and provide a notification channel.
type subRequest struct {
	route *Route
	doneC chan bool
}

type router struct {
	routes       map[protocol.Path][]*Route // mapping the path to the route slice
	handleC      chan *protocol.Message
	subscribeC   chan subRequest
	unsubscribeC chan subRequest
	stopC        chan bool      // Channel that signals stop of the router
	stopping     bool           // Flag: the router is in stopping process and no incoming messages are accepted
	wg           sync.WaitGroup // Add any operation that we need to wait upon here

	accessManager auth.AccessManager
	messageStore  store.MessageStore
	kvStore       store.KVStore
	cluster       *cluster.Cluster
}

// NewRouter returns a pointer to Router
func NewRouter(accessManager auth.AccessManager, messageStore store.MessageStore, kvStore store.KVStore, cluster *cluster.Cluster) Router {
	return &router{
		routes: make(map[protocol.Path][]*Route),

		handleC:      make(chan *protocol.Message, handleChannelCapacity),
		subscribeC:   make(chan subRequest, subscribeChannelCapacity),
		unsubscribeC: make(chan subRequest, unsubscribeChannelCapacity),
		stopC:        make(chan bool, 1),

		accessManager: accessManager,
		messageStore:  messageStore,
		kvStore:       kvStore,
		cluster:       cluster,
	}
}

func (router *router) Start() error {
	router.panicIfInternalDependenciesAreNil()
	resetRouterMetrics()
	logger.Info("Starting router")

	go func() {
		router.wg.Add(1)
		for {
			if router.stopping && router.channelsAreEmpty() {
				router.closeRoutes()
				router.wg.Done()
				return
			}

			func() {
				defer protocol.PanicLogger()

				select {
				case message := <-router.handleC:
					router.routeMessage(message)
					runtime.Gosched()
				case subscriber := <-router.subscribeC:
					router.subscribe(subscriber.route)
					subscriber.doneC <- true
				case unsubscriber := <-router.unsubscribeC:
					router.unsubscribe(unsubscriber.route)
					unsubscriber.doneC <- true
				case <-router.stopC:
					router.stopping = true
				}
			}()
		}
	}()

	return nil
}

// Stop stops the router by closing the stop channel, and waiting on the WaitGroup
func (router *router) Stop() error {
	logger.Info("Stopping router")

	router.stopC <- true
	router.wg.Wait()
	return nil
}

func (router *router) Check() error {
	if router.accessManager == nil || router.messageStore == nil || router.kvStore == nil {
		logger.WithFields(log.Fields{
			"err": ErrServiceNotProvided,
		}).Error("Some mandatory services are not provided")
		return ErrServiceNotProvided
	}
	err := router.messageStore.Check()
	if err != nil {
		logger.WithFields(log.Fields{
			"err": err,
		}).Error("MessageStore check failed")
		return err
	}
	err = router.kvStore.Check()
	if err != nil {
		log.WithFields(log.Fields{
			"module": "router",
			"err":    err,
		}).Error("KVStore check failed")
		return err
	}
	return nil
}

// HandleMessage stores the message in the MessageStore(and gets a new ID for it iff the message was created locally)
// and then passes it to: the internal channel, and asynchronously to the cluster (if available).
func (router *router) HandleMessage(message *protocol.Message) error {
	logger.WithFields(log.Fields{
		"userID": message.UserID,
		"path":   message.Path,
	}).Debug("HandleMessage")

	mTotalMessagesIncoming.Add(1)
	if err := router.isStopping(); err != nil {
		logger.WithFields(log.Fields{
			"err": err,
		}).Error("Router is stopping")
		return err
	}

	if !router.accessManager.IsAllowed(auth.WRITE, message.UserID, message.Path) {
		return &PermissionDeniedError{message.UserID, auth.WRITE, message.Path}
	}

	lenMessage := int64(len(message.Bytes()))
	mTotalMessagesIncomingBytes.Add(lenMessage)

	msgPathPartition := message.Path.Partition()
	if router.cluster == nil || (router.cluster != nil && message.NodeID == router.cluster.Config.ID) {
		// for a new locally-generated message, we need to generate a new message-ID
		txCallback := func(msgId uint64) []byte {
			message.ID = msgId
			message.Time = time.Now().Unix()
			return message.Bytes()
		}
		if err := router.messageStore.StoreTx(msgPathPartition, txCallback); err != nil {
			logger.WithFields(log.Fields{
				"err":          err,
				"msgPartition": msgPathPartition,
			}).Error("Error storing new local message in partition")
			mTotalMessageStoreErrors.Add(1)
			return err
		}
	} else {
		if err := router.messageStore.Store(msgPathPartition, message.ID, message.Bytes()); err != nil {
			logger.WithFields(log.Fields{
				"err":          err,
				"msgPartition": msgPathPartition,
			}).Error("Error storing message from cluster in partition")
			mTotalMessageStoreErrors.Add(1)
			return err
		}
	}
	mTotalMessagesStoredBytes.Add(lenMessage)

	router.handleOverloadedChannel()

	router.handleC <- message

	if router.cluster != nil && message.NodeID == router.cluster.Config.ID {
		go router.cluster.BroadcastMessage(message)
	}

	return nil
}

// Subscribe adds a route to the subscribers. If there is already a route with same Application Id and Path, it will be replaced.
func (router *router) Subscribe(r *Route) (*Route, error) {
	logger.WithFields(log.Fields{
		"accessManager": router.accessManager,
		"userID":        r.UserID,
		"path":          r.Path,
	}).Debug("Subscribe")

	if err := router.isStopping(); err != nil {
		return nil, err
	}

	accessAllowed := router.accessManager.IsAllowed(auth.READ, r.UserID, r.Path)
	if !accessAllowed {
		return r, &PermissionDeniedError{UserID: r.UserID, AccessType: auth.READ, Path: r.Path}
	}
	req := subRequest{
		route: r,
		doneC: make(chan bool),
	}
	router.subscribeC <- req
	<-req.doneC
	return r, nil
}

func (router *router) Unsubscribe(r *Route) {
	logger.WithFields(log.Fields{
		"accessManager": router.accessManager,
		"userID":        r.UserID,
		"path":          r.Path,
	}).Debug("Unsubscribe")

	req := subRequest{
		route: r,
		doneC: make(chan bool),
	}
	router.unsubscribeC <- req
	<-req.doneC
}

func (router *router) subscribe(r *Route) {
	logger.WithFields(log.Fields{"userID": r.UserID, "path": r.Path}).Debug("Internal subscribe")
	mTotalSubscriptionAttempts.Add(1)

	slice, present := router.routes[r.Path]
	var removed bool
	if present {
		// Try to remove, to avoid double subscriptions of the same app
		slice, removed = removeIfMatching(slice, r)
	} else {
		// Path not present yet. Initialize the slice
		slice = make([]*Route, 0, 1)
		router.routes[r.Path] = slice
		mCurrentRoutes.Add(1)
	}
	router.routes[r.Path] = append(slice, r)
	if removed {
		mTotalDuplicateSubscriptionsAttempts.Add(1)
	} else {
		mTotalSubscriptions.Add(1)
		mCurrentSubscriptions.Add(1)
	}
}

func (router *router) unsubscribe(r *Route) {
	logger.WithFields(log.Fields{"userID": r.UserID, "path": r.Path}).Debug("Internal unsubscribe")
	mTotalUnsubscriptionAttempts.Add(1)

	slice, present := router.routes[r.Path]
	if !present {
		mTotalInvalidTopicOnUnsubscriptionAttempts.Add(1)
		return
	}
	var removed bool
	router.routes[r.Path], removed = removeIfMatching(slice, r)
	if removed {
		mTotalUnsubscriptions.Add(1)
		mCurrentSubscriptions.Add(-1)
	} else {
		mTotalInvalidUnsubscriptionAttempts.Add(1)
	}
	if len(router.routes[r.Path]) == 0 {
		delete(router.routes, r.Path)
		mCurrentRoutes.Add(-1)
	}
}

func (router *router) panicIfInternalDependenciesAreNil() {
	if router.accessManager == nil || router.kvStore == nil || router.messageStore == nil {
		panic(fmt.Sprintf("router: the internal dependencies marked with `true` are not set: AccessManager=%v, KVStore=%v, MessageStore=%v",
			router.accessManager == nil, router.kvStore == nil, router.messageStore == nil))
	}
}

func (router *router) channelsAreEmpty() bool {
	return len(router.handleC) == 0 && len(router.subscribeC) == 0 && len(router.unsubscribeC) == 0
}

func (router *router) isStopping() error {
	if router.stopping {
		return &ModuleStoppingError{"Router"}
	}
	return nil
}

func (router *router) routeMessage(message *protocol.Message) {
	logger.WithFields(log.Fields{
		"msgMetadata": message.Metadata(),
	}).Debug("Called routeMessage for data")
	mTotalMessagesRouted.Add(1)

	matched := false
	for path, pathRoutes := range router.routes {
		if matchesTopic(message.Path, path) {
			matched = true
			for _, route := range pathRoutes {
				if err := route.Deliver(message); err == ErrInvalidRoute {
					// Unsubscribe invalid routes
					router.unsubscribe(route)
				}
			}
		}
	}

	if !matched {
		logger.WithField("topic", message.Path).Debug("No route matched.")
		mTotalMessagesNotMatchingTopic.Add(1)
	}
}

func (router *router) closeRoutes() {
	logger.Debug("Called closeRoutes")

	for _, currentRouteList := range router.routes {
		for _, route := range currentRouteList {
			router.unsubscribe(route)
			log.WithFields(log.Fields{"module": "router", "route": route.String()}).Debug("Closing route")
			route.Close()
		}
	}
}

func (router *router) handleOverloadedChannel() {
	if float32(len(router.handleC))/float32(cap(router.handleC)) > overloadedHandleChannelRatio {
		logger.WithFields(log.Fields{
			"currentLength": len(router.handleC),
			"maxCapacity":   cap(router.handleC),
		}).Warn("handleC channel is almost full")
		mTotalOverloadedHandleChannel.Add(1)
	}
}

// matchesTopic checks whether the supplied routePath matches the message topic
func matchesTopic(messagePath, routePath protocol.Path) bool {
	messagePathLen := len(string(messagePath))
	routePathLen := len(string(routePath))
	return strings.HasPrefix(string(messagePath), string(routePath)) &&
		(messagePathLen == routePathLen ||
			(messagePathLen > routePathLen && string(messagePath)[routePathLen] == '/'))
}

// removeIfMatching removes a route from the supplied list, based on same ApplicationID id and same path (if existing)
// returns: the (possibly updated) slide, and a boolean value (true if route was removed, false otherwise)
func removeIfMatching(slice []*Route, route *Route) ([]*Route, bool) {
	position := -1
	for p, r := range slice {
		if r.ApplicationID == route.ApplicationID && r.Path == route.Path {
			position = p
		}
	}
	if position == -1 {
		return slice, false
	}
	return append(slice[:position], slice[position+1:]...), true
}

// AccessManager returns the `accessManager` provided for the router
func (router *router) AccessManager() (auth.AccessManager, error) {
	if router.accessManager == nil {
		return nil, ErrServiceNotProvided
	}
	return router.accessManager, nil
}

// MessageStore returns the `messageStore` provided for the router
func (router *router) MessageStore() (store.MessageStore, error) {
	if router.messageStore == nil {
		return nil, ErrServiceNotProvided
	}
	return router.messageStore, nil
}

// KVStore returns the `kvStore` provided for the router
func (router *router) KVStore() (store.KVStore, error) {
	if router.kvStore == nil {
		return nil, ErrServiceNotProvided
	}
	return router.kvStore, nil
}

func (router *router) Fetch(req store.FetchRequest) error {
	if err := router.isStopping(); err != nil {
		return err
	}
	router.messageStore.Fetch(req)
	return nil
}

// Cluster returns the `cluster` provided for the router, or nil if no cluster was set-up
func (router *router) Cluster() *cluster.Cluster {
	return router.cluster
}
<commit_msg>improved logging<commit_after>package server

import (
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/smancke/guble/server/cluster"

	log "github.com/Sirupsen/logrus"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server/auth"
	"github.com/smancke/guble/store"
)

const (
	overloadedHandleChannelRatio = 0.9
	handleChannelCapacity        = 500
	subscribeChannelCapacity     = 10
	unsubscribeChannelCapacity   = 10
)

// Router interface provides a mechanism for PubSub messaging
type Router interface {
	AccessManager() (auth.AccessManager, error)
	MessageStore() (store.MessageStore, error)
	KVStore() (store.KVStore, error)
	Cluster() *cluster.Cluster

	Fetch(store.FetchRequest) error

	Subscribe(r *Route) (*Route, error)
	Unsubscribe(r *Route)
	HandleMessage(message *protocol.Message) error
}

// Helper struct to pass `Route` to subscription channel and provide a notification channel.
type subRequest struct {
	route *Route
	doneC chan bool
}

type router struct {
	routes       map[protocol.Path][]*Route // mapping the path to the route slice
	handleC      chan *protocol.Message
	subscribeC   chan subRequest
	unsubscribeC chan subRequest
	stopC        chan bool      // Channel that signals stop of the router
	stopping     bool           // Flag: the router is in stopping process and no incoming messages are accepted
	wg           sync.WaitGroup // Add any operation that we need to wait upon here

	accessManager auth.AccessManager
	messageStore  store.MessageStore
	kvStore       store.KVStore
	cluster       *cluster.Cluster
}

// NewRouter returns a pointer to Router
func NewRouter(accessManager auth.AccessManager, messageStore store.MessageStore, kvStore store.KVStore, cluster *cluster.Cluster) Router {
	return &router{
		routes: make(map[protocol.Path][]*Route),

		handleC:      make(chan *protocol.Message, handleChannelCapacity),
		subscribeC:   make(chan subRequest, subscribeChannelCapacity),
		unsubscribeC: make(chan subRequest, unsubscribeChannelCapacity),
		stopC:        make(chan bool, 1),

		accessManager: accessManager,
		messageStore:  messageStore,
		kvStore:       kvStore,
		cluster:       cluster,
	}
}

func (router *router) Start() error {
	router.panicIfInternalDependenciesAreNil()
	resetRouterMetrics()
	logger.Info("Starting router")

	go func() {
		router.wg.Add(1)
		for {
			if router.stopping && router.channelsAreEmpty() {
				router.closeRoutes()
				router.wg.Done()
				return
			}

			func() {
				defer protocol.PanicLogger()

				select {
				case message := <-router.handleC:
					router.routeMessage(message)
					runtime.Gosched()
				case subscriber := <-router.subscribeC:
					router.subscribe(subscriber.route)
					subscriber.doneC <- true
				case unsubscriber := <-router.unsubscribeC:
					router.unsubscribe(unsubscriber.route)
					unsubscriber.doneC <- true
				case <-router.stopC:
					router.stopping = true
				}
			}()
		}
	}()

	return nil
}

// Stop stops the router by closing the stop channel, and waiting on the WaitGroup
func (router *router) Stop() error {
	logger.Info("Stopping router")

	router.stopC <- true
	router.wg.Wait()
	return nil
}

func (router *router) Check() error {
	if router.accessManager == nil || router.messageStore == nil || router.kvStore == nil {
		logger.WithField("err", ErrServiceNotProvided).Error("Some mandatory services are not provided")
		return ErrServiceNotProvided
	}
	err := router.messageStore.Check()
	if err != nil {
		logger.WithField("err", err).Error("MessageStore check failed")
		return err
	}
	err = router.kvStore.Check()
	if err != nil {
		log.WithFields(log.Fields{
			"module": "router",
			"err":    err,
		}).Error("KVStore check failed")
		return err
	}
	return nil
}

// HandleMessage stores the message in the MessageStore(and gets a new ID for it iff the message was created locally)
// and then passes it to: the internal channel, and asynchronously to the cluster (if available).
func (router *router) HandleMessage(message *protocol.Message) error {
	logger.WithFields(log.Fields{
		"userID": message.UserID,
		"path":   message.Path,
	}).Debug("HandleMessage")

	mTotalMessagesIncoming.Add(1)
	if err := router.isStopping(); err != nil {
		logger.WithFields(log.Fields{
			"err": err,
		}).Error("Router is stopping")
		return err
	}

	if !router.accessManager.IsAllowed(auth.WRITE, message.UserID, message.Path) {
		return &PermissionDeniedError{message.UserID, auth.WRITE, message.Path}
	}

	lenMessage := int64(len(message.Bytes()))
	mTotalMessagesIncomingBytes.Add(lenMessage)

	msgPathPartition := message.Path.Partition()
	if router.cluster == nil || (router.cluster != nil && message.NodeID == router.cluster.Config.ID) {
		// for a new locally-generated message, we need to generate a new message-ID
		txCallback := func(msgId uint64) []byte {
			message.ID = msgId
			message.Time = time.Now().Unix()
			return message.Bytes()
		}
		if err := router.messageStore.StoreTx(msgPathPartition, txCallback); err != nil {
			logger.WithFields(log.Fields{
				"err":          err,
				"msgPartition": msgPathPartition,
			}).Error("Error storing new local message in partition")
			mTotalMessageStoreErrors.Add(1)
			return err
		}
	} else {
		if err := router.messageStore.Store(msgPathPartition, message.ID, message.Bytes()); err != nil {
			logger.WithFields(log.Fields{
				"err":          err,
				"msgPartition": msgPathPartition,
			}).Error("Error storing message from cluster in partition")
			mTotalMessageStoreErrors.Add(1)
			return err
		}
	}
	mTotalMessagesStoredBytes.Add(lenMessage)

	router.handleOverloadedChannel()

	router.handleC <- message

	if router.cluster != nil && message.NodeID == router.cluster.Config.ID {
		go router.cluster.BroadcastMessage(message)
	}

	return nil
}

// Subscribe adds a route to the subscribers. If there is already a route with same Application Id and Path, it will be replaced.
func (router *router) Subscribe(r *Route) (*Route, error) {
	logger.WithFields(log.Fields{
		"accessManager": router.accessManager,
		"userID":        r.UserID,
		"path":          r.Path,
	}).Debug("Subscribe")

	if err := router.isStopping(); err != nil {
		return nil, err
	}

	accessAllowed := router.accessManager.IsAllowed(auth.READ, r.UserID, r.Path)
	if !accessAllowed {
		return r, &PermissionDeniedError{UserID: r.UserID, AccessType: auth.READ, Path: r.Path}
	}
	req := subRequest{
		route: r,
		doneC: make(chan bool),
	}
	router.subscribeC <- req
	<-req.doneC
	return r, nil
}

func (router *router) Unsubscribe(r *Route) {
	logger.WithFields(log.Fields{
		"accessManager": router.accessManager,
		"userID":        r.UserID,
		"path":          r.Path,
	}).Debug("Unsubscribe")

	req := subRequest{
		route: r,
		doneC: make(chan bool),
	}
	router.unsubscribeC <- req
	<-req.doneC
}

func (router *router) subscribe(r *Route) {
	logger.WithFields(log.Fields{"userID": r.UserID, "path": r.Path}).Debug("Internal subscribe")
	mTotalSubscriptionAttempts.Add(1)

	slice, present := router.routes[r.Path]
	var removed bool
	if present {
		// Try to remove, to avoid double subscriptions of the same app
		slice, removed = removeIfMatching(slice, r)
	} else {
		// Path not present yet. Initialize the slice
		slice = make([]*Route, 0, 1)
		router.routes[r.Path] = slice
		mCurrentRoutes.Add(1)
	}
	router.routes[r.Path] = append(slice, r)
	if removed {
		mTotalDuplicateSubscriptionsAttempts.Add(1)
	} else {
		mTotalSubscriptions.Add(1)
		mCurrentSubscriptions.Add(1)
	}
}

func (router *router) unsubscribe(r *Route) {
	logger.WithFields(log.Fields{"userID": r.UserID, "path": r.Path}).Debug("Internal unsubscribe")
	mTotalUnsubscriptionAttempts.Add(1)

	slice, present := router.routes[r.Path]
	if !present {
		mTotalInvalidTopicOnUnsubscriptionAttempts.Add(1)
		return
	}
	var removed bool
	router.routes[r.Path], removed = removeIfMatching(slice, r)
	if removed {
		mTotalUnsubscriptions.Add(1)
		mCurrentSubscriptions.Add(-1)
	} else {
		mTotalInvalidUnsubscriptionAttempts.Add(1)
	}
	if len(router.routes[r.Path]) == 0 {
		delete(router.routes, r.Path)
		mCurrentRoutes.Add(-1)
	}
}

func (router *router) panicIfInternalDependenciesAreNil() {
	if router.accessManager == nil || router.kvStore == nil || router.messageStore == nil {
		panic(fmt.Sprintf("router: the internal dependencies marked with `true` are not set: AccessManager=%v, KVStore=%v, MessageStore=%v",
			router.accessManager == nil, router.kvStore == nil, router.messageStore == nil))
	}
}

func (router *router) channelsAreEmpty() bool {
	return len(router.handleC) == 0 && len(router.subscribeC) == 0 && len(router.unsubscribeC) == 0
}

func (router *router) isStopping() error {
	if router.stopping {
		return &ModuleStoppingError{"Router"}
	}
	return nil
}

func (router *router) routeMessage(message *protocol.Message) {
	logger.WithFields(log.Fields{
		"msgMetadata": message.Metadata(),
	}).Debug("Called routeMessage for data")
	mTotalMessagesRouted.Add(1)

	matched := false
	for path, pathRoutes := range router.routes {
		if matchesTopic(message.Path, path) {
			matched = true
			for _, route := range pathRoutes {
				if err := route.Deliver(message); err == ErrInvalidRoute {
					// Unsubscribe invalid routes
					router.unsubscribe(route)
				}
			}
		}
	}

	if !matched {
		logger.WithField("topic", message.Path).Debug("No route matched.")
		mTotalMessagesNotMatchingTopic.Add(1)
	}
}

func (router *router) closeRoutes() {
	logger.Debug("Called closeRoutes")

	for _, currentRouteList := range router.routes {
		for _, route := range currentRouteList {
			router.unsubscribe(route)
			log.WithFields(log.Fields{"module": "router", "route": route.String()}).Debug("Closing route")
			route.Close()
		}
	}
}

func (router *router) handleOverloadedChannel() {
	if float32(len(router.handleC))/float32(cap(router.handleC)) > overloadedHandleChannelRatio {
		logger.WithFields(log.Fields{
			"currentLength": len(router.handleC),
			"maxCapacity":   cap(router.handleC),
		}).Warn("handleC channel is almost full")
		mTotalOverloadedHandleChannel.Add(1)
	}
}

// matchesTopic checks whether the supplied routePath matches the message topic
func matchesTopic(messagePath, routePath protocol.Path) bool {
	messagePathLen := len(string(messagePath))
	routePathLen := len(string(routePath))
	return strings.HasPrefix(string(messagePath), string(routePath)) &&
		(messagePathLen == routePathLen ||
			(messagePathLen > routePathLen && string(messagePath)[routePathLen] == '/'))
}

// removeIfMatching removes a route from the supplied list, based on same ApplicationID id and same path (if existing)
// returns: the (possibly updated) slide, and a boolean value (true if route was removed, false otherwise)
func removeIfMatching(slice []*Route, route *Route) ([]*Route, bool) {
	position := -1
	for p, r := range slice {
		if r.ApplicationID == route.ApplicationID && r.Path == route.Path {
			position = p
		}
	}
	if position == -1 {
		return slice, false
	}
	return append(slice[:position], slice[position+1:]...), true
}

// AccessManager returns the `accessManager` provided for the router
func (router *router) AccessManager() (auth.AccessManager, error) {
	if router.accessManager == nil {
		return nil, ErrServiceNotProvided
	}
	return router.accessManager, nil
}

// MessageStore returns the `messageStore` provided for the router
func (router *router) MessageStore() (store.MessageStore, error) {
	if router.messageStore == nil {
		return nil, ErrServiceNotProvided
	}
	return router.messageStore, nil
}

// KVStore returns the `kvStore` provided for the router
func (router *router) KVStore() (store.KVStore, error) {
	if router.kvStore == nil {
		return nil, ErrServiceNotProvided
	}
	return router.kvStore, nil
}

func (router *router) Fetch(req store.FetchRequest) error {
	if err := router.isStopping(); err != nil {
		return err
	}
	router.messageStore.Fetch(req)
	return nil
}

// Cluster returns the `cluster` provided for the router, or nil if no cluster was set-up
func (router *router) Cluster() *cluster.Cluster {
	return router.cluster
}
<|endoftext|>


<commit_before>package autotls

import (
	"context"
	"crypto/tls"
	"log"
	"net/http"

	"golang.org/x/crypto/acme/autocert"
	"golang.org/x/sync/errgroup"
)

func run(ctx context.Context, r http.Handler, domain ...string) error {
	var g errgroup.Group

	s1 := &http.Server{
		Addr:    ":http",
		Handler: http.HandlerFunc(redirect),
	}
	s2 := &http.Server{
		Handler: r,
	}

	g.Go(func() error {
		return s1.ListenAndServe()
	})
	g.Go(func() error {
		return s2.Serve(autocert.NewListener(domain...))
	})

	g.Go(func() error {
		if ctx == nil {
			return nil
		}
		<-ctx.Done()

		var gShutdown errgroup.Group
		gShutdown.Go(func() error {
			return s1.Shutdown(context.Background())
		})
		gShutdown.Go(func() error {
			return s2.Shutdown(context.Background())
		})

		return gShutdown.Wait()
	})
	return g.Wait()
}

// Run support 1-line LetsEncrypt HTTPS servers with graceful shutdown
func RunWithContext(ctx context.Context, r http.Handler, domain ...string) error {
	return run(ctx, r, domain...)
}

// Run support 1-line LetsEncrypt HTTPS servers
func Run(r http.Handler, domain ...string) error {
	return run(nil, r, domain...)
}

// RunWithManager support custom autocert manager
func RunWithManager(r http.Handler, m *autocert.Manager) error {
	return RunWithManagerAndTLSConfig(r, m, *m.TLSConfig())
}

// RunWithManagerAndTLSConfig support custom autocert manager and tls.Config
func RunWithManagerAndTLSConfig(r http.Handler, m *autocert.Manager, tlsc tls.Config) error {
	var g errgroup.Group
	if m.Cache == nil {
		var e error
		m.Cache, e = getCacheDir()
		if e != nil {
			log.Println(e)
		}
	}
	defaultTLSConfig := m.TLSConfig()
	tlsc.GetCertificate = defaultTLSConfig.GetCertificate
	tlsc.NextProtos = defaultTLSConfig.NextProtos
	s := &http.Server{
		Addr:      ":https",
		TLSConfig: &tlsc,
		Handler:   r,
	}
	g.Go(func() error {
		return http.ListenAndServe(":http", m.HTTPHandler(http.HandlerFunc(redirect)))
	})
	g.Go(func() error {
		return s.ListenAndServeTLS("", "")
	})
	return g.Wait()
}

func redirect(w http.ResponseWriter, req *http.Request) {
	target := "https://" + req.Host + req.RequestURI

	http.Redirect(w, req, target, http.StatusMovedPermanently)
}
<commit_msg>chore: add default todo context<commit_after>package autotls

import (
	"context"
	"crypto/tls"
	"log"
	"net/http"

	"golang.org/x/crypto/acme/autocert"
	"golang.org/x/sync/errgroup"
)

type tlsContextKey string

var ctxKey = tlsContextKey("autls")

func run(ctx context.Context, r http.Handler, domain ...string) error {
	var g errgroup.Group

	s1 := &http.Server{
		Addr:    ":http",
		Handler: http.HandlerFunc(redirect),
	}
	s2 := &http.Server{
		Handler: r,
	}

	g.Go(func() error {
		return s1.ListenAndServe()
	})
	g.Go(func() error {
		return s2.Serve(autocert.NewListener(domain...))
	})

	g.Go(func() error {
		if v := ctx.Value(ctxKey); v != nil {
			return nil
		}

		<-ctx.Done()

		var gShutdown errgroup.Group
		gShutdown.Go(func() error {
			return s1.Shutdown(context.Background())
		})
		gShutdown.Go(func() error {
			return s2.Shutdown(context.Background())
		})

		return gShutdown.Wait()
	})
	return g.Wait()
}

// Run support 1-line LetsEncrypt HTTPS servers with graceful shutdown
func RunWithContext(ctx context.Context, r http.Handler, domain ...string) error {
	return run(ctx, r, domain...)
}

// Run support 1-line LetsEncrypt HTTPS servers
func Run(r http.Handler, domain ...string) error {
	ctx := context.WithValue(context.Background(), ctxKey, "done")
	return run(ctx, r, domain...)
}

// RunWithManager support custom autocert manager
func RunWithManager(r http.Handler, m *autocert.Manager) error {
	return RunWithManagerAndTLSConfig(r, m, *m.TLSConfig())
}

// RunWithManagerAndTLSConfig support custom autocert manager and tls.Config
func RunWithManagerAndTLSConfig(r http.Handler, m *autocert.Manager, tlsc tls.Config) error {
	var g errgroup.Group
	if m.Cache == nil {
		var e error
		m.Cache, e = getCacheDir()
		if e != nil {
			log.Println(e)
		}
	}
	defaultTLSConfig := m.TLSConfig()
	tlsc.GetCertificate = defaultTLSConfig.GetCertificate
	tlsc.NextProtos = defaultTLSConfig.NextProtos
	s := &http.Server{
		Addr:      ":https",
		TLSConfig: &tlsc,
		Handler:   r,
	}
	g.Go(func() error {
		return http.ListenAndServe(":http", m.HTTPHandler(http.HandlerFunc(redirect)))
	})
	g.Go(func() error {
		return s.ListenAndServeTLS("", "")
	})
	return g.Wait()
}

func redirect(w http.ResponseWriter, req *http.Request) {
	target := "https://" + req.Host + req.RequestURI

	http.Redirect(w, req, target, http.StatusMovedPermanently)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/fsouza/go-dockerclient"
	"github.com/garyburd/redigo/redis"
)

// DockerClient for running code
var DockerClient *docker.Client

// Runner runs the code
type Runner struct {
	Lang    string `json:"lang"`
	Source  string `json:"source"`
	Version string `json:"version"`
	Timeout int    `json:"timeout"` // How long is the code going to run
}

// Runnerthrottle Limit the max throttle for runner
var Runnerthrottle chan struct{}

// Run the code in the container
func (r *Runner) Run(output messages, conn redis.Conn, uuid string) {
	Runnerthrottle <- struct{}{}
	defer func() { <-Runnerthrottle }()

	container, err := r.createContainer(uuid)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Container %s cannot be created - %v\n", uuid, err)
		return
	}

	stdoutReader, stdoutWriter := io.Pipe()
	stdinReader, stdinWriter := io.Pipe()

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v", err)
	}

	defer stdinWriter.Close()
	defer stdoutWriter.Close()

	go pipeStdin(conn, uuid, stdinWriter)
	go pipeStdout(stdoutReader, output)

	// Start running the container
	err = r.startContainer(container.ID)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Container %s cannot be started - %v\n", uuid, err)
		return
	}
	defer DockerClient.RemoveContainer(docker.RemoveContainerOptions{ID: container.ID, Force: true})

	successChan := make(chan struct{})
	errorChan := make(chan error)

	go func() {
		_, err := DockerClient.WaitContainer(container.ID)
		if err != nil {
			fmt.Printf("Error: %v\n", err)
			errorChan <- err
			return
		}
		successChan <- struct{}{}
	}()

	go func() {
		err = r.attachContainer(container.ID, stdoutWriter, stdinReader)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: Container %s cannot be attached - %v\n", uuid, err)
		}
	}()

	select {
	case <-successChan:
		fmt.Fprintf(os.Stdout, "Container %s is executed successfully\n", uuid)
	case err := <-errorChan:
		fmt.Fprintf(os.Stdout, "Container %s failed caused by - %v\n", uuid, err)
	case <-time.After(time.Duration(r.Timeout) * time.Second):
		msg := fmt.Sprintf("Container %s is terminated caused by 15 sec timeout\n", uuid)
		fmt.Fprintf(os.Stderr, msg)
		output <- msg
	}
}

func pipeStdin(conn redis.Conn, uuid string, stdin *io.PipeWriter) {
	psc := redis.PubSubConn{Conn: conn}
	psc.Subscribe(uuid + "#stdin")

	defer func() {
		psc.Unsubscribe(uuid + "#stdin")
		psc.Close()
		conn.Close()
	}()

StdinSubscriptionLoop:
	for {
		switch n := psc.Receive().(type) {
		case redis.Message:
			fmt.Printf("Message: %s %s\n", n.Channel, n.Data)
			stdin.Write(n.Data)
		case error:
			break StdinSubscriptionLoop
		}
	}
	fmt.Println("Stdin subscription closed")
}

func pipeStdout(stdout *io.PipeReader, output messages) {
	buffer := make([]byte, 512)
	for {
		n, err := stdout.Read(buffer)
		if err != nil {
			if err != io.EOF {
				stdout.Close()
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			}

			close(output)
			break
		}

		data := buffer[0:n]
		output <- string(data)

		// Clear the buffer
		for i := 0; i < n; i++ {
			buffer[i] = 0
		}
	}
}

func getDockerClient() (*docker.Client, error) {
	dockerHost := os.Getenv("DOCKER_HOST")

	// If DOCKER_HOST exists in env, using docker-machine
	if dockerHost != "" {
		return docker.NewClientFromEnv()
	}

	// Otherwise using sock connection
	//TODO: Deal with the TLS case (even though you are not using it for now)
	endpoint := "unix:///var/run/docker.sock"
	return docker.NewClient(endpoint)
}

var imageMapper = map[string]string{
	"ruby":   "koderunr-ruby",
	"python": "koderunr-python",
	"go":     "koderunr-go",
	"c":      "koderunr-c",
	"elixir": "koderunr-erl",
}

func (r *Runner) image() string {
	return imageMapper[r.Lang]
}

func (r *Runner) createContainer(uuid string) (*docker.Container, error) {
	cmd := []string{r.Source, uuid}

	if r.Version != "" {
		cmd = append(cmd, r.Version)
	}
	return DockerClient.CreateContainer(docker.CreateContainerOptions{
		Name: uuid,
		Config: &docker.Config{
			Image:           r.image(),
			NetworkDisabled: true,
			OpenStdin:       true,
			Cmd:             cmd,
		},
	})
}

func (r *Runner) startContainer(containerID string) error {
	return DockerClient.StartContainer(containerID, &docker.HostConfig{
		CPUQuota: 15000,
		Memory:   50 * 1024 * 1024,
	})
}

func (r *Runner) attachContainer(containerID string, stdoutWriter *io.PipeWriter, stdinReader *io.PipeReader) error {
	return DockerClient.AttachToContainer(docker.AttachToContainerOptions{
		Container:    containerID,
		Stdin:        true,
		Stdout:       true,
		Stderr:       true,
		Stream:       true,
		OutputStream: stdoutWriter,
		ErrorStream:  stdoutWriter,
		InputStream:  stdinReader,
	})
}
<commit_msg>Increasing the CPU quota because it's too slow when doing compiling.<commit_after>package main

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/fsouza/go-dockerclient"
	"github.com/garyburd/redigo/redis"
)

// DockerClient for running code
var DockerClient *docker.Client

// Runner runs the code
type Runner struct {
	Lang    string `json:"lang"`
	Source  string `json:"source"`
	Version string `json:"version"`
	Timeout int    `json:"timeout"` // How long is the code going to run
}

// Runnerthrottle Limit the max throttle for runner
var Runnerthrottle chan struct{}

// Run the code in the container
func (r *Runner) Run(output messages, conn redis.Conn, uuid string) {
	Runnerthrottle <- struct{}{}
	defer func() { <-Runnerthrottle }()

	container, err := r.createContainer(uuid)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Container %s cannot be created - %v\n", uuid, err)
		return
	}

	stdoutReader, stdoutWriter := io.Pipe()
	stdinReader, stdinWriter := io.Pipe()

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v", err)
	}

	defer stdinWriter.Close()
	defer stdoutWriter.Close()

	go pipeStdin(conn, uuid, stdinWriter)
	go pipeStdout(stdoutReader, output)

	// Start running the container
	err = r.startContainer(container.ID)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Container %s cannot be started - %v\n", uuid, err)
		return
	}
	defer DockerClient.RemoveContainer(docker.RemoveContainerOptions{ID: container.ID, Force: true})

	successChan := make(chan struct{})
	errorChan := make(chan error)

	go func() {
		_, err := DockerClient.WaitContainer(container.ID)
		if err != nil {
			fmt.Printf("Error: %v\n", err)
			errorChan <- err
			return
		}
		successChan <- struct{}{}
	}()

	go func() {
		err = r.attachContainer(container.ID, stdoutWriter, stdinReader)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: Container %s cannot be attached - %v\n", uuid, err)
		}
	}()

	select {
	case <-successChan:
		fmt.Fprintf(os.Stdout, "Container %s is executed successfully\n", uuid)
	case err := <-errorChan:
		fmt.Fprintf(os.Stdout, "Container %s failed caused by - %v\n", uuid, err)
	case <-time.After(time.Duration(r.Timeout) * time.Second):
		msg := fmt.Sprintf("Container %s is terminated caused by 15 sec timeout\n", uuid)
		fmt.Fprintf(os.Stderr, msg)
		output <- msg
	}
}

func pipeStdin(conn redis.Conn, uuid string, stdin *io.PipeWriter) {
	psc := redis.PubSubConn{Conn: conn}
	psc.Subscribe(uuid + "#stdin")

	defer func() {
		psc.Unsubscribe(uuid + "#stdin")
		psc.Close()
		conn.Close()
	}()

StdinSubscriptionLoop:
	for {
		switch n := psc.Receive().(type) {
		case redis.Message:
			fmt.Printf("Message: %s %s\n", n.Channel, n.Data)
			stdin.Write(n.Data)
		case error:
			break StdinSubscriptionLoop
		}
	}
	fmt.Println("Stdin subscription closed")
}

func pipeStdout(stdout *io.PipeReader, output messages) {
	buffer := make([]byte, 512)
	for {
		n, err := stdout.Read(buffer)
		if err != nil {
			if err != io.EOF {
				stdout.Close()
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			}

			close(output)
			break
		}

		data := buffer[0:n]
		output <- string(data)

		// Clear the buffer
		for i := 0; i < n; i++ {
			buffer[i] = 0
		}
	}
}

func getDockerClient() (*docker.Client, error) {
	dockerHost := os.Getenv("DOCKER_HOST")

	// If DOCKER_HOST exists in env, using docker-machine
	if dockerHost != "" {
		return docker.NewClientFromEnv()
	}

	// Otherwise using sock connection
	//TODO: Deal with the TLS case (even though you are not using it for now)
	endpoint := "unix:///var/run/docker.sock"
	return docker.NewClient(endpoint)
}

var imageMapper = map[string]string{
	"ruby":   "koderunr-ruby",
	"python": "koderunr-python",
	"go":     "koderunr-go",
	"c":      "koderunr-c",
	"elixir": "koderunr-erl",
}

func (r *Runner) image() string {
	return imageMapper[r.Lang]
}

func (r *Runner) createContainer(uuid string) (*docker.Container, error) {
	cmd := []string{r.Source, uuid}

	if r.Version != "" {
		cmd = append(cmd, r.Version)
	}
	return DockerClient.CreateContainer(docker.CreateContainerOptions{
		Name: uuid,
		Config: &docker.Config{
			Image:           r.image(),
			NetworkDisabled: true,
			OpenStdin:       true,
			Cmd:             cmd,
		},
	})
}

func (r *Runner) startContainer(containerID string) error {
	return DockerClient.StartContainer(containerID, &docker.HostConfig{
		CPUQuota: 40000,
		Memory:   50 * 1024 * 1024,
	})
}

func (r *Runner) attachContainer(containerID string, stdoutWriter *io.PipeWriter, stdinReader *io.PipeReader) error {
	return DockerClient.AttachToContainer(docker.AttachToContainerOptions{
		Container:    containerID,
		Stdin:        true,
		Stdout:       true,
		Stderr:       true,
		Stream:       true,
		OutputStream: stdoutWriter,
		ErrorStream:  stdoutWriter,
		InputStream:  stdinReader,
	})
}
<|endoftext|>


<commit_before>package server

import (
	"os"
	"os/signal"
	"syscall"

	log "github.com/golang/glog"
	"github.com/pborman/uuid"
)

type Server interface {
	Config() options
	Init(...Option)
	Handle(Handler) error
	NewHandler(interface{}) Handler
	NewSubscriber(string, interface{}) Subscriber
	Subscribe(Subscriber) error
	Register() error
	Deregister() error
	Start() error
	Stop() error
}

type Option func(*options)

var (
	DefaultAddress        = ":0"
	DefaultName           = "go-server"
	DefaultVersion        = "1.0.0"
	DefaultId             = uuid.NewUUID().String()
	DefaultServer  Server = newRpcServer()
)

// Returns config options for the default service
func Config() options {
	return DefaultServer.Config()
}

// Initialises the default server with options passed in
func Init(opt ...Option) {
	if DefaultServer == nil {
		DefaultServer = newRpcServer(opt...)
	}
	DefaultServer.Init(opt...)
}

// Returns a new server with options passed in
func NewServer(opt ...Option) Server {
	return newRpcServer(opt...)
}

// Creates a new subscriber interface with the given topic
// and handler using the default server
func NewSubscriber(topic string, h interface{}) Subscriber {
	return DefaultServer.NewSubscriber(topic, h)
}

// Creates a new handler interface using the default server
func NewHandler(h interface{}) Handler {
	return DefaultServer.NewHandler(h)
}

// Registers a handler interface with the default server to
// handle inbound requests
func Handle(h Handler) error {
	return DefaultServer.Handle(h)
}

// Registers a subscriber interface with the default server
// which subscribes to specified topic with the broker
func Subscribe(s Subscriber) error {
	return DefaultServer.Subscribe(s)
}

// Registers the default server with the discovery system
func Register() error {
	return DefaultServer.Register()
}

// Deregisters the default server from the discovery system
func Deregister() error {
	return DefaultServer.Deregister()
}

// Blocking run starts the default server and waits for a kill
// signal before exiting. Also registers/deregisters the server
func Run() error {
	if err := Start(); err != nil {
		return err
	}

	if err := DefaultServer.Register(); err != nil {
		return err
	}

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)
	log.Infof("Received signal %s", <-ch)

	if err := DefaultServer.Deregister(); err != nil {
		return err
	}

	return Stop()
}

// Starts the default server
func Start() error {
	config := DefaultServer.Config()
	log.Infof("Starting server %s id %s", config.Name(), config.Id())
	return DefaultServer.Start()
}

// Stops the default server
func Stop() error {
	log.Infof("Stopping server")
	return DefaultServer.Stop()
}
<commit_msg>add comment<commit_after>/*
Server represents a server instance in go-micro which handles synchronous
requests via handlers and asynchronous requests via subscribers that
register with a broker.

The server combines the all the packages in go-micro to create a whole unit
used for building applications including discovery, client/server communication
and pub/sub.
*/
package server

import (
	"os"
	"os/signal"
	"syscall"

	log "github.com/golang/glog"
	"github.com/pborman/uuid"
)

type Server interface {
	Config() options
	Init(...Option)
	Handle(Handler) error
	NewHandler(interface{}) Handler
	NewSubscriber(string, interface{}) Subscriber
	Subscribe(Subscriber) error
	Register() error
	Deregister() error
	Start() error
	Stop() error
}

type Option func(*options)

var (
	DefaultAddress        = ":0"
	DefaultName           = "go-server"
	DefaultVersion        = "1.0.0"
	DefaultId             = uuid.NewUUID().String()
	DefaultServer  Server = newRpcServer()
)

// Returns config options for the default service
func Config() options {
	return DefaultServer.Config()
}

// Initialises the default server with options passed in
func Init(opt ...Option) {
	if DefaultServer == nil {
		DefaultServer = newRpcServer(opt...)
	}
	DefaultServer.Init(opt...)
}

// Returns a new server with options passed in
func NewServer(opt ...Option) Server {
	return newRpcServer(opt...)
}

// Creates a new subscriber interface with the given topic
// and handler using the default server
func NewSubscriber(topic string, h interface{}) Subscriber {
	return DefaultServer.NewSubscriber(topic, h)
}

// Creates a new handler interface using the default server
func NewHandler(h interface{}) Handler {
	return DefaultServer.NewHandler(h)
}

// Registers a handler interface with the default server to
// handle inbound requests
func Handle(h Handler) error {
	return DefaultServer.Handle(h)
}

// Registers a subscriber interface with the default server
// which subscribes to specified topic with the broker
func Subscribe(s Subscriber) error {
	return DefaultServer.Subscribe(s)
}

// Registers the default server with the discovery system
func Register() error {
	return DefaultServer.Register()
}

// Deregisters the default server from the discovery system
func Deregister() error {
	return DefaultServer.Deregister()
}

// Blocking run starts the default server and waits for a kill
// signal before exiting. Also registers/deregisters the server
func Run() error {
	if err := Start(); err != nil {
		return err
	}

	if err := DefaultServer.Register(); err != nil {
		return err
	}

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)
	log.Infof("Received signal %s", <-ch)

	if err := DefaultServer.Deregister(); err != nil {
		return err
	}

	return Stop()
}

// Starts the default server
func Start() error {
	config := DefaultServer.Config()
	log.Infof("Starting server %s id %s", config.Name(), config.Id())
	return DefaultServer.Start()
}

// Stops the default server
func Stop() error {
	log.Infof("Stopping server")
	return DefaultServer.Stop()
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"net/http"

	"github.com/gorilla/mux"
	k8sProxy "github.com/rancher/rancher/k8s/proxy"
	managementapi "github.com/rancher/rancher/pkg/api/server"
	"github.com/rancher/rancher/pkg/auth/providers/publicapi"
	authrequests "github.com/rancher/rancher/pkg/auth/requests"
	"github.com/rancher/rancher/pkg/auth/tokens"
	"github.com/rancher/rancher/pkg/proxy"
	"github.com/rancher/rancher/pkg/remotedialer"
	"github.com/rancher/rancher/pkg/tunnel"
	"github.com/rancher/rancher/server/capabilities"
	"github.com/rancher/rancher/server/ui"
	managementSchema "github.com/rancher/types/apis/management.cattle.io/v3/schema"
	"github.com/rancher/types/config"
)

var (
	whiteList = []string{
		"*.amazonaws.com",
		"*.amazonaws.com.cn",
		"forums.rancher.com",
		"api.exoscale.ch",
		"api.ubiquityhosting.com",
		"api.digitalocean.com",
		"*.otc.t-systems.com",
		"api.profitbricks.com",
	}
)

type Server struct {
	Tunneler *remotedialer.Server
}

func New(ctx context.Context, httpPort, httpsPort int, management *config.ManagementContext) (*Server, error) {
	var result http.Handler
	tokenAPI, err := tokens.NewAPIHandler(ctx, management)
	if err != nil {
		return nil, err
	}

	publicAPI, err := publicapi.NewHandler(ctx, management)
	if err != nil {
		return nil, err
	}

	managementAPI, err := managementapi.New(ctx, httpPort, httpsPort, management, func() http.Handler {
		return result
	})

	if err != nil {
		return nil, err
	}

	k8sProxy, err := k8sProxy.New(management)
	if err != nil {
		return nil, err
	}

	tunnel := tunnel.NewTunneler(management)

	authedAPIs := newAuthed(tokenAPI, managementAPI, k8sProxy)

	authedHandler, err := authrequests.NewAuthenticationFilter(ctx, management, authedAPIs)
	if err != nil {
		return nil, err
	}

	unauthed := mux.NewRouter()
	unauthed.Handle("/", ui.UI(managementAPI))
	unauthed.Handle("/v3/settings/cacerts", authedAPIs).Methods(http.MethodGet)
	unauthed.PathPrefix("/v3-public").Handler(publicAPI)
	unauthed.NotFoundHandler = ui.UI(http.NotFoundHandler())
	unauthed.PathPrefix("/v3/connect").Handler(tunnel)
	unauthed.PathPrefix("/v3").Handler(authedHandler)
	unauthed.PathPrefix("/meta").Handler(authedHandler)
	unauthed.PathPrefix("/k8s/clusters/").Handler(authedHandler)

	uiContent := ui.Content()
	unauthed.PathPrefix("/assets").Handler(uiContent)
	unauthed.PathPrefix("/translations").Handler(uiContent)
	unauthed.Handle("/humans.txt", uiContent)
	unauthed.Handle("/index.html", uiContent)
	unauthed.Handle("/robots.txt", uiContent)
	unauthed.Handle("/VERSION.txt", uiContent)

	registerHealth(unauthed)

	result = unauthed
	return &Server{
		Tunneler: tunnel,
	}, nil
}

func newAuthed(tokenAPI http.Handler, managementAPI http.Handler, k8sproxy http.Handler) *mux.Router {
	authed := mux.NewRouter()
	authed.PathPrefix("/meta/proxy").Handler(newProxy())
	authed.PathPrefix("/meta").Handler(managementAPI)
	authed.PathPrefix("/v3/gkeMachineTypes").Handler(capabilities.NewGKEMachineTypesHandler())
	authed.PathPrefix("/v3/gkeVersions").Handler(capabilities.NewGKEVersionsHandler())
	authed.PathPrefix("/v3/gkeZones").Handler(capabilities.NewGKEZonesHandler())
	authed.PathPrefix("/v3/identit").Handler(tokenAPI)
	authed.PathPrefix("/v3/token").Handler(tokenAPI)
	authed.PathPrefix("/k8s/clusters/").Handler(k8sproxy)
	authed.PathPrefix(managementSchema.Version.Path).Handler(managementAPI)

	return authed
}

func newProxy() http.Handler {
	return proxy.NewProxy("/proxy/", func() []string {
		return whiteList
	})
}
<commit_msg>add packet to whitelist proxy<commit_after>package server

import (
	"context"
	"net/http"

	"github.com/gorilla/mux"
	k8sProxy "github.com/rancher/rancher/k8s/proxy"
	managementapi "github.com/rancher/rancher/pkg/api/server"
	"github.com/rancher/rancher/pkg/auth/providers/publicapi"
	authrequests "github.com/rancher/rancher/pkg/auth/requests"
	"github.com/rancher/rancher/pkg/auth/tokens"
	"github.com/rancher/rancher/pkg/proxy"
	"github.com/rancher/rancher/pkg/remotedialer"
	"github.com/rancher/rancher/pkg/tunnel"
	"github.com/rancher/rancher/server/capabilities"
	"github.com/rancher/rancher/server/ui"
	managementSchema "github.com/rancher/types/apis/management.cattle.io/v3/schema"
	"github.com/rancher/types/config"
)

var (
	whiteList = []string{
		"*.amazonaws.com",
		"*.amazonaws.com.cn",
		"forums.rancher.com",
		"api.exoscale.ch",
		"api.ubiquityhosting.com",
		"api.digitalocean.com",
		"*.otc.t-systems.com",
		"api.profitbricks.com",
		"api.packet.net",
	}
)

type Server struct {
	Tunneler *remotedialer.Server
}

func New(ctx context.Context, httpPort, httpsPort int, management *config.ManagementContext) (*Server, error) {
	var result http.Handler
	tokenAPI, err := tokens.NewAPIHandler(ctx, management)
	if err != nil {
		return nil, err
	}

	publicAPI, err := publicapi.NewHandler(ctx, management)
	if err != nil {
		return nil, err
	}

	managementAPI, err := managementapi.New(ctx, httpPort, httpsPort, management, func() http.Handler {
		return result
	})

	if err != nil {
		return nil, err
	}

	k8sProxy, err := k8sProxy.New(management)
	if err != nil {
		return nil, err
	}

	tunnel := tunnel.NewTunneler(management)

	authedAPIs := newAuthed(tokenAPI, managementAPI, k8sProxy)

	authedHandler, err := authrequests.NewAuthenticationFilter(ctx, management, authedAPIs)
	if err != nil {
		return nil, err
	}

	unauthed := mux.NewRouter()
	unauthed.Handle("/", ui.UI(managementAPI))
	unauthed.Handle("/v3/settings/cacerts", authedAPIs).Methods(http.MethodGet)
	unauthed.PathPrefix("/v3-public").Handler(publicAPI)
	unauthed.NotFoundHandler = ui.UI(http.NotFoundHandler())
	unauthed.PathPrefix("/v3/connect").Handler(tunnel)
	unauthed.PathPrefix("/v3").Handler(authedHandler)
	unauthed.PathPrefix("/meta").Handler(authedHandler)
	unauthed.PathPrefix("/k8s/clusters/").Handler(authedHandler)

	uiContent := ui.Content()
	unauthed.PathPrefix("/assets").Handler(uiContent)
	unauthed.PathPrefix("/translations").Handler(uiContent)
	unauthed.Handle("/humans.txt", uiContent)
	unauthed.Handle("/index.html", uiContent)
	unauthed.Handle("/robots.txt", uiContent)
	unauthed.Handle("/VERSION.txt", uiContent)

	registerHealth(unauthed)

	result = unauthed
	return &Server{
		Tunneler: tunnel,
	}, nil
}

func newAuthed(tokenAPI http.Handler, managementAPI http.Handler, k8sproxy http.Handler) *mux.Router {
	authed := mux.NewRouter()
	authed.PathPrefix("/meta/proxy").Handler(newProxy())
	authed.PathPrefix("/meta").Handler(managementAPI)
	authed.PathPrefix("/v3/gkeMachineTypes").Handler(capabilities.NewGKEMachineTypesHandler())
	authed.PathPrefix("/v3/gkeVersions").Handler(capabilities.NewGKEVersionsHandler())
	authed.PathPrefix("/v3/gkeZones").Handler(capabilities.NewGKEZonesHandler())
	authed.PathPrefix("/v3/identit").Handler(tokenAPI)
	authed.PathPrefix("/v3/token").Handler(tokenAPI)
	authed.PathPrefix("/k8s/clusters/").Handler(k8sproxy)
	authed.PathPrefix(managementSchema.Version.Path).Handler(managementAPI)

	return authed
}

func newProxy() http.Handler {
	return proxy.NewProxy("/proxy/", func() []string {
		return whiteList
	})
}
<|endoftext|>


<commit_before>/*

Copyright 2017 Google 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 serverhelpers

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"strings"

	"github.com/golang/glog"
	"google.golang.org/grpc"
)

var (
	// ListenAddress is the grpc listen address
	ListenAddress = flag.String("listen", "", "GRPC listen address")

	serverCert = flag.String("server-cert", "", "server TLS cert")
	serverKey  = flag.String("server-key", "", "server TLS key")
	clientCA   = flag.String("client-ca", "", "client CA")
)

func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// TODO(tamird): point to merged gRPC code rather than a PR.
		// This is a partial recreation of gRPC's internal checks https://github.com/grpc/grpc-go/pull/514/files#diff-95e9a25b738459a2d3030e1e6fa2a718R61
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			grpcServer.ServeHTTP(w, r)
		} else {
			otherHandler.ServeHTTP(w, r)
		}
	})
}

// ListenAndServe starts grpc server
func ListenAndServe(grpcServer *grpc.Server, otherHandler http.Handler) error {
	lis, err := net.Listen("tcp", *ListenAddress)
	if err != nil {
		return fmt.Errorf("failed to listen: %v", err)
	}

	if *serverCert != "" {
		serverCertKeypair, err := tls.LoadX509KeyPair(*serverCert, *serverKey)
		if err != nil {
			return fmt.Errorf("failed to load server tls cert/key: %v", err)
		}

		var clientCertPool *x509.CertPool
		if *clientCA != "" {
			caCert, err := ioutil.ReadFile(*clientCA)
			if err != nil {
				return fmt.Errorf("failed to load client ca: %v", err)
			}
			clientCertPool = x509.NewCertPool()
			clientCertPool.AppendCertsFromPEM(caCert)
		}

		var h http.Handler
		if otherHandler == nil {
			h = grpcServer
		} else {
			h = grpcHandlerFunc(grpcServer, otherHandler)
		}

		httpsServer := &http.Server{
			Handler: h,
			TLSConfig: &tls.Config{
				Certificates: []tls.Certificate{serverCertKeypair},
				NextProtos:   []string{"h2"},
			},
		}

		if clientCertPool != nil {
			httpsServer.TLSConfig.ClientCAs = clientCertPool
			httpsServer.TLSConfig.ClientAuth = tls.RequireAndVerifyClientCert
		} else {
			glog.Warningf("no client ca provided for grpc server")
		}

		glog.Infof("serving on %v", *ListenAddress)
		err = httpsServer.Serve(tls.NewListener(lis, httpsServer.TLSConfig))
		return fmt.Errorf("failed to serve: %v", err)
	}

	glog.Warningf("serving INSECURE on %v", *ListenAddress)
	err = grpcServer.Serve(lis)
	return fmt.Errorf("failed to serve: %v", err)
}
<commit_msg>Added a helper func to create grpc server with all the goodies and credstore<commit_after>/*

Copyright 2017 Google 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 serverhelpers

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"strings"

	"github.com/golang/glog"
	"github.com/google/credstore/client"
	"github.com/grpc-ecosystem/go-grpc-middleware"
	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
	opentracing "github.com/opentracing/opentracing-go"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
)

var (
	// ListenAddress is the grpc listen address
	ListenAddress = flag.String("listen", "", "GRPC listen address")

	serverCert = flag.String("server-cert", "", "server TLS cert")
	serverKey  = flag.String("server-key", "", "server TLS key")
	clientCA   = flag.String("client-ca", "", "client CA")

	credStoreAddress = flag.String("credstore-address", "", "credstore grpc address")
	credStoreCA      = flag.String("credstore-ca", "", "credstore server ca")
)

func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// TODO(tamird): point to merged gRPC code rather than a PR.
		// This is a partial recreation of gRPC's internal checks https://github.com/grpc/grpc-go/pull/514/files#diff-95e9a25b738459a2d3030e1e6fa2a718R61
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			grpcServer.ServeHTTP(w, r)
		} else {
			otherHandler.ServeHTTP(w, r)
		}
	})
}

// ListenAndServe starts grpc server
func ListenAndServe(grpcServer *grpc.Server, otherHandler http.Handler) error {
	lis, err := net.Listen("tcp", *ListenAddress)
	if err != nil {
		return fmt.Errorf("failed to listen: %v", err)
	}

	if *serverCert != "" {
		serverCertKeypair, err := tls.LoadX509KeyPair(*serverCert, *serverKey)
		if err != nil {
			return fmt.Errorf("failed to load server tls cert/key: %v", err)
		}

		var clientCertPool *x509.CertPool
		if *clientCA != "" {
			caCert, err := ioutil.ReadFile(*clientCA)
			if err != nil {
				return fmt.Errorf("failed to load client ca: %v", err)
			}
			clientCertPool = x509.NewCertPool()
			clientCertPool.AppendCertsFromPEM(caCert)
		}

		var h http.Handler
		if otherHandler == nil {
			h = grpcServer
		} else {
			h = grpcHandlerFunc(grpcServer, otherHandler)
		}

		httpsServer := &http.Server{
			Handler: h,
			TLSConfig: &tls.Config{
				Certificates: []tls.Certificate{serverCertKeypair},
				NextProtos:   []string{"h2"},
			},
		}

		if clientCertPool != nil {
			httpsServer.TLSConfig.ClientCAs = clientCertPool
			httpsServer.TLSConfig.ClientAuth = tls.RequireAndVerifyClientCert
		} else {
			glog.Warningf("no client ca provided for grpc server")
		}

		glog.Infof("serving on %v", *ListenAddress)
		err = httpsServer.Serve(tls.NewListener(lis, httpsServer.TLSConfig))
		return fmt.Errorf("failed to serve: %v", err)
	}

	glog.Warningf("serving INSECURE on %v", *ListenAddress)
	err = grpcServer.Serve(lis)
	return fmt.Errorf("failed to serve: %v", err)
}

// NewServer creates a new GRPC server stub with credstore auth (if requested).
func NewServer() (*grpc.Server, *client.CredstoreClient, error) {
	var grpcServer *grpc.Server
	var cc *client.CredstoreClient

	if *credStoreAddress != "" {
		var err error
		cc, err = client.NewCredstoreClient(context.Background(), *credStoreAddress, *credStoreCA)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to init credstore: %v", err)
		}

		glog.Infof("enabled credstore auth")
		grpcServer = grpc.NewServer(
			grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(
				otgrpc.OpenTracingServerInterceptor(opentracing.GlobalTracer()),
				grpc_prometheus.UnaryServerInterceptor,
				client.CredStoreTokenInterceptor(cc.SigningKey()),
				client.CredStoreMethodAuthInterceptor(),
			)))
	} else {
		grpcServer = grpc.NewServer(
			grpc.UnaryInterceptor(
				otgrpc.OpenTracingServerInterceptor(opentracing.GlobalTracer())))
	}

	reflection.Register(grpcServer)
	grpc_prometheus.Register(grpcServer)

	return grpcServer, cc, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/go-ini/ini"
	"github.com/miekg/dns"
)

var addresses = make(map[string]string)

const dnsSuffix string = "aws."

var awsRegion = flag.String("region", "ap-southeast-2", "AWS region for API access")
var port = flag.Int("port", 10053, "UDP Port to listen for DNS requests on")
var refreshInterval = flag.Int("refresh", 5, "Number of minutes between refreshing hosts")

func main() {
	flag.Parse()
	go updateAddresses()

	server, addr, _, _ := setupServer()
	dns.HandleFunc(dnsSuffix, awsDNSServer)
	wg := &sync.WaitGroup{}
	wg.Add(1)
	defer server.Shutdown()
	fmt.Println("Serving on ", addr)
	wg.Wait()
}

func updateAddresses() {
	ticker := time.NewTicker(time.Duration(*refreshInterval) * time.Minute)
	populateAddresses()
	select {
	case <-ticker.C:
		fmt.Println("Updating addresses.")
		populateAddresses()
	}

}

func setupServer() (*dns.Server, string, chan struct{}, error) {
	pc, err := net.ListenPacket("udp", fmt.Sprintf("127.0.0.1:%d", *port))
	if err != nil {
		panic(err)
		// return nil, "", nil, err
	}
	server := &dns.Server{PacketConn: pc, ReadTimeout: time.Hour, WriteTimeout: time.Hour}

	waitLock := sync.Mutex{}
	waitLock.Lock()
	server.NotifyStartedFunc = waitLock.Unlock

	fin := make(chan struct{}, 0)

	go func() {
		server.ActivateAndServe()
		close(fin)
		pc.Close()
	}()

	waitLock.Lock()
	return server, pc.LocalAddr().String(), fin, nil
}

func awsDNSServer(w dns.ResponseWriter, req *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(req)

	record := strings.ToLower(req.Question[0].Name)
	fmt.Println("DNS Request: ", record)

	// Lookup the address
	if len(addresses[record]) > 0 {
		m.Extra = make([]dns.RR, 1)
		m.Extra[0] = &dns.A{Hdr: dns.RR_Header{Name: m.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0}, A: net.ParseIP(addresses[record])}
		w.WriteMsg(m)
	} else {
		m.SetRcode(req, dns.RcodeNameError)
		m.Authoritative = true

		w.WriteMsg(m)
	}
}

func populateAddresses() {
	for _, profile := range getAvailableAwsProfiles() {
		fmt.Printf("Loading hosts for AWS profile %s\n", strings.ToLower(profile))
		svc := ec2.New(session.New(), &aws.Config{Credentials: credentials.NewSharedCredentials("", strings.ToLower(profile)), Region: aws.String(*awsRegion)})

		// Call the DescribeInstances Operation
		resp, err := svc.DescribeInstances(nil)
		if err != nil {
			fmt.Printf("Unable to load instance details for profile %s: %s\n", profile, err)
			continue
		}

		// resp has all of the response data, pull out instance IDs:
		fmt.Println("> Number of instances: ", len(resp.Reservations))
		for idx := range resp.Reservations {
			for _, inst := range resp.Reservations[idx].Instances {
				name := getNameTagVal(inst.Tags)
				record := strings.ToLower(fmt.Sprintf("%s.%s", *inst.InstanceId, dnsSuffix))
				addresses[record] = *inst.PrivateIpAddress
				fmt.Printf("Added address %s: %s\n", record, *inst.PrivateIpAddress)
				if len(name) != 0 {
					record := strings.ToLower(fmt.Sprintf("%s.%s", parameterizeString(name), dnsSuffix))
					addresses[record] = *inst.PrivateIpAddress
					fmt.Printf("Added address %s: %s\n", record, *inst.PrivateIpAddress)
				}
			}
		}
	}
}

func parameterizeString(input string) string {
	input = strings.ToLower(input)
	re := regexp.MustCompile("[^a-z0-9-_.]+")
	return re.ReplaceAllString(input, "-")
}

func getAvailableAwsProfiles() []string {
	homeDir := os.Getenv("HOME")
	file := filepath.Join(homeDir, ".aws", "credentials")
	config, err := ini.Load(file)
	if err != nil {
		fmt.Printf("Error loading aws credentials file to discover profiles: %s\n", err)
		return []string{"default"}
	}
	return config.SectionStrings()
}

func getNameTagVal(tags []*ec2.Tag) string {
	for _, tag := range tags {
		if *tag.Key == "Name" {
			return *tag.Value
		}
	}

	return ""
}
<commit_msg>Added support for triggering a reload via a DNS query<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/go-ini/ini"
	"github.com/miekg/dns"
)

var addresses = make(map[string]string)

const dnsSuffix string = "aws."

var awsRegion = flag.String("region", "ap-southeast-2", "AWS region for API access")
var port = flag.Int("port", 10053, "UDP Port to listen for DNS requests on")
var refreshInterval = flag.Int("refresh", 5, "Number of minutes between refreshing hosts")

func main() {
	flag.Parse()
	go updateAddresses()

	server, addr, _, _ := setupServer()
	dns.HandleFunc(dnsSuffix, awsDNSServer)
	wg := &sync.WaitGroup{}
	wg.Add(1)
	defer server.Shutdown()
	fmt.Println("Serving on ", addr)
	wg.Wait()
}

func updateAddresses() {
	ticker := time.NewTicker(time.Duration(*refreshInterval) * time.Minute)
	populateAddresses()
	select {
	case <-ticker.C:
		fmt.Println("Updating addresses.")
		populateAddresses()
	}

}

func setupServer() (*dns.Server, string, chan struct{}, error) {
	pc, err := net.ListenPacket("udp", fmt.Sprintf("127.0.0.1:%d", *port))
	if err != nil {
		panic(err)
		// return nil, "", nil, err
	}
	server := &dns.Server{PacketConn: pc, ReadTimeout: time.Hour, WriteTimeout: time.Hour}

	waitLock := sync.Mutex{}
	waitLock.Lock()
	server.NotifyStartedFunc = waitLock.Unlock

	fin := make(chan struct{}, 0)

	go func() {
		server.ActivateAndServe()
		close(fin)
		pc.Close()
	}()

	waitLock.Lock()
	return server, pc.LocalAddr().String(), fin, nil
}

func awsDNSServer(w dns.ResponseWriter, req *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(req)

	record := strings.ToLower(req.Question[0].Name)
	fmt.Println("DNS Request: ", record)

	if record == "reload-me.aws." {
		go populateAddresses()

		m.Extra = make([]dns.RR, 1)
		m.Extra[0] = &dns.TXT{Hdr: dns.RR_Header{Name: m.Question[0].Name, Rrtype: dns.TypeTXT, Class: dns.ClassNONE, Ttl: 0}, Txt: []string{"Reloaded OK"}}
		w.WriteMsg(m)
	}

	// Lookup the address
	if len(addresses[record]) > 0 {
		m.Extra = make([]dns.RR, 1)
		m.Extra[0] = &dns.A{Hdr: dns.RR_Header{Name: m.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0}, A: net.ParseIP(addresses[record])}
		w.WriteMsg(m)
	} else {
		m.SetRcode(req, dns.RcodeNameError)
		m.Authoritative = true

		w.WriteMsg(m)
	}
}

func populateAddresses() {
	for _, profile := range getAvailableAwsProfiles() {
		fmt.Printf("Loading hosts for AWS profile %s\n", strings.ToLower(profile))
		svc := ec2.New(session.New(), &aws.Config{Credentials: credentials.NewSharedCredentials("", strings.ToLower(profile)), Region: aws.String(*awsRegion)})

		// Call the DescribeInstances Operation
		resp, err := svc.DescribeInstances(nil)
		if err != nil {
			fmt.Printf("Unable to load instance details for profile %s: %s\n", profile, err)
			continue
		}

		// resp has all of the response data, pull out instance IDs:
		fmt.Println("> Number of instances: ", len(resp.Reservations))
		for idx := range resp.Reservations {
			for _, inst := range resp.Reservations[idx].Instances {
				name := getNameTagVal(inst.Tags)
				record := strings.ToLower(fmt.Sprintf("%s.%s", *inst.InstanceId, dnsSuffix))
				addresses[record] = *inst.PrivateIpAddress
				fmt.Printf("Added address %s: %s\n", record, *inst.PrivateIpAddress)
				if len(name) != 0 {
					record := strings.ToLower(fmt.Sprintf("%s.%s", parameterizeString(name), dnsSuffix))
					addresses[record] = *inst.PrivateIpAddress
					fmt.Printf("Added address %s: %s\n", record, *inst.PrivateIpAddress)
				}
			}
		}
	}
}

func parameterizeString(input string) string {
	input = strings.ToLower(input)
	re := regexp.MustCompile("[^a-z0-9-_.]+")
	return re.ReplaceAllString(input, "-")
}

func getAvailableAwsProfiles() []string {
	homeDir := os.Getenv("HOME")
	file := filepath.Join(homeDir, ".aws", "credentials")
	config, err := ini.Load(file)
	if err != nil {
		fmt.Printf("Error loading aws credentials file to discover profiles: %s\n", err)
		return []string{"default"}
	}
	return config.SectionStrings()
}

func getNameTagVal(tags []*ec2.Tag) string {
	for _, tag := range tags {
		if *tag.Key == "Name" {
			return *tag.Value
		}
	}

	return ""
}
<|endoftext|>


<commit_before>/*
 Copyright 2017 Crunchy Data Solutions, 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 main

import (
	"io/ioutil"
	"log"
	"os"
	"strings"
)

var logger *log.Logger

/*
this program reads an input file, and removes the master database,
replacing the slave database as the new master, and writes out the file to
the same input filename..it is used to rewrite a pgbouncer.ini file after
a failover has been performed..
*/
func main() {
	logger = log.New(os.Stdout, "logger: ", log.Lshortfile|log.Ldate|log.Ltime)
	var VERSION = os.Getenv("CCP_VERSION")

	logger.Println("bounce " + VERSION + ": starting")

	args := os.Args[1:]
	if len(args) < 1 {
		logger.Println("error:  no file passed on command line")
		os.Exit(2)
	}

	var filename = args[0]
	logger.Println("opening " + filename)

	content, err := ioutil.ReadFile(filename)
	if err != nil {
		logger.Println(err.Error())
		os.Exit(2)
	}
	lines := strings.Split(string(content), "\n")

	//logger.Println("here is the file....")
	var slaveflipped = false
	var sectionfound = false
	var outputString string
	for i := 0; i < len(lines); i++ {
		if strings.Contains(lines[i], "[databases]") {
			//logger.Println("found database section")
			//logger.Println(lines[i])
			outputString += lines[i] + "\n"
			sectionfound = true
		} else {
			if strings.Contains(lines[i], "[") {
				sectionfound = false
			}
			if sectionfound {
				if len(lines[i]) < 1 {
					//logger.Println(lines[i])
					outputString += lines[i] + "\n"
				} else {
					//logger.Println(" process this one-" + lines[i])
					if strings.Contains(lines[i], "master") {
						//logger.Println("dropping this one " + lines[i])
					} else {
						//logger.Println("should be slave " + lines[i])

						//logger.Println("master " + lines[i][pos:length])
						if slaveflipped == false {
							outputString += replace(lines[i])
							slaveflipped = true
						} else {
							outputString += lines[i] + "\n"
						}

					}
				}
			} else {
				//logger.Println(lines[i])
				outputString += lines[i] + "\n"
			}
		}

	}
	err = ioutil.WriteFile(filename, []byte(outputString), 0644)
	if err != nil {
		logger.Println(err.Error())
		os.Exit(2)
	}

}

func replace(instring string) string {
	var pos = strings.Index(instring, "=")
	var length = len(instring)
	//var OSE_PROJECT = os.Getenv("OSE_PROJECT")
	//var PG_MASTER_SERVICE = os.Getenv("PG_MASTER_SERVICE")

	return "master " + instring[pos:length] + "\n"
}
<commit_msg>slaveflipped -> replicaflipped<commit_after>/*
 Copyright 2017 Crunchy Data Solutions, 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 main

import (
	"io/ioutil"
	"log"
	"os"
	"strings"
)

var logger *log.Logger

/*
This program reads an input file, and removes the master database,
replacing the replica database as the new master. It then writes out the file to
the same input filename. It is used to rewrite a pgbouncer.ini file after
a failover has been performed.
*/
func main() {
	logger = log.New(os.Stdout, "logger: ", log.Lshortfile|log.Ldate|log.Ltime)
	var VERSION = os.Getenv("CCP_VERSION")

	logger.Println("bounce " + VERSION + ": starting")

	args := os.Args[1:]
	if len(args) < 1 {
		logger.Println("error:  no file passed on command line")
		os.Exit(2)
	}

	var filename = args[0]
	logger.Println("opening " + filename)

	content, err := ioutil.ReadFile(filename)
	if err != nil {
		logger.Println(err.Error())
		os.Exit(2)
	}
	lines := strings.Split(string(content), "\n")

	//logger.Println("here is the file....")
	var replicaflipped = false
	var sectionfound = false
	var outputString string
	for i := 0; i < len(lines); i++ {
		if strings.Contains(lines[i], "[databases]") {
			//logger.Println("found database section")
			//logger.Println(lines[i])
			outputString += lines[i] + "\n"
			sectionfound = true
		} else {
			if strings.Contains(lines[i], "[") {
				sectionfound = false
			}
			if sectionfound {
				if len(lines[i]) < 1 {
					//logger.Println(lines[i])
					outputString += lines[i] + "\n"
				} else {
					//logger.Println(" process this one-" + lines[i])
					if strings.Contains(lines[i], "master") {
						//logger.Println("dropping this one " + lines[i])
					} else {
						//logger.Println("should be replica " + lines[i])

						//logger.Println("master " + lines[i][pos:length])
						if replicaflipped == false {
							outputString += replace(lines[i])
							replicaflipped = true
						} else {
							outputString += lines[i] + "\n"
						}

					}
				}
			} else {
				//logger.Println(lines[i])
				outputString += lines[i] + "\n"
			}
		}

	}
	err = ioutil.WriteFile(filename, []byte(outputString), 0644)
	if err != nil {
		logger.Println(err.Error())
		os.Exit(2)
	}

}

func replace(instring string) string {
	var pos = strings.Index(instring, "=")
	var length = len(instring)
	//var OSE_PROJECT = os.Getenv("OSE_PROJECT")
	//var PG_MASTER_SERVICE = os.Getenv("PG_MASTER_SERVICE")

	return "master " + instring[pos:length] + "\n"
}
<|endoftext|>