<commit_before>// Copyright 2015 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"database/sql"
	"flag"
	"fmt"
	"io"
	"os"
	"runtime"
	"strings"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/peterh/liner"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/util/errors2"
	"github.com/pingcap/tidb/util/printer"
)

var (
	logLevel = flag.String("L", "error", "log level")
	store    = flag.String("store", "goleveldb", "the name for the registered storage, e.g. memory, goleveldb, boltdb")
	dbPath   = flag.String("dbpath", "test", "db path")

	line        *liner.State
	historyPath = "/tmp/tidb_interpreter"
)

func openHistory() {
	if f, err := os.Open(historyPath); err == nil {
		line.ReadHistory(f)
		f.Close()
	}
}

func saveHistory() {
	if f, err := os.Create(historyPath); err == nil {
		line.WriteHistory(f)
		f.Close()
	}
}

func executeLine(tx *sql.Tx, txnLine string) error {
	if tidb.IsQuery(txnLine) {
		rows, err := tx.Query(txnLine)
		if err != nil {
			return errors.Trace(err)
		}
		defer rows.Close()
		cols, err := rows.Columns()
		if err != nil {
			return errors.Trace(err)
		}

		values := make([][]byte, len(cols))
		scanArgs := make([]interface{}, len(values))
		for i := range values {
			scanArgs[i] = &values[i]
		}

		var datas [][]string
		for rows.Next() {
			err := rows.Scan(scanArgs...)
			if err != nil {
				return errors.Trace(err)
			}

			data := make([]string, len(cols))
			for i, value := range values {
				if value == nil {
					data[i] = "NULL"
				} else {
					data[i] = string(value)
				}
			}

			datas = append(datas, data)
		}

		// For `cols` and `datas[i]` always has the same length,
		// no need to check return validity.
		result, _ := printer.GetPrintResult(cols, datas)
		fmt.Printf("%s", result)

		if err := rows.Err(); err != nil {
			return errors.Trace(err)
		}
	} else {
		// TODO: rows affected and last insert id
		_, err := tx.Exec(txnLine)
		if err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func mayExit(err error, line string) {
	if errors2.ErrorEqual(err, liner.ErrPromptAborted) || errors2.ErrorEqual(err, io.EOF) {
		fmt.Println("\nBye")
		saveHistory()
		os.Exit(0)
	}
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
}

func readStatement(prompt string) (string, error) {
	var ret string
	for {
		l, err := line.Prompt(prompt)
		if err != nil {
			return "", err
		}
		if strings.HasSuffix(l, ";") == false {
			ret += l + "\n"
			prompt = "   -> "
			continue
		}
		return ret + l, nil
	}
}

func main() {
	printer.PrintTiDBInfo()

	flag.Parse()
	log.SetLevelByString(*logLevel)
	// support for signal notify
	runtime.GOMAXPROCS(runtime.NumCPU())

	line = liner.NewLiner()
	defer line.Close()

	line.SetCtrlCAborts(true)
	openHistory()

	mdb, err := sql.Open(tidb.DriverName, *store+"://"+*dbPath)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}

	for {
		l, err := readStatement("tidb> ")
		mayExit(err, l)
		line.AppendHistory(l)

		// if we're in transaction
		if strings.HasPrefix(l, "BEGIN") || strings.HasPrefix(l, "begin") {
			tx, err := mdb.Begin()
			if err != nil {
				log.Error(errors.ErrorStack(err))
				continue
			}
			for {
				txnLine, err := readStatement(">> ")
				mayExit(err, txnLine)
				line.AppendHistory(txnLine)

				if !strings.HasSuffix(txnLine, ";") {
					txnLine += ";"
				}

				if strings.HasPrefix(txnLine, "COMMIT") || strings.HasPrefix(txnLine, "commit") {
					err := tx.Commit()
					if err != nil {
						log.Error(errors.ErrorStack(err))
						tx.Rollback()
					}
					break
				}
				// normal sql statement
				err = executeLine(tx, txnLine)
				if err != nil {
					log.Error(errors.ErrorStack(err))
					tx.Rollback()
					break
				}
			}
		} else {
			tx, err := mdb.Begin()
			if err != nil {
				log.Error(errors.ErrorStack(err))
				continue
			}
			err = executeLine(tx, l)
			if err != nil {
				log.Error(errors.ErrorStack(err))
				tx.Rollback()
			} else {
				tx.Commit()
			}
		}
	}
}
<commit_msg>interpreter: mayExit return boolean instead of Exit(0)<commit_after>// Copyright 2015 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"database/sql"
	"flag"
	"fmt"
	"io"
	"os"
	"runtime"
	"strings"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/peterh/liner"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/util/errors2"
	"github.com/pingcap/tidb/util/printer"
)

var (
	logLevel = flag.String("L", "error", "log level")
	store    = flag.String("store", "goleveldb", "the name for the registered storage, e.g. memory, goleveldb, boltdb")
	dbPath   = flag.String("dbpath", "test", "db path")

	line        *liner.State
	historyPath = "/tmp/tidb_interpreter"
)

func openHistory() {
	if f, err := os.Open(historyPath); err == nil {
		line.ReadHistory(f)
		f.Close()
	}
}

func saveHistory() {
	if f, err := os.Create(historyPath); err == nil {
		line.WriteHistory(f)
		f.Close()
	}
}

func executeLine(tx *sql.Tx, txnLine string) error {
	if tidb.IsQuery(txnLine) {
		rows, err := tx.Query(txnLine)
		if err != nil {
			return errors.Trace(err)
		}
		defer rows.Close()
		cols, err := rows.Columns()
		if err != nil {
			return errors.Trace(err)
		}

		values := make([][]byte, len(cols))
		scanArgs := make([]interface{}, len(values))
		for i := range values {
			scanArgs[i] = &values[i]
		}

		var datas [][]string
		for rows.Next() {
			err := rows.Scan(scanArgs...)
			if err != nil {
				return errors.Trace(err)
			}

			data := make([]string, len(cols))
			for i, value := range values {
				if value == nil {
					data[i] = "NULL"
				} else {
					data[i] = string(value)
				}
			}

			datas = append(datas, data)
		}

		// For `cols` and `datas[i]` always has the same length,
		// no need to check return validity.
		result, _ := printer.GetPrintResult(cols, datas)
		fmt.Printf("%s", result)

		if err := rows.Err(); err != nil {
			return errors.Trace(err)
		}
	} else {
		// TODO: rows affected and last insert id
		_, err := tx.Exec(txnLine)
		if err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func mayExit(err error, l string) bool {
	if errors2.ErrorEqual(err, liner.ErrPromptAborted) || errors2.ErrorEqual(err, io.EOF) {
		fmt.Println("\nBye")
		saveHistory()
		line.Close()
		return true
	}
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	return false
}

func readStatement(prompt string) (string, error) {
	var ret string
	for {
		l, err := line.Prompt(prompt)
		if err != nil {
			return "", err
		}
		if strings.HasSuffix(l, ";") == false {
			ret += l + "\n"
			prompt = "   -> "
			continue
		}
		return ret + l, nil
	}
}

func main() {
	printer.PrintTiDBInfo()

	flag.Parse()
	log.SetLevelByString(*logLevel)
	// support for signal notify
	runtime.GOMAXPROCS(runtime.NumCPU())

	line = liner.NewLiner()
	defer line.Close()

	line.SetCtrlCAborts(true)
	openHistory()

	mdb, err := sql.Open(tidb.DriverName, *store+"://"+*dbPath)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}

	for {
		l, err := readStatement("tidb> ")
		if mayExit(err, l) {
			return
		}
		line.AppendHistory(l)

		// if we're in transaction
		if strings.HasPrefix(l, "BEGIN") || strings.HasPrefix(l, "begin") {
			tx, err := mdb.Begin()
			if err != nil {
				log.Error(errors.ErrorStack(err))
				continue
			}
			for {
				txnLine, err := readStatement(">> ")
				if mayExit(err, txnLine) {
					return
				}
				line.AppendHistory(txnLine)

				if !strings.HasSuffix(txnLine, ";") {
					txnLine += ";"
				}

				if strings.HasPrefix(txnLine, "COMMIT") || strings.HasPrefix(txnLine, "commit") {
					err := tx.Commit()
					if err != nil {
						log.Error(errors.ErrorStack(err))
						tx.Rollback()
					}
					break
				}
				// normal sql statement
				err = executeLine(tx, txnLine)
				if err != nil {
					log.Error(errors.ErrorStack(err))
					tx.Rollback()
					break
				}
			}
		} else {
			tx, err := mdb.Begin()
			if err != nil {
				log.Error(errors.ErrorStack(err))
				continue
			}
			err = executeLine(tx, l)
			if err != nil {
				log.Error(errors.ErrorStack(err))
				tx.Rollback()
			} else {
				tx.Commit()
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"./Godeps/_workspace/src/golang.org/x/net/websocket"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
    "io"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

type Listeners struct {
    connections []*websocket.Conn
    message chan interface{}
}

func (l *Listeners) Broadcast() {
    for {
        jsonData := <- l.message

        if jsonData == nil {
            continue
        }

        buf, err := json.Marshal(jsonData)

        if err != nil {
            fmt.Println(err.Error())
            continue
        }

        for _, ws := range l.connections {
            if ws != nil {
                _, err2 := ws.Write(buf)

                if err2 != nil {
                    fmt.Println(err2.Error())
                }
            }
        }
    }
}

func (l *Listeners) Add(ws *websocket.Conn) {
    l.connections = append(l.connections, ws)
    io.Copy(ws, ws)
}

func main() {
	urlStr := flag.String("url", "http://localhost:8080", "MMS Backend URL")
	flag.Parse()

	mmsUrl, err := url.Parse(*urlStr)

	if err != nil {
        panic("Could not parse url: " + err.Error())
    }

    message := make(chan interface{})
    listeners := Listeners{connections: make([]*websocket.Conn, 0), message: message}

    go listeners.Broadcast()

    go func() {
        http.Handle("/", websocket.Handler(listeners.Add))
        http.ListenAndServe(":12345", nil)
    }()

	proxy := httputil.NewSingleHostReverseProxy(mmsUrl)

	// override director to intercept request body
	oldDirector := proxy.Director
	proxy.Director = func(request *http.Request) {
		fmt.Printf("%s %s %s\n", request.Method, request.URL, request.Header.Get("Content-Type"))

		if strings.Index(request.Header.Get("Content-Type"), "application/json") == 0 {
			buf, _ := ioutil.ReadAll(request.Body)
			var jsonData interface{}
			err := json.Unmarshal(buf, &jsonData)

			if err == nil {
                requestUrlStr := request.URL.String()
                queryParams, _ := url.ParseQuery(request.URL.RawQuery)

                // copy jsonData to websocket listeners
                wrapper := make(map[string]interface{})
                wrapper["remoteAddr"] = request.RemoteAddr
                wrapper["ah"] = queryParams.Get("ah")
                wrapper["type"] = nil

                if strings.Index(requestUrlStr, "/agents/api/automation/metrics") == 0 {
                    wrapper["type"] = "metrics"
                } else if strings.Index(requestUrlStr, "/agents/api/automation/status") == 0 {
                    wrapper["type"] = "status"
                } else if strings.Index(requestUrlStr, "/agents/api/automation/log") == 0 {
                    wrapper["type"] = "log"
                }

                wrapper["content"] = jsonData
                message <- wrapper
				//fmt.Println(jsonData)
			}

			request.Body = ioutil.NopCloser(bytes.NewReader(buf))
		}

		oldDirector(request)
	}

	http.ListenAndServe(":9090", proxy)
}
<commit_msg>Serve static files and collapse all endpoints into same HTTP instance<commit_after>package main

import (
	"./Godeps/_workspace/src/golang.org/x/net/websocket"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
)

type Listeners struct {
	// FIXME this sucks and should be replaced with something the prunes dead/closed connections
	connections []*websocket.Conn
	message     chan interface{}
}

func (l *Listeners) Broadcast() {
	for {
		jsonData := <-l.message

		if jsonData == nil {
			continue
		}

		buf, err := json.Marshal(jsonData)

		if err != nil {
			fmt.Println(err.Error())
			continue
		}

		for _, ws := range l.connections {
			if ws != nil {
				_, err2 := ws.Write(buf)

				if err2 != nil {
					fmt.Println(err2.Error())
				}
			}
		}
	}
}

func (l *Listeners) Add(ws *websocket.Conn) {
	l.connections = append(l.connections, ws)
	io.Copy(ws, ws)
}

type SelectiveProxy struct {
	prefixes       map[string]http.Handler
	defaultHandler http.Handler
}

func (s SelectiveProxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	requestUrlStr := req.URL.String()

	for prefix, handler := range s.prefixes {
		// FIXME this is crappy and should be replaced with a trie or something just better
		if strings.Index(requestUrlStr, prefix) == 0 {
			// call the first handler whose prefix matches the starf of the requestUrlStr
			handler.ServeHTTP(w, req)
			return
		}
	}

	s.defaultHandler.ServeHTTP(w, req)
}

func main() {
	urlStr := flag.String("url", "http://localhost:8080", "MMS Backend URL")
	proxyPort := flag.Int("port", 9090, "Reverse Proxy Port")
	staticPath := flag.String("staticPath", "../webapp/dist", "Cyclops Frontend")
	flag.Parse()

	mmsUrl, err := url.Parse(*urlStr)

	if err != nil {
		panic("Could not parse url: " + err.Error())
	}

	message := make(chan interface{})
	listeners := Listeners{connections: make([]*websocket.Conn, 0), message: message}
	go listeners.Broadcast()
	proxy := httputil.NewSingleHostReverseProxy(mmsUrl)

	// override director to intercept request body
	oldDirector := proxy.Director
	proxy.Director = func(request *http.Request) {
		fmt.Printf("%s %s %s\n", request.Method, request.URL, request.Header.Get("Content-Type"))

		if strings.Index(request.Header.Get("Content-Type"), "application/json") == 0 {
			buf, _ := ioutil.ReadAll(request.Body)
			var jsonData interface{}
			err := json.Unmarshal(buf, &jsonData)

			if err == nil {
				requestUrlStr := request.URL.String()
				queryParams, _ := url.ParseQuery(request.URL.RawQuery)

				// copy jsonData to websocket listeners
				wrapper := make(map[string]interface{})
				wrapper["remoteAddr"] = request.RemoteAddr
				wrapper["ah"] = queryParams.Get("ah")
				wrapper["type"] = nil

				if strings.Index(requestUrlStr, "/agents/api/automation/metrics") == 0 {
					wrapper["type"] = "metrics"
				} else if strings.Index(requestUrlStr, "/agents/api/automation/status") == 0 {
					wrapper["type"] = "status"
				} else if strings.Index(requestUrlStr, "/agents/api/automation/log") == 0 {
					wrapper["type"] = "log"
				}

				wrapper["content"] = jsonData
				message <- wrapper
			}

			request.Body = ioutil.NopCloser(bytes.NewReader(buf))
		}

		oldDirector(request)
	}

	prefixes := make(map[string]http.Handler)
	prefixes["/ws"] = websocket.Handler(listeners.Add)
	prefixes["/agents/api/automation/"] = proxy
	selectiveProxy := SelectiveProxy{prefixes: prefixes, defaultHandler: http.FileServer(http.Dir(*staticPath))}
	http.ListenAndServe(":"+strconv.Itoa(*proxyPort), selectiveProxy)
}
<|endoftext|>


<commit_before>package storage

import (
	"context"

	"github.com/pkg/errors"

	resourcetypes "github.com/projecteru2/core/resources/types"
	"github.com/projecteru2/core/scheduler"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
)

type storageRequest struct {
	request int64
	limit   int64
}

// MakeRequest .
func MakeRequest(opts types.ResourceOptions) (resourcetypes.ResourceRequest, error) {
	sr := &storageRequest{
		request: opts.StorageRequest,
		limit:   opts.StorageLimit,
	}
	// add volume request / limit to storage request / limit
	if len(opts.VolumeRequest) != 0 && len(opts.VolumeLimit) != len(opts.VolumeRequest) {
		return nil, errors.Wrapf(types.ErrBadVolume, "volume request and limit must be the same length")
	}
	for idx := range opts.VolumeLimit {
		if len(opts.VolumeRequest) > 0 {
			sr.request += opts.VolumeRequest[idx].SizeInBytes
			sr.limit += utils.Max(opts.VolumeLimit[idx].SizeInBytes, opts.VolumeRequest[idx].SizeInBytes)
		} else {
			sr.request += opts.VolumeLimit[idx].SizeInBytes
			sr.limit += opts.VolumeLimit[idx].SizeInBytes
		}
	}

	return sr, sr.Validate()
}

// Type .
func (s storageRequest) Type() types.ResourceType {
	return types.ResourceStorage
}

// Validate .
func (s *storageRequest) Validate() error {
	if s.limit < 0 || s.request < 0 {
		return errors.Wrap(types.ErrBadStorage, "storage limit or request less than 0")
	}
	if s.limit > 0 && s.request == 0 {
		s.request = s.limit
	}
	if s.limit > 0 && s.request > 0 && s.request > s.limit {
		s.limit = s.request // softlimit storage size
	}
	return nil
}

// MakeScheduler .
func (s storageRequest) MakeScheduler() resourcetypes.SchedulerV2 {
	return func(ctx context.Context, scheduleInfos []resourcetypes.ScheduleInfo) (plans resourcetypes.ResourcePlans, total int, err error) {
		schedulerV1, err := scheduler.GetSchedulerV1()
		if err != nil {
			return
		}

		scheduleInfos, total, err = schedulerV1.SelectStorageNodes(ctx, scheduleInfos, s.request)
		return ResourcePlans{
			request:  s.request,
			limit:    s.limit,
			capacity: resourcetypes.GetCapacity(scheduleInfos),
		}, total, err
	}
}

// Rate .
func (s storageRequest) Rate(node types.Node) float64 {
	return float64(s.request) / float64(node.InitStorageCap)
}

// ResourcePlans .
type ResourcePlans struct {
	request  int64
	limit    int64
	capacity map[string]int
}

// Type .
func (rp ResourcePlans) Type() types.ResourceType {
	return types.ResourceStorage
}

// Capacity .
func (rp ResourcePlans) Capacity() map[string]int {
	return rp.capacity
}

// ApplyChangesOnNode .
func (rp ResourcePlans) ApplyChangesOnNode(node *types.Node, indices ...int) {
	node.StorageCap -= int64(len(indices)) * rp.request
}

// RollbackChangesOnNode .
func (rp ResourcePlans) RollbackChangesOnNode(node *types.Node, indices ...int) {
	node.StorageCap += int64(len(indices)) * rp.request
}

// Dispense .
func (rp ResourcePlans) Dispense(opts resourcetypes.DispenseOptions, r *types.ResourceMeta) (*types.ResourceMeta, error) {
	if rp.capacity[opts.Node.Name] <= opts.Index {
		return nil, errors.WithStack(types.ErrInsufficientCap)
	}
	r.StorageLimit = rp.limit
	r.StorageRequest = rp.request
	return r, nil
}
<commit_msg>fix storage issue (#575)<commit_after>package storage

import (
	"context"

	"github.com/pkg/errors"

	resourcetypes "github.com/projecteru2/core/resources/types"
	"github.com/projecteru2/core/scheduler"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
)

type storageRequest struct {
	request int64
	limit   int64
}

// MakeRequest .
func MakeRequest(opts types.ResourceOptions) (resourcetypes.ResourceRequest, error) {
	sr := &storageRequest{
		request: opts.StorageRequest,
		limit:   opts.StorageLimit,
	}
	if sr.limit > 0 && sr.request == 0 {
		sr.request = sr.limit
	}
	// add volume request / limit to storage request / limit
	if len(opts.VolumeRequest) != 0 && len(opts.VolumeLimit) != len(opts.VolumeRequest) {
		return nil, errors.Wrapf(types.ErrBadVolume, "volume request and limit must be the same length")
	}
	for idx := range opts.VolumeLimit {
		if len(opts.VolumeRequest) > 0 {
			sr.request += opts.VolumeRequest[idx].SizeInBytes
			sr.limit += utils.Max(opts.VolumeLimit[idx].SizeInBytes, opts.VolumeRequest[idx].SizeInBytes)
		} else {
			sr.request += opts.VolumeLimit[idx].SizeInBytes
			sr.limit += opts.VolumeLimit[idx].SizeInBytes
		}
	}

	return sr, sr.Validate()
}

// Type .
func (s storageRequest) Type() types.ResourceType {
	return types.ResourceStorage
}

// Validate .
func (s *storageRequest) Validate() error {
	if s.limit < 0 || s.request < 0 {
		return errors.Wrap(types.ErrBadStorage, "storage limit or request less than 0")
	}
	if s.limit > 0 && s.request == 0 {
		s.request = s.limit
	}
	if s.limit > 0 && s.request > 0 && s.request > s.limit {
		s.limit = s.request // softlimit storage size
	}
	return nil
}

// MakeScheduler .
func (s storageRequest) MakeScheduler() resourcetypes.SchedulerV2 {
	return func(ctx context.Context, scheduleInfos []resourcetypes.ScheduleInfo) (plans resourcetypes.ResourcePlans, total int, err error) {
		schedulerV1, err := scheduler.GetSchedulerV1()
		if err != nil {
			return
		}

		scheduleInfos, total, err = schedulerV1.SelectStorageNodes(ctx, scheduleInfos, s.request)
		return ResourcePlans{
			request:  s.request,
			limit:    s.limit,
			capacity: resourcetypes.GetCapacity(scheduleInfos),
		}, total, err
	}
}

// Rate .
func (s storageRequest) Rate(node types.Node) float64 {
	return float64(s.request) / float64(node.InitStorageCap)
}

// ResourcePlans .
type ResourcePlans struct {
	request  int64
	limit    int64
	capacity map[string]int
}

// Type .
func (rp ResourcePlans) Type() types.ResourceType {
	return types.ResourceStorage
}

// Capacity .
func (rp ResourcePlans) Capacity() map[string]int {
	return rp.capacity
}

// ApplyChangesOnNode .
func (rp ResourcePlans) ApplyChangesOnNode(node *types.Node, indices ...int) {
	node.StorageCap -= int64(len(indices)) * rp.request
}

// RollbackChangesOnNode .
func (rp ResourcePlans) RollbackChangesOnNode(node *types.Node, indices ...int) {
	node.StorageCap += int64(len(indices)) * rp.request
}

// Dispense .
func (rp ResourcePlans) Dispense(opts resourcetypes.DispenseOptions, r *types.ResourceMeta) (*types.ResourceMeta, error) {
	if rp.capacity[opts.Node.Name] <= opts.Index {
		return nil, errors.WithStack(types.ErrInsufficientCap)
	}
	r.StorageLimit = rp.limit
	r.StorageRequest = rp.request
	return r, nil
}
<|endoftext|>


<commit_before>package jsonclient

import (
	"encoding/json"
	"fmt"
	"math"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strconv"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/google/certificate-transparency/go/testdata"
)

func TestNewJSONClient(t *testing.T) {
	tests := []struct {
		pubKey string
		errstr string
	}{
		{"", "no PEM block"},
		{"bogus", "no PEM block"},
		{testdata.RsaPublicKeyPEM, ""},
		{testdata.EcdsaPublicKeyPEM, ""},
		{testdata.DsaPublicKeyPEM, "Unsupported public key type"},
		{testdata.RsaPublicKeyPEM + "bogus", "extra data found"},
	}
	for _, test := range tests {
		client, err := New("http://127.0.0.1", nil, test.pubKey)
		if test.errstr != "" {
			if err == nil {
				t.Errorf("New()=%p,nil; want error %q", client, test.errstr)
			} else if !strings.Contains(err.Error(), test.errstr) {
				t.Errorf("New()=nil,%q; want error %q", err.Error(), test.errstr)
			}
			continue
		}
		if err != nil {
			t.Errorf("New()=nil,%q; want no error", err.Error())
		} else if client == nil {
			t.Errorf("New()=nil,nil; want client")
		}
	}
}

type TestStruct struct {
	TreeSize  int    `json:"tree_size"`
	Timestamp int    `json:"timestamp"`
	Data      string `json:"data"`
}

type TestParams struct {
	RespCode int `json:"rc"`
}

func MockServer(t *testing.T, failCount int, retryAfter int) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/struct/path":
			fmt.Fprintf(w, `{"tree_size": 11, "timestamp": 99}`)
		case "/struct/params":
			var s TestStruct
			if r.Method == http.MethodGet {
				s.TreeSize, _ = strconv.Atoi(r.FormValue("tree_size"))
				s.Timestamp, _ = strconv.Atoi(r.FormValue("timestamp"))
				s.Data = r.FormValue("data")
			} else {
				decoder := json.NewDecoder(r.Body)
				err := decoder.Decode(&s)
				if err != nil {
					panic("Failed to decode: " + err.Error())
				}
				defer r.Body.Close()
			}
			fmt.Fprintf(w, `{"tree_size": %d, "timestamp": %d, "data": "%s"}`, s.TreeSize, s.Timestamp, s.Data)
		case "/error":
			var params TestParams
			if r.Method == http.MethodGet {
				params.RespCode, _ = strconv.Atoi(r.FormValue("rc"))
			} else {
				decoder := json.NewDecoder(r.Body)
				err := decoder.Decode(&params)
				if err != nil {
					panic("Failed to decode: " + err.Error())
				}
				defer r.Body.Close()
			}
			http.Error(w, "error page", params.RespCode)
		case "/malformed":
			fmt.Fprintf(w, `{"tree_size": 11, "timestamp": 99`) // no closing }
		case "/retry":
			if failCount > 0 {
				failCount--
				if retryAfter != 0 {
					if retryAfter > 0 {
						w.Header().Add("Retry-After", strconv.Itoa(retryAfter))
					}
					w.WriteHeader(http.StatusServiceUnavailable)
				} else {
					w.WriteHeader(http.StatusRequestTimeout)
				}
			} else {
				fmt.Fprintf(w, `{"tree_size": 11, "timestamp": 99}`)
			}
		default:
			t.Fatalf("Unhandled URL path: %s", r.URL.Path)
		}
	}))
}

func TestGetAndParse(t *testing.T) {
	tests := []struct {
		uri    string
		params map[string]string
		status int
		result TestStruct
		errstr string
	}{
		{uri: "[invalid-uri]", errstr: "too many colons"},
		{uri: "/short%", errstr: "invalid URL escape"},
		{uri: "/malformed", status: http.StatusOK, errstr: "unexpected EOF"},
		{uri: "/error", params: map[string]string{"rc": "404"}, status: http.StatusNotFound},
		{uri: "/error", params: map[string]string{"rc": "403"}, status: http.StatusForbidden},
		{uri: "/struct/path", status: http.StatusOK, result: TestStruct{11, 99, ""}},
		{
			uri:    "/struct/params",
			status: http.StatusOK,
			params: map[string]string{"tree_size": "42", "timestamp": "88", "data": "abcd"},
			result: TestStruct{42, 88, "abcd"},
		},
	}

	ts := MockServer(t, -1, 0)
	defer ts.Close()

	logClient, err := NewWithoutVerification(ts.URL, &http.Client{})
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()

	for _, test := range tests {
		var result TestStruct
		httpRsp, err := logClient.GetAndParse(ctx, test.uri, test.params, &result)
		if test.errstr != "" {
			if err == nil {
				t.Errorf("GetAndParse(%q)=%+v,nil; want error %q", test.uri, result, test.errstr)
			} else if !strings.Contains(err.Error(), test.errstr) {
				t.Errorf("GetAndParse(%q)=nil,%q; want error %q", test.uri, err.Error(), test.errstr)
			}
			continue
		}
		if httpRsp.StatusCode != test.status {
			t.Errorf("GetAndParse('%s') got status %d; want %d", test.uri, httpRsp.StatusCode, test.status)
		}
		if test.status == http.StatusOK {
			if err != nil {
				t.Errorf("GetAndParse(%q)=nil,%q; want %+v", test.uri, err.Error(), result)
			}
			if !reflect.DeepEqual(result, test.result) {
				t.Errorf("GetAndParse(%q)=%+v,nil; want %+v", test.uri, result, test.result)
			}
		}
	}
}

func TestPostAndParse(t *testing.T) {
	tests := []struct {
		uri     string
		request interface{}
		status  int
		result  TestStruct
		errstr  string
	}{
		{uri: "[invalid-uri]", errstr: "too many colons"},
		{uri: "/short%", errstr: "invalid URL escape"},
		{uri: "/struct/params", request: json.Number(`invalid`), errstr: "invalid number literal"},
		{uri: "/malformed", status: http.StatusOK, errstr: "unexpected end of JSON"},
		{uri: "/error", request: TestParams{RespCode: 404}, status: http.StatusNotFound},
		{uri: "/error", request: TestParams{RespCode: 403}, status: http.StatusForbidden},
		{uri: "/struct/path", status: http.StatusOK, result: TestStruct{11, 99, ""}},
		{
			uri:     "/struct/params",
			status:  http.StatusOK,
			request: TestStruct{42, 88, "abcd"},
			result:  TestStruct{42, 88, "abcd"},
		},
	}

	ts := MockServer(t, -1, 0)
	defer ts.Close()

	logClient, err := NewWithoutVerification(ts.URL, &http.Client{})
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()

	for _, test := range tests {
		var result TestStruct
		httpRsp, err := logClient.PostAndParse(ctx, test.uri, test.request, &result)
		if test.errstr != "" {
			if err == nil {
				t.Errorf("PostAndParse(%q)=%+v,nil; want %q", test.uri, result, test.errstr)
			} else if !strings.Contains(err.Error(), test.errstr) {
				t.Errorf("PostAndParse(%q)=nil,%q; want error %q", test.uri, err.Error(), test.errstr)
			}
			continue
		}
		if httpRsp.StatusCode != test.status {
			t.Errorf("PostAndParse(%q) got status %d; want %d", test.uri, httpRsp.StatusCode, test.status)
		}
		if test.status == http.StatusOK {
			if err != nil {
				t.Errorf("PostAndParse(%q)=nil,%q; want %+v", test.uri, err.Error(), test.result)
			}
			if !reflect.DeepEqual(result, test.result) {
				t.Errorf("PostAndParse(%q)=%+v,nil; want %+v", test.uri, result, test.result)
			}
		}
	}
}

func TestPostAndParseWithRetry(t *testing.T) {
	leeway := time.Millisecond * 100
	jiffy := time.Millisecond

	tests := []struct {
		uri          string
		request      interface{}
		deadlineSecs int // -1 indicates no deadline
		expected     time.Duration
		retryAfter   int // -1 indicates generate 503 with no Retry-After
		failCount    int
		errstr       string
	}{
		{"/retry", nil, -1, jiffy, 0, 0, ""},
		{"/error", TestParams{RespCode: 418}, 2, jiffy, 0, 0, "teapot"},
		{"/short%", nil, 2, 2 * time.Second, 0, 0, "deadline exceeded"},
		{"/retry", nil, -1, 7 * time.Second, -1, 3, ""},
		{"/retry", nil, 6, 5 * time.Second, 5, 1, ""},
		{"/retry", nil, 5, 5 * time.Second, 10, 1, "deadline exceeded"},
		{"/retry", nil, 10, 5 * time.Second, 1, 5, ""},
		{"/retry", nil, 1, 10 * jiffy, 0, 10, ""},
	}
	for _, test := range tests {
		ts := MockServer(t, test.failCount, test.retryAfter)
		defer ts.Close()

		logClient, err := NewWithoutVerification(ts.URL, &http.Client{})
		if err != nil {
			t.Fatal(err)
		}
		ctx := context.Background()
		if test.deadlineSecs >= 0 {
			ctx, _ = context.WithDeadline(context.Background(), time.Now().Add(time.Duration(test.deadlineSecs)*time.Second))
		}

		var result TestStruct
		started := time.Now()
		httpRsp, err := logClient.PostAndParseWithRetry(ctx, test.uri, test.request, &result)
		took := time.Since(started)

		if math.Abs(float64(took-test.expected)) > float64(leeway) {
			t.Errorf("PostAndParseWithRetry() took %s; want ~%s", took, test.expected)
		}
		if test.errstr != "" {
			if err == nil {
				t.Errorf("PostAndParseWithRetry()=%+v,nil; want error %q", result, test.errstr)
			} else if !strings.Contains(err.Error(), test.errstr) {
				t.Errorf("PostAndParseWithRetry()=nil,%q; want error %q", err.Error(), test.errstr)
			}
			continue
		}
		if err != nil {
			t.Errorf("PostAndParseWithRetry()=nil,%q; want no error", err.Error())
		} else if httpRsp.StatusCode != http.StatusOK {
			t.Errorf("PostAndParseWithRetry() got status %d; want OK(404)", httpRsp.StatusCode)
		}
	}
}

func TestContextRequired(t *testing.T) {
	ts := MockServer(t, -1, 0)
	defer ts.Close()

	logClient, err := NewWithoutVerification(ts.URL, &http.Client{})
	if err != nil {
		t.Fatal(err)
	}
	var result TestStruct
	_, err = logClient.GetAndParse(nil, "/struct/path", nil, &result)
	if err == nil {
		t.Errorf("GetAndParse() succeeded with empty Context")
	}
	_, err = logClient.PostAndParse(nil, "/struct/path", nil, &result)
	if err == nil {
		t.Errorf("PostAndParse() succeeded with empty Context")
	}
	_, err = logClient.PostAndParseWithRetry(nil, "/struct/path", nil, &result)
	if err == nil {
		t.Errorf("PostAndParseWithRetry() succeeded with empty Context")
	}
}
<commit_msg>go/jsonclient: fix test for Go 1.8<commit_after>package jsonclient

import (
	"encoding/json"
	"fmt"
	"math"
	"net/http"
	"net/http/httptest"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/google/certificate-transparency/go/testdata"
)

func TestNewJSONClient(t *testing.T) {
	tests := []struct {
		pubKey string
		errstr string
	}{
		{"", "no PEM block"},
		{"bogus", "no PEM block"},
		{testdata.RsaPublicKeyPEM, ""},
		{testdata.EcdsaPublicKeyPEM, ""},
		{testdata.DsaPublicKeyPEM, "Unsupported public key type"},
		{testdata.RsaPublicKeyPEM + "bogus", "extra data found"},
	}
	for _, test := range tests {
		client, err := New("http://127.0.0.1", nil, test.pubKey)
		if test.errstr != "" {
			if err == nil {
				t.Errorf("New()=%p,nil; want error %q", client, test.errstr)
			} else if !strings.Contains(err.Error(), test.errstr) {
				t.Errorf("New()=nil,%q; want error %q", err.Error(), test.errstr)
			}
			continue
		}
		if err != nil {
			t.Errorf("New()=nil,%q; want no error", err.Error())
		} else if client == nil {
			t.Errorf("New()=nil,nil; want client")
		}
	}
}

type TestStruct struct {
	TreeSize  int    `json:"tree_size"`
	Timestamp int    `json:"timestamp"`
	Data      string `json:"data"`
}

type TestParams struct {
	RespCode int `json:"rc"`
}

func MockServer(t *testing.T, failCount int, retryAfter int) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/struct/path":
			fmt.Fprintf(w, `{"tree_size": 11, "timestamp": 99}`)
		case "/struct/params":
			var s TestStruct
			if r.Method == http.MethodGet {
				s.TreeSize, _ = strconv.Atoi(r.FormValue("tree_size"))
				s.Timestamp, _ = strconv.Atoi(r.FormValue("timestamp"))
				s.Data = r.FormValue("data")
			} else {
				decoder := json.NewDecoder(r.Body)
				err := decoder.Decode(&s)
				if err != nil {
					panic("Failed to decode: " + err.Error())
				}
				defer r.Body.Close()
			}
			fmt.Fprintf(w, `{"tree_size": %d, "timestamp": %d, "data": "%s"}`, s.TreeSize, s.Timestamp, s.Data)
		case "/error":
			var params TestParams
			if r.Method == http.MethodGet {
				params.RespCode, _ = strconv.Atoi(r.FormValue("rc"))
			} else {
				decoder := json.NewDecoder(r.Body)
				err := decoder.Decode(&params)
				if err != nil {
					panic("Failed to decode: " + err.Error())
				}
				defer r.Body.Close()
			}
			http.Error(w, "error page", params.RespCode)
		case "/malformed":
			fmt.Fprintf(w, `{"tree_size": 11, "timestamp": 99`) // no closing }
		case "/retry":
			if failCount > 0 {
				failCount--
				if retryAfter != 0 {
					if retryAfter > 0 {
						w.Header().Add("Retry-After", strconv.Itoa(retryAfter))
					}
					w.WriteHeader(http.StatusServiceUnavailable)
				} else {
					w.WriteHeader(http.StatusRequestTimeout)
				}
			} else {
				fmt.Fprintf(w, `{"tree_size": 11, "timestamp": 99}`)
			}
		default:
			t.Fatalf("Unhandled URL path: %s", r.URL.Path)
		}
	}))
}

func TestGetAndParse(t *testing.T) {
	rc := regexp.MustCompile
	tests := []struct {
		uri    string
		params map[string]string
		status int
		result TestStruct
		errstr *regexp.Regexp
	}{
		{uri: "[invalid-uri]", errstr: rc("too many colons|unexpected .* in address")},
		{uri: "/short%", errstr: rc("invalid URL escape")},
		{uri: "/malformed", status: http.StatusOK, errstr: rc("unexpected EOF")},
		{uri: "/error", params: map[string]string{"rc": "404"}, status: http.StatusNotFound},
		{uri: "/error", params: map[string]string{"rc": "403"}, status: http.StatusForbidden},
		{uri: "/struct/path", status: http.StatusOK, result: TestStruct{11, 99, ""}},
		{
			uri:    "/struct/params",
			status: http.StatusOK,
			params: map[string]string{"tree_size": "42", "timestamp": "88", "data": "abcd"},
			result: TestStruct{42, 88, "abcd"},
		},
	}

	ts := MockServer(t, -1, 0)
	defer ts.Close()

	logClient, err := NewWithoutVerification(ts.URL, &http.Client{})
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()

	for _, test := range tests {
		var result TestStruct
		httpRsp, err := logClient.GetAndParse(ctx, test.uri, test.params, &result)
		if test.errstr != nil {
			if err == nil {
				t.Errorf("GetAndParse(%q)=%+v,nil; want error matching %q", test.uri, result, test.errstr)
			} else if !test.errstr.MatchString(err.Error()) {
				t.Errorf("GetAndParse(%q)=nil,%q; want error matching %q", test.uri, err.Error(), test.errstr)
			}
			continue
		}
		if httpRsp.StatusCode != test.status {
			t.Errorf("GetAndParse('%s') got status %d; want %d", test.uri, httpRsp.StatusCode, test.status)
		}
		if test.status == http.StatusOK {
			if err != nil {
				t.Errorf("GetAndParse(%q)=nil,%q; want %+v", test.uri, err.Error(), result)
			}
			if !reflect.DeepEqual(result, test.result) {
				t.Errorf("GetAndParse(%q)=%+v,nil; want %+v", test.uri, result, test.result)
			}
		}
	}
}

func TestPostAndParse(t *testing.T) {
	rc := regexp.MustCompile
	tests := []struct {
		uri     string
		request interface{}
		status  int
		result  TestStruct
		errstr  *regexp.Regexp
	}{
		{uri: "[invalid-uri]", errstr: rc("too many colons|unexpected .* in address")},
		{uri: "/short%", errstr: rc("invalid URL escape")},
		{uri: "/struct/params", request: json.Number(`invalid`), errstr: rc("invalid number literal")},
		{uri: "/malformed", status: http.StatusOK, errstr: rc("unexpected end of JSON")},
		{uri: "/error", request: TestParams{RespCode: 404}, status: http.StatusNotFound},
		{uri: "/error", request: TestParams{RespCode: 403}, status: http.StatusForbidden},
		{uri: "/struct/path", status: http.StatusOK, result: TestStruct{11, 99, ""}},
		{
			uri:     "/struct/params",
			status:  http.StatusOK,
			request: TestStruct{42, 88, "abcd"},
			result:  TestStruct{42, 88, "abcd"},
		},
	}

	ts := MockServer(t, -1, 0)
	defer ts.Close()

	logClient, err := NewWithoutVerification(ts.URL, &http.Client{})
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()

	for _, test := range tests {
		var result TestStruct
		httpRsp, err := logClient.PostAndParse(ctx, test.uri, test.request, &result)
		if test.errstr != nil {
			if err == nil {
				t.Errorf("PostAndParse(%q)=%+v,nil; want error matching %q", test.uri, result, test.errstr)
			} else if !test.errstr.MatchString(err.Error()) {
				t.Errorf("PostAndParse(%q)=nil,%q; want error matching %q", test.uri, err.Error(), test.errstr)
			}
			continue
		}
		if httpRsp.StatusCode != test.status {
			t.Errorf("PostAndParse(%q) got status %d; want %d", test.uri, httpRsp.StatusCode, test.status)
		}
		if test.status == http.StatusOK {
			if err != nil {
				t.Errorf("PostAndParse(%q)=nil,%q; want %+v", test.uri, err.Error(), test.result)
			}
			if !reflect.DeepEqual(result, test.result) {
				t.Errorf("PostAndParse(%q)=%+v,nil; want %+v", test.uri, result, test.result)
			}
		}
	}
}

func TestPostAndParseWithRetry(t *testing.T) {
	leeway := time.Millisecond * 100
	jiffy := time.Millisecond

	tests := []struct {
		uri          string
		request      interface{}
		deadlineSecs int // -1 indicates no deadline
		expected     time.Duration
		retryAfter   int // -1 indicates generate 503 with no Retry-After
		failCount    int
		errstr       string
	}{
		{"/retry", nil, -1, jiffy, 0, 0, ""},
		{"/error", TestParams{RespCode: 418}, 2, jiffy, 0, 0, "teapot"},
		{"/short%", nil, 2, 2 * time.Second, 0, 0, "deadline exceeded"},
		{"/retry", nil, -1, 7 * time.Second, -1, 3, ""},
		{"/retry", nil, 6, 5 * time.Second, 5, 1, ""},
		{"/retry", nil, 5, 5 * time.Second, 10, 1, "deadline exceeded"},
		{"/retry", nil, 10, 5 * time.Second, 1, 5, ""},
		{"/retry", nil, 1, 10 * jiffy, 0, 10, ""},
	}
	for _, test := range tests {
		ts := MockServer(t, test.failCount, test.retryAfter)
		defer ts.Close()

		logClient, err := NewWithoutVerification(ts.URL, &http.Client{})
		if err != nil {
			t.Fatal(err)
		}
		ctx := context.Background()
		if test.deadlineSecs >= 0 {
			ctx, _ = context.WithDeadline(context.Background(), time.Now().Add(time.Duration(test.deadlineSecs)*time.Second))
		}

		var result TestStruct
		started := time.Now()
		httpRsp, err := logClient.PostAndParseWithRetry(ctx, test.uri, test.request, &result)
		took := time.Since(started)

		if math.Abs(float64(took-test.expected)) > float64(leeway) {
			t.Errorf("PostAndParseWithRetry() took %s; want ~%s", took, test.expected)
		}
		if test.errstr != "" {
			if err == nil {
				t.Errorf("PostAndParseWithRetry()=%+v,nil; want error %q", result, test.errstr)
			} else if !strings.Contains(err.Error(), test.errstr) {
				t.Errorf("PostAndParseWithRetry()=nil,%q; want error %q", err.Error(), test.errstr)
			}
			continue
		}
		if err != nil {
			t.Errorf("PostAndParseWithRetry()=nil,%q; want no error", err.Error())
		} else if httpRsp.StatusCode != http.StatusOK {
			t.Errorf("PostAndParseWithRetry() got status %d; want OK(404)", httpRsp.StatusCode)
		}
	}
}

func TestContextRequired(t *testing.T) {
	ts := MockServer(t, -1, 0)
	defer ts.Close()

	logClient, err := NewWithoutVerification(ts.URL, &http.Client{})
	if err != nil {
		t.Fatal(err)
	}
	var result TestStruct
	_, err = logClient.GetAndParse(nil, "/struct/path", nil, &result)
	if err == nil {
		t.Errorf("GetAndParse() succeeded with empty Context")
	}
	_, err = logClient.PostAndParse(nil, "/struct/path", nil, &result)
	if err == nil {
		t.Errorf("PostAndParse() succeeded with empty Context")
	}
	_, err = logClient.PostAndParseWithRetry(nil, "/struct/path", nil, &result)
	if err == nil {
		t.Errorf("PostAndParseWithRetry() succeeded with empty Context")
	}
}
<|endoftext|>


<commit_before>package sqlite

/*
#cgo pkg-config: sqlite3

#include <sqlite3.h>
#include <stdlib.h>
#include <string.h>

#ifndef SQLITE_OPEN_READWRITE
# define SQLITE_OPEN_READWRITE 0
#endif

#ifndef SQLITE_OPEN_FULLMUTEX
# define SQLITE_OPEN_FULLMUTEX 0
#endif

static int
_sqlite3_open_v2(const char *filename, sqlite3 **ppDb, int flags, const char *zVfs) {
#ifdef SQLITE_OPEN_URI
  return sqlite3_open_v2(filename, ppDb, flags | SQLITE_OPEN_URI, zVfs);
#else
  return sqlite3_open_v2(filename, ppDb, flags, zVfs);
#endif
}

static int
_sqlite3_bind_text(sqlite3_stmt *stmt, int n, char *p, int np) {
  return sqlite3_bind_text(stmt, n, p, np, SQLITE_TRANSIENT);
}

static int
_sqlite3_bind_blob(sqlite3_stmt *stmt, int n, void *p, int np) {
  return sqlite3_bind_blob(stmt, n, p, np, SQLITE_TRANSIENT);
}

#include <stdio.h>
#include <stdint.h>

static long
_sqlite3_last_insert_rowid(sqlite3* db) {
  return (long) sqlite3_last_insert_rowid(db);
}

static long
_sqlite3_changes(sqlite3* db) {
  return (long) sqlite3_changes(db);
}

*/
import "C"
import (
	"database/sql"
	"database/sql/driver"
	"errors"
	"io"
	"strings"
	"time"
	"unsafe"
)

// Timestamp formats understood by both this module and SQLite.
// The first format in the slice will be used when saving time values
// into the database. When parsing a string from a timestamp or
// datetime column, the formats are tried in order.
var SQLiteTimestampFormats = []string{
	"2006-01-02 15:04:05.999999999",
	"2006-01-02T15:04:05.999999999",
	"2006-01-02 15:04:05",
	"2006-01-02T15:04:05",
	"2006-01-02 15:04",
	"2006-01-02T15:04",
	"2006-01-02",
}

func init() {
	sql.Register("sqlite3", &SQLiteDriver{})
}

// Driver struct.
type SQLiteDriver struct {
}

// Conn struct.
type SQLiteConn struct {
	db *C.sqlite3
}

// Tx struct.
type SQLiteTx struct {
	c *SQLiteConn
}

// Stmt struct.
type SQLiteStmt struct {
	c      *SQLiteConn
	s      *C.sqlite3_stmt
	t      string
	closed bool
}

// Result struct.
type SQLiteResult struct {
	id      int64
	changes int64
}

// Rows struct.
type SQLiteRows struct {
	s        *SQLiteStmt
	nc       int
	cols     []string
	decltype []string
}

// Commit transaction.
func (tx *SQLiteTx) Commit() error {
	if err := tx.c.exec("COMMIT"); err != nil {
		return err
	}
	return nil
}

// Rollback transaction.
func (tx *SQLiteTx) Rollback() error {
	if err := tx.c.exec("ROLLBACK"); err != nil {
		return err
	}
	return nil
}

func (c *SQLiteConn) exec(cmd string) error {
	pcmd := C.CString(cmd)
	defer C.free(unsafe.Pointer(pcmd))
	rv := C.sqlite3_exec(c.db, pcmd, nil, nil, nil)
	if rv != C.SQLITE_OK {
		return errors.New(C.GoString(C.sqlite3_errmsg(c.db)))
	}
	return nil
}

// Begin transaction.
func (c *SQLiteConn) Begin() (driver.Tx, error) {
	if err := c.exec("BEGIN"); err != nil {
		return nil, err
	}
	return &SQLiteTx{c}, nil
}

// Open database and return a new connection.
// You can specify DSN string with URI filename.
//   test.db
//   file:test.db?cache=shared&mode=memory
//   :memory:
//   file::memory:
func (d *SQLiteDriver) Open(dsn string) (driver.Conn, error) {
	if C.sqlite3_threadsafe() == 0 {
		return nil, errors.New("sqlite library was not compiled for thread-safe operation")
	}

	var db *C.sqlite3
	name := C.CString(dsn)
	defer C.free(unsafe.Pointer(name))
	rv := C._sqlite3_open_v2(name, &db,
		C.SQLITE_OPEN_FULLMUTEX|
			C.SQLITE_OPEN_READWRITE|
			C.SQLITE_OPEN_CREATE,
		nil)
	if rv != 0 {
		return nil, errors.New(C.GoString(C.sqlite3_errmsg(db)))
	}
	if db == nil {
		return nil, errors.New("sqlite succeeded without returning a database")
	}

	rv = C.sqlite3_busy_timeout(db, 5000)
	if rv != C.SQLITE_OK {
		return nil, errors.New(C.GoString(C.sqlite3_errmsg(db)))
	}

	return &SQLiteConn{db}, nil
}

// Close the connection.
func (c *SQLiteConn) Close() error {
	s := C.sqlite3_next_stmt(c.db, nil)
	for s != nil {
		C.sqlite3_finalize(s)
		s = C.sqlite3_next_stmt(c.db, nil)
	}
	rv := C.sqlite3_close(c.db)
	if rv != C.SQLITE_OK {
		return errors.New("error while closing sqlite database connection")
	}
	c.db = nil
	return nil
}

// Prepare query string. Return a new statement.
func (c *SQLiteConn) Prepare(query string) (driver.Stmt, error) {
	pquery := C.CString(query)
	defer C.free(unsafe.Pointer(pquery))
	var s *C.sqlite3_stmt
	var perror *C.char
	rv := C.sqlite3_prepare_v2(c.db, pquery, -1, &s, &perror)
	if rv != C.SQLITE_OK {
		return nil, errors.New(C.GoString(C.sqlite3_errmsg(c.db)))
	}
	var t string
	if perror != nil && C.strlen(perror) > 0 {
		t = C.GoString(perror)
	}
	return &SQLiteStmt{c: c, s: s, t: t}, nil
}

// Close the statement.
func (s *SQLiteStmt) Close() error {
	if s.closed {
		return nil
	}
	s.closed = true
	if s.c == nil || s.c.db == nil {
		return errors.New("sqlite statement with already closed database connection")
	}
	rv := C.sqlite3_finalize(s.s)
	if rv != C.SQLITE_OK {
		return errors.New(C.GoString(C.sqlite3_errmsg(s.c.db)))
	}
	return nil
}

// Return a number of parameters.
func (s *SQLiteStmt) NumInput() int {
	return int(C.sqlite3_bind_parameter_count(s.s))
}

func (s *SQLiteStmt) bind(args []driver.Value) error {
	rv := C.sqlite3_reset(s.s)
	if rv != C.SQLITE_ROW && rv != C.SQLITE_OK && rv != C.SQLITE_DONE {
		return errors.New(C.GoString(C.sqlite3_errmsg(s.c.db)))
	}

	for i, v := range args {
		n := C.int(i + 1)
		switch v := v.(type) {
		case nil:
			rv = C.sqlite3_bind_null(s.s, n)
		case string:
			if len(v) == 0 {
				b := []byte{0}
				rv = C._sqlite3_bind_text(s.s, n, (*C.char)(unsafe.Pointer(&b[0])), C.int(0))
			} else {
				b := []byte(v)
				rv = C._sqlite3_bind_text(s.s, n, (*C.char)(unsafe.Pointer(&b[0])), C.int(len(b)))
			}
		case int:
			rv = C.sqlite3_bind_int64(s.s, n, C.sqlite3_int64(v))
		case int32:
			rv = C.sqlite3_bind_int(s.s, n, C.int(v))
		case int64:
			rv = C.sqlite3_bind_int64(s.s, n, C.sqlite3_int64(v))
		case byte:
			rv = C.sqlite3_bind_int(s.s, n, C.int(v))
		case bool:
			if bool(v) {
				rv = C.sqlite3_bind_int(s.s, n, 1)
			} else {
				rv = C.sqlite3_bind_int(s.s, n, 0)
			}
		case float32:
			rv = C.sqlite3_bind_double(s.s, n, C.double(v))
		case float64:
			rv = C.sqlite3_bind_double(s.s, n, C.double(v))
		case []byte:
			var p *byte
			if len(v) > 0 {
				p = &v[0]
			}
			rv = C._sqlite3_bind_blob(s.s, n, unsafe.Pointer(p), C.int(len(v)))
		case time.Time:
			b := []byte(v.UTC().Format(SQLiteTimestampFormats[0]))
			rv = C._sqlite3_bind_text(s.s, n, (*C.char)(unsafe.Pointer(&b[0])), C.int(len(b)))
		}
		if rv != C.SQLITE_OK {
			return errors.New(C.GoString(C.sqlite3_errmsg(s.c.db)))
		}
	}
	return nil
}

// Query the statment with arguments. Return records.
func (s *SQLiteStmt) Query(args []driver.Value) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	return &SQLiteRows{s, int(C.sqlite3_column_count(s.s)), nil, nil}, nil
}

// Return last inserted ID.
func (r *SQLiteResult) LastInsertId() (int64, error) {
	return r.id, nil
}

// Return how many rows affected.
func (r *SQLiteResult) RowsAffected() (int64, error) {
	return r.changes, nil
}

// Execute the statement with arguments. Return result object.
func (s *SQLiteStmt) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rv := C.sqlite3_step(s.s)
	if rv != C.SQLITE_ROW && rv != C.SQLITE_OK && rv != C.SQLITE_DONE {
		return nil, errors.New(C.GoString(C.sqlite3_errmsg(s.c.db)))
	}

	res := &SQLiteResult{
		int64(C._sqlite3_last_insert_rowid(s.c.db)),
		int64(C._sqlite3_changes(s.c.db)),
	}
	return res, nil
}

// Close the rows.
func (rc *SQLiteRows) Close() error {
	rv := C.sqlite3_reset(rc.s.s)
	if rv != C.SQLITE_OK {
		return errors.New(C.GoString(C.sqlite3_errmsg(rc.s.c.db)))
	}
	return nil
}

// Return column names.
func (rc *SQLiteRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			rc.cols[i] = C.GoString(C.sqlite3_column_name(rc.s.s, C.int(i)))
		}
	}
	return rc.cols
}

// Move cursor to next.
func (rc *SQLiteRows) Next(dest []driver.Value) error {
	rv := C.sqlite3_step(rc.s.s)
	if rv == C.SQLITE_DONE {
		return io.EOF
	}
	if rv != C.SQLITE_ROW {
		return errors.New(C.GoString(C.sqlite3_errmsg(rc.s.c.db)))
	}

	if rc.decltype == nil {
		rc.decltype = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			rc.decltype[i] = strings.ToLower(C.GoString(C.sqlite3_column_decltype(rc.s.s, C.int(i))))
		}
	}

	for i := range dest {
		switch C.sqlite3_column_type(rc.s.s, C.int(i)) {
		case C.SQLITE_INTEGER:
			val := int64(C.sqlite3_column_int64(rc.s.s, C.int(i)))
			switch rc.decltype[i] {
			case "timestamp", "datetime":
				dest[i] = time.Unix(val, 0)
			case "boolean":
				dest[i] = val > 0
			default:
				dest[i] = val
			}
		case C.SQLITE_FLOAT:
			dest[i] = float64(C.sqlite3_column_double(rc.s.s, C.int(i)))
		case C.SQLITE_BLOB:
			p := C.sqlite3_column_blob(rc.s.s, C.int(i))
			n := int(C.sqlite3_column_bytes(rc.s.s, C.int(i)))
			switch dest[i].(type) {
			case sql.RawBytes:
				dest[i] = (*[1 << 30]byte)(unsafe.Pointer(p))[0:n]
			default:
				slice := make([]byte, n)
				copy(slice[:], (*[1 << 30]byte)(unsafe.Pointer(p))[0:n])
				dest[i] = slice
			}
		case C.SQLITE_NULL:
			dest[i] = nil
		case C.SQLITE_TEXT:
			var err error
			s := C.GoString((*C.char)(unsafe.Pointer(C.sqlite3_column_text(rc.s.s, C.int(i)))))

			switch rc.decltype[i] {
			case "timestamp", "datetime":
				for _, format := range SQLiteTimestampFormats {
					if dest[i], err = time.Parse(format, s); err == nil {
						break
					}
				}
				if err != nil {
					// The column is a time value, so return the zero time on parse failure.
					dest[i] = time.Time{}
				}
			default:
				dest[i] = s
			}

		}
	}
	return nil
}
<commit_msg>Update go-sqlite3 to fix hanging issue when inserting into snippet_search<commit_after>package sqlite

/*
#cgo pkg-config: sqlite3

#include <sqlite3.h>
#include <stdlib.h>
#include <string.h>

#ifndef SQLITE_OPEN_READWRITE
# define SQLITE_OPEN_READWRITE 0
#endif

#ifndef SQLITE_OPEN_FULLMUTEX
# define SQLITE_OPEN_FULLMUTEX 0
#endif

static int
_sqlite3_open_v2(const char *filename, sqlite3 **ppDb, int flags, const char *zVfs) {
#ifdef SQLITE_OPEN_URI
  return sqlite3_open_v2(filename, ppDb, flags | SQLITE_OPEN_URI, zVfs);
#else
  return sqlite3_open_v2(filename, ppDb, flags, zVfs);
#endif
}

static int
_sqlite3_bind_text(sqlite3_stmt *stmt, int n, char *p, int np) {
  return sqlite3_bind_text(stmt, n, p, np, SQLITE_TRANSIENT);
}

static int
_sqlite3_bind_blob(sqlite3_stmt *stmt, int n, void *p, int np) {
  return sqlite3_bind_blob(stmt, n, p, np, SQLITE_TRANSIENT);
}

#include <stdio.h>
#include <stdint.h>

static long
_sqlite3_last_insert_rowid(sqlite3* db) {
  return (long) sqlite3_last_insert_rowid(db);
}

static long
_sqlite3_changes(sqlite3* db) {
  return (long) sqlite3_changes(db);
}

*/
import "C"
import (
	"database/sql"
	"database/sql/driver"
	"errors"
	"io"
	"strings"
	"time"
	"unsafe"
)

// Timestamp formats understood by both this module and SQLite.
// The first format in the slice will be used when saving time values
// into the database. When parsing a string from a timestamp or
// datetime column, the formats are tried in order.
var SQLiteTimestampFormats = []string{
	"2006-01-02 15:04:05.999999999",
	"2006-01-02T15:04:05.999999999",
	"2006-01-02 15:04:05",
	"2006-01-02T15:04:05",
	"2006-01-02 15:04",
	"2006-01-02T15:04",
	"2006-01-02",
}

func init() {
	sql.Register("sqlite3", &SQLiteDriver{})
}

// Driver struct.
type SQLiteDriver struct {
}

// Conn struct.
type SQLiteConn struct {
	db *C.sqlite3
}

// Tx struct.
type SQLiteTx struct {
	c *SQLiteConn
}

// Stmt struct.
type SQLiteStmt struct {
	c      *SQLiteConn
	s      *C.sqlite3_stmt
	t      string
	closed bool
}

// Result struct.
type SQLiteResult struct {
	id      int64
	changes int64
}

// Rows struct.
type SQLiteRows struct {
	s        *SQLiteStmt
	nc       int
	cols     []string
	decltype []string
}

// Commit transaction.
func (tx *SQLiteTx) Commit() error {
	if err := tx.c.exec("COMMIT"); err != nil {
		return err
	}
	return nil
}

// Rollback transaction.
func (tx *SQLiteTx) Rollback() error {
	if err := tx.c.exec("ROLLBACK"); err != nil {
		return err
	}
	return nil
}

func (c *SQLiteConn) exec(cmd string) error {
	pcmd := C.CString(cmd)
	defer C.free(unsafe.Pointer(pcmd))
	rv := C.sqlite3_exec(c.db, pcmd, nil, nil, nil)
	if rv != C.SQLITE_OK {
		return errors.New(C.GoString(C.sqlite3_errmsg(c.db)))
	}
	return nil
}

// Begin transaction.
func (c *SQLiteConn) Begin() (driver.Tx, error) {
	if err := c.exec("BEGIN"); err != nil {
		return nil, err
	}
	return &SQLiteTx{c}, nil
}

// Open database and return a new connection.
// You can specify DSN string with URI filename.
//   test.db
//   file:test.db?cache=shared&mode=memory
//   :memory:
//   file::memory:
func (d *SQLiteDriver) Open(dsn string) (driver.Conn, error) {
	if C.sqlite3_threadsafe() == 0 {
		return nil, errors.New("sqlite library was not compiled for thread-safe operation")
	}

	var db *C.sqlite3
	name := C.CString(dsn)
	defer C.free(unsafe.Pointer(name))
	rv := C._sqlite3_open_v2(name, &db,
		C.SQLITE_OPEN_FULLMUTEX|
			C.SQLITE_OPEN_READWRITE|
			C.SQLITE_OPEN_CREATE,
		nil)
	if rv != 0 {
		return nil, errors.New(C.GoString(C.sqlite3_errmsg(db)))
	}
	if db == nil {
		return nil, errors.New("sqlite succeeded without returning a database")
	}

	rv = C.sqlite3_busy_timeout(db, 5000)
	if rv != C.SQLITE_OK {
		return nil, errors.New(C.GoString(C.sqlite3_errmsg(db)))
	}

	return &SQLiteConn{db}, nil
}

// Close the connection.
func (c *SQLiteConn) Close() error {
	// s := C.sqlite3_next_stmt(c.db, nil)
	// for s != nil {
	// 	C.sqlite3_finalize(s)
	// 	s = C.sqlite3_next_stmt(c.db, nil)
	// }
	rv := C.sqlite3_close(c.db)
	if rv != C.SQLITE_OK {
		return errors.New("error while closing sqlite database connection")
	}
	c.db = nil
	return nil
}

// Prepare query string. Return a new statement.
func (c *SQLiteConn) Prepare(query string) (driver.Stmt, error) {
	pquery := C.CString(query)
	defer C.free(unsafe.Pointer(pquery))
	var s *C.sqlite3_stmt
	var perror *C.char
	rv := C.sqlite3_prepare_v2(c.db, pquery, -1, &s, &perror)
	if rv != C.SQLITE_OK {
		return nil, errors.New(C.GoString(C.sqlite3_errmsg(c.db)))
	}
	var t string
	if perror != nil && C.strlen(perror) > 0 {
		t = C.GoString(perror)
	}
	return &SQLiteStmt{c: c, s: s, t: t}, nil
}

// Close the statement.
func (s *SQLiteStmt) Close() error {
	if s.closed {
		return nil
	}
	s.closed = true
	if s.c == nil || s.c.db == nil {
		return errors.New("sqlite statement with already closed database connection")
	}
	rv := C.sqlite3_finalize(s.s)
	if rv != C.SQLITE_OK {
		return errors.New(C.GoString(C.sqlite3_errmsg(s.c.db)))
	}
	return nil
}

// Return a number of parameters.
func (s *SQLiteStmt) NumInput() int {
	return int(C.sqlite3_bind_parameter_count(s.s))
}

func (s *SQLiteStmt) bind(args []driver.Value) error {
	rv := C.sqlite3_reset(s.s)
	if rv != C.SQLITE_ROW && rv != C.SQLITE_OK && rv != C.SQLITE_DONE {
		return errors.New(C.GoString(C.sqlite3_errmsg(s.c.db)))
	}

	for i, v := range args {
		n := C.int(i + 1)
		switch v := v.(type) {
		case nil:
			rv = C.sqlite3_bind_null(s.s, n)
		case string:
			if len(v) == 0 {
				b := []byte{0}
				rv = C._sqlite3_bind_text(s.s, n, (*C.char)(unsafe.Pointer(&b[0])), C.int(0))
			} else {
				b := []byte(v)
				rv = C._sqlite3_bind_text(s.s, n, (*C.char)(unsafe.Pointer(&b[0])), C.int(len(b)))
			}
		case int:
			rv = C.sqlite3_bind_int64(s.s, n, C.sqlite3_int64(v))
		case int32:
			rv = C.sqlite3_bind_int(s.s, n, C.int(v))
		case int64:
			rv = C.sqlite3_bind_int64(s.s, n, C.sqlite3_int64(v))
		case byte:
			rv = C.sqlite3_bind_int(s.s, n, C.int(v))
		case bool:
			if bool(v) {
				rv = C.sqlite3_bind_int(s.s, n, 1)
			} else {
				rv = C.sqlite3_bind_int(s.s, n, 0)
			}
		case float32:
			rv = C.sqlite3_bind_double(s.s, n, C.double(v))
		case float64:
			rv = C.sqlite3_bind_double(s.s, n, C.double(v))
		case []byte:
			var p *byte
			if len(v) > 0 {
				p = &v[0]
			}
			rv = C._sqlite3_bind_blob(s.s, n, unsafe.Pointer(p), C.int(len(v)))
		case time.Time:
			b := []byte(v.UTC().Format(SQLiteTimestampFormats[0]))
			rv = C._sqlite3_bind_text(s.s, n, (*C.char)(unsafe.Pointer(&b[0])), C.int(len(b)))
		}
		if rv != C.SQLITE_OK {
			return errors.New(C.GoString(C.sqlite3_errmsg(s.c.db)))
		}
	}
	return nil
}

// Query the statment with arguments. Return records.
func (s *SQLiteStmt) Query(args []driver.Value) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	return &SQLiteRows{s, int(C.sqlite3_column_count(s.s)), nil, nil}, nil
}

// Return last inserted ID.
func (r *SQLiteResult) LastInsertId() (int64, error) {
	return r.id, nil
}

// Return how many rows affected.
func (r *SQLiteResult) RowsAffected() (int64, error) {
	return r.changes, nil
}

// Execute the statement with arguments. Return result object.
func (s *SQLiteStmt) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rv := C.sqlite3_step(s.s)
	if rv != C.SQLITE_ROW && rv != C.SQLITE_OK && rv != C.SQLITE_DONE {
		return nil, errors.New(C.GoString(C.sqlite3_errmsg(s.c.db)))
	}

	res := &SQLiteResult{
		int64(C._sqlite3_last_insert_rowid(s.c.db)),
		int64(C._sqlite3_changes(s.c.db)),
	}
	return res, nil
}

// Close the rows.
func (rc *SQLiteRows) Close() error {
	rv := C.sqlite3_reset(rc.s.s)
	if rv != C.SQLITE_OK {
		return errors.New(C.GoString(C.sqlite3_errmsg(rc.s.c.db)))
	}
	return nil
}

// Return column names.
func (rc *SQLiteRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			rc.cols[i] = C.GoString(C.sqlite3_column_name(rc.s.s, C.int(i)))
		}
	}
	return rc.cols
}

// Move cursor to next.
func (rc *SQLiteRows) Next(dest []driver.Value) error {
	rv := C.sqlite3_step(rc.s.s)
	if rv == C.SQLITE_DONE {
		return io.EOF
	}
	if rv != C.SQLITE_ROW {
		return errors.New(C.GoString(C.sqlite3_errmsg(rc.s.c.db)))
	}

	if rc.decltype == nil {
		rc.decltype = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			rc.decltype[i] = strings.ToLower(C.GoString(C.sqlite3_column_decltype(rc.s.s, C.int(i))))
		}
	}

	for i := range dest {
		switch C.sqlite3_column_type(rc.s.s, C.int(i)) {
		case C.SQLITE_INTEGER:
			val := int64(C.sqlite3_column_int64(rc.s.s, C.int(i)))
			switch rc.decltype[i] {
			case "timestamp", "datetime":
				dest[i] = time.Unix(val, 0)
			case "boolean":
				dest[i] = val > 0
			default:
				dest[i] = val
			}
		case C.SQLITE_FLOAT:
			dest[i] = float64(C.sqlite3_column_double(rc.s.s, C.int(i)))
		case C.SQLITE_BLOB:
			p := C.sqlite3_column_blob(rc.s.s, C.int(i))
			n := int(C.sqlite3_column_bytes(rc.s.s, C.int(i)))
			switch dest[i].(type) {
			case sql.RawBytes:
				dest[i] = (*[1 << 30]byte)(unsafe.Pointer(p))[0:n]
			default:
				slice := make([]byte, n)
				copy(slice[:], (*[1 << 30]byte)(unsafe.Pointer(p))[0:n])
				dest[i] = slice
			}
		case C.SQLITE_NULL:
			dest[i] = nil
		case C.SQLITE_TEXT:
			var err error
			s := C.GoString((*C.char)(unsafe.Pointer(C.sqlite3_column_text(rc.s.s, C.int(i)))))

			switch rc.decltype[i] {
			case "timestamp", "datetime":
				for _, format := range SQLiteTimestampFormats {
					if dest[i], err = time.Parse(format, s); err == nil {
						break
					}
				}
				if err != nil {
					// The column is a time value, so return the zero time on parse failure.
					dest[i] = time.Time{}
				}
			default:
				dest[i] = s
			}

		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package servenv

import (
	"crypto/tls"
	"flag"
	"fmt"
	"math"
	"net"
	"time"

	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"

	"vitess.io/vitess/go/trace"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/reflection"

	"context"

	"vitess.io/vitess/go/vt/grpccommon"
	"vitess.io/vitess/go/vt/grpcoptionaltls"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/vttls"
)

// This file handles gRPC server, on its own port.
// Clients register servers, based on service map:
//
// servenv.RegisterGRPCFlags()
// servenv.OnRun(func() {
//   if servenv.GRPCCheckServiceMap("XXX") {
//     pb.RegisterXXX(servenv.GRPCServer, XXX)
//   }
// }
//
// Note servenv.GRPCServer can only be used in servenv.OnRun,
// and not before, as it is initialized right before calling OnRun.
var (
	// GRPCPort is the port to listen on for gRPC. If not set or zero, don't listen.
	GRPCPort = flag.Int("grpc_port", 0, "Port to listen on for gRPC calls")

	// GRPCCert is the cert to use if TLS is enabled
	GRPCCert = flag.String("grpc_cert", "", "server certificate to use for gRPC connections, requires grpc_key, enables TLS")

	// GRPCKey is the key to use if TLS is enabled
	GRPCKey = flag.String("grpc_key", "", "server private key to use for gRPC connections, requires grpc_cert, enables TLS")

	// GRPCCA is the CA to use if TLS is enabled
	GRPCCA = flag.String("grpc_ca", "", "server CA to use for gRPC connections, requires TLS, and enforces client certificate check")

	// GRPCCRL is the CRL (Certificate Revocation List) to use if TLS is enabled
	GRPCCRL = flag.String("grpc_crl", "", "path to a certificate revocation list in PEM format, client certificates will be further verified against this file during TLS handshake")

	GRPCEnableOptionalTLS = flag.Bool("grpc_enable_optional_tls", false, "enable optional TLS mode when a server accepts both TLS and plain-text connections on the same port")

	// GRPCServerCA if specified will combine server cert and server CA
	GRPCServerCA = flag.String("grpc_server_ca", "", "path to server CA in PEM format, which will be combine with server cert, return full certificate chain to clients")

	// GRPCAuth which auth plugin to use (at the moment now only static is supported)
	GRPCAuth = flag.String("grpc_auth_mode", "", "Which auth plugin implementation to use (eg: static)")

	// GRPCServer is the global server to serve gRPC.
	GRPCServer *grpc.Server

	// GRPCMaxConnectionAge is the maximum age of a client connection, before GoAway is sent.
	// This is useful for L4 loadbalancing to ensure rebalancing after scaling.
	GRPCMaxConnectionAge = flag.Duration("grpc_max_connection_age", time.Duration(math.MaxInt64), "Maximum age of a client connection before GoAway is sent.")

	// GRPCMaxConnectionAgeGrace is an additional grace period after GRPCMaxConnectionAge, after which
	// connections are forcibly closed.
	GRPCMaxConnectionAgeGrace = flag.Duration("grpc_max_connection_age_grace", time.Duration(math.MaxInt64), "Additional grace period after grpc_max_connection_age, after which connections are forcibly closed.")

	// GRPCInitialConnWindowSize ServerOption that sets window size for a connection.
	// The lower bound for window size is 64K and any value smaller than that will be ignored.
	GRPCInitialConnWindowSize = flag.Int("grpc_server_initial_conn_window_size", 0, "gRPC server initial connection window size")

	// GRPCInitialWindowSize ServerOption that sets window size for stream.
	// The lower bound for window size is 64K and any value smaller than that will be ignored.
	GRPCInitialWindowSize = flag.Int("grpc_server_initial_window_size", 0, "gRPC server initial window size")

	// EnforcementPolicy MinTime that sets the keepalive enforcement policy on the server.
	// This is the minimum amount of time a client should wait before sending a keepalive ping.
	GRPCKeepAliveEnforcementPolicyMinTime = flag.Duration("grpc_server_keepalive_enforcement_policy_min_time", 10*time.Second, "gRPC server minimum keepalive time")

	// EnforcementPolicy PermitWithoutStream - If true, server allows keepalive pings
	// even when there are no active streams (RPCs). If false, and client sends ping when
	// there are no active streams, server will send GOAWAY and close the connection.
	GRPCKeepAliveEnforcementPolicyPermitWithoutStream = flag.Bool("grpc_server_keepalive_enforcement_policy_permit_without_stream", false, "gRPC server permit client keepalive pings even when there are no active streams (RPCs)")

	authPlugin Authenticator
)

// isGRPCEnabled returns true if gRPC server is set
func isGRPCEnabled() bool {
	if GRPCPort != nil && *GRPCPort != 0 {
		return true
	}

	if SocketFile != nil && *SocketFile != "" {
		return true
	}

	return false
}

// createGRPCServer create the gRPC server we will be using.
// It has to be called after flags are parsed, but before
// services register themselves.
func createGRPCServer() {
	// skip if not registered
	if !isGRPCEnabled() {
		log.Infof("Skipping gRPC server creation")
		return
	}

	grpccommon.EnableTracingOpt()

	var opts []grpc.ServerOption
	if GRPCPort != nil && *GRPCCert != "" && *GRPCKey != "" {
		config, err := vttls.ServerConfig(*GRPCCert, *GRPCKey, *GRPCCA, *GRPCCRL, *GRPCServerCA, tls.VersionTLS12)
		if err != nil {
			log.Exitf("Failed to log gRPC cert/key/ca: %v", err)
		}

		// create the creds server options
		creds := credentials.NewTLS(config)
		if *GRPCEnableOptionalTLS {
			log.Warning("Optional TLS is active. Plain-text connections will be accepted")
			creds = grpcoptionaltls.New(creds)
		}
		opts = []grpc.ServerOption{grpc.Creds(creds)}
	}
	// Override the default max message size for both send and receive
	// (which is 4 MiB in gRPC 1.0.0).
	// Large messages can occur when users try to insert or fetch very big
	// rows. If they hit the limit, they'll see the following error:
	// grpc: received message length XXXXXXX exceeding the max size 4194304
	// Note: For gRPC 1.0.0 it's sufficient to set the limit on the server only
	// because it's not enforced on the client side.
	log.Infof("Setting grpc max message size to %d", *grpccommon.MaxMessageSize)
	opts = append(opts, grpc.MaxRecvMsgSize(*grpccommon.MaxMessageSize))
	opts = append(opts, grpc.MaxSendMsgSize(*grpccommon.MaxMessageSize))

	if *GRPCInitialConnWindowSize != 0 {
		log.Infof("Setting grpc server initial conn window size to %d", int32(*GRPCInitialConnWindowSize))
		opts = append(opts, grpc.InitialConnWindowSize(int32(*GRPCInitialConnWindowSize)))
	}

	if *GRPCInitialWindowSize != 0 {
		log.Infof("Setting grpc server initial window size to %d", int32(*GRPCInitialWindowSize))
		opts = append(opts, grpc.InitialWindowSize(int32(*GRPCInitialWindowSize)))
	}

	ep := keepalive.EnforcementPolicy{
		MinTime:             *GRPCKeepAliveEnforcementPolicyMinTime,
		PermitWithoutStream: *GRPCKeepAliveEnforcementPolicyPermitWithoutStream,
	}
	opts = append(opts, grpc.KeepaliveEnforcementPolicy(ep))

	if GRPCMaxConnectionAge != nil {
		ka := keepalive.ServerParameters{
			MaxConnectionAge: *GRPCMaxConnectionAge,
		}
		if GRPCMaxConnectionAgeGrace != nil {
			ka.MaxConnectionAgeGrace = *GRPCMaxConnectionAgeGrace
		}
		opts = append(opts, grpc.KeepaliveParams(ka))
	}

	opts = append(opts, interceptors()...)

	GRPCServer = grpc.NewServer(opts...)
}

// We can only set a ServerInterceptor once, so we chain multiple interceptors into one
func interceptors() []grpc.ServerOption {
	interceptors := &serverInterceptorBuilder{}

	if *GRPCAuth != "" {
		log.Infof("enabling auth plugin %v", *GRPCAuth)
		pluginInitializer := GetAuthenticator(*GRPCAuth)
		authPluginImpl, err := pluginInitializer()
		if err != nil {
			log.Fatalf("Failed to load auth plugin: %v", err)
		}
		authPlugin = authPluginImpl
		interceptors.Add(authenticatingStreamInterceptor, authenticatingUnaryInterceptor)
	}

	if *grpccommon.EnableGRPCPrometheus {
		interceptors.Add(grpc_prometheus.StreamServerInterceptor, grpc_prometheus.UnaryServerInterceptor)
	}

	trace.AddGrpcServerOptions(interceptors.Add)

	return interceptors.Build()
}

func serveGRPC() {
	if *grpccommon.EnableGRPCPrometheus {
		grpc_prometheus.Register(GRPCServer)
		grpc_prometheus.EnableHandlingTimeHistogram()
	}
	// skip if not registered
	if GRPCPort == nil || *GRPCPort == 0 {
		return
	}

	// register reflection to support list calls :)
	reflection.Register(GRPCServer)

	// register health service to support health checks
	healthServer := health.NewServer()
	healthpb.RegisterHealthServer(GRPCServer, healthServer)

	// listen on the port
	log.Infof("Listening for gRPC calls on port %v", *GRPCPort)
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *GRPCPort))
	if err != nil {
		log.Exitf("Cannot listen on port %v for gRPC: %v", *GRPCPort, err)
	}

	// and serve on it
	// NOTE: Before we call Serve(), all services must have registered themselves
	//       with "GRPCServer". This is the case because go/vt/servenv/run.go
	//       runs all OnRun() hooks after createGRPCServer() and before
	//       serveGRPC(). If this was not the case, the binary would crash with
	//       the error "grpc: Server.RegisterService after Server.Serve".
	go func() {
		err := GRPCServer.Serve(listener)
		if err != nil {
			log.Exitf("Failed to start grpc server: %v", err)
		}
	}()

	OnTermSync(func() {
		log.Info("Initiated graceful stop of gRPC server")
		GRPCServer.GracefulStop()
		log.Info("gRPC server stopped")
	})
}

// GRPCCheckServiceMap returns if we should register a gRPC service
// (and also logs how to enable / disable it)
func GRPCCheckServiceMap(name string) bool {
	// Silently fail individual services if gRPC is not enabled in
	// the first place (either on a grpc port or on the socket file)
	if !isGRPCEnabled() {
		return false
	}

	// then check ServiceMap
	return CheckServiceMap("grpc", name)
}

func authenticatingStreamInterceptor(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	newCtx, err := authPlugin.Authenticate(stream.Context(), info.FullMethod)

	if err != nil {
		return err
	}

	wrapped := WrapServerStream(stream)
	wrapped.WrappedContext = newCtx
	return handler(srv, wrapped)
}

func authenticatingUnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	newCtx, err := authPlugin.Authenticate(ctx, info.FullMethod)
	if err != nil {
		return nil, err
	}

	return handler(newCtx, req)
}

// WrappedServerStream is based on the service stream wrapper from: https://github.com/grpc-ecosystem/go-grpc-middleware
type WrappedServerStream struct {
	grpc.ServerStream
	WrappedContext context.Context
}

// Context returns the wrapper's WrappedContext, overwriting the nested grpc.ServerStream.Context()
func (w *WrappedServerStream) Context() context.Context {
	return w.WrappedContext
}

// WrapServerStream returns a ServerStream that has the ability to overwrite context.
func WrapServerStream(stream grpc.ServerStream) *WrappedServerStream {
	if existing, ok := stream.(*WrappedServerStream); ok {
		return existing
	}
	return &WrappedServerStream{ServerStream: stream, WrappedContext: stream.Context()}
}

// serverInterceptorBuilder chains together multiple ServerInterceptors
type serverInterceptorBuilder struct {
	streamInterceptors []grpc.StreamServerInterceptor
	unaryInterceptors  []grpc.UnaryServerInterceptor
}

// Add adds interceptors to the builder
func (collector *serverInterceptorBuilder) Add(s grpc.StreamServerInterceptor, u grpc.UnaryServerInterceptor) {
	collector.streamInterceptors = append(collector.streamInterceptors, s)
	collector.unaryInterceptors = append(collector.unaryInterceptors, u)
}

// AddUnary adds a single unary interceptor to the builder
func (collector *serverInterceptorBuilder) AddUnary(u grpc.UnaryServerInterceptor) {
	collector.unaryInterceptors = append(collector.unaryInterceptors, u)
}

// Build returns DialOptions to add to the grpc.Dial call
func (collector *serverInterceptorBuilder) Build() []grpc.ServerOption {
	log.Infof("Building interceptors with %d unary interceptors and %d stream interceptors", len(collector.unaryInterceptors), len(collector.streamInterceptors))
	switch len(collector.unaryInterceptors) + len(collector.streamInterceptors) {
	case 0:
		return []grpc.ServerOption{}
	default:
		return []grpc.ServerOption{
			grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(collector.unaryInterceptors...)),
			grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(collector.streamInterceptors...)),
		}
	}
}
<commit_msg>Set serving status to registered services<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package servenv

import (
	"crypto/tls"
	"flag"
	"fmt"
	"math"
	"net"
	"time"

	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"

	"vitess.io/vitess/go/trace"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/reflection"

	"context"

	"vitess.io/vitess/go/vt/grpccommon"
	"vitess.io/vitess/go/vt/grpcoptionaltls"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/vttls"
)

// This file handles gRPC server, on its own port.
// Clients register servers, based on service map:
//
// servenv.RegisterGRPCFlags()
// servenv.OnRun(func() {
//   if servenv.GRPCCheckServiceMap("XXX") {
//     pb.RegisterXXX(servenv.GRPCServer, XXX)
//   }
// }
//
// Note servenv.GRPCServer can only be used in servenv.OnRun,
// and not before, as it is initialized right before calling OnRun.
var (
	// GRPCPort is the port to listen on for gRPC. If not set or zero, don't listen.
	GRPCPort = flag.Int("grpc_port", 0, "Port to listen on for gRPC calls")

	// GRPCCert is the cert to use if TLS is enabled
	GRPCCert = flag.String("grpc_cert", "", "server certificate to use for gRPC connections, requires grpc_key, enables TLS")

	// GRPCKey is the key to use if TLS is enabled
	GRPCKey = flag.String("grpc_key", "", "server private key to use for gRPC connections, requires grpc_cert, enables TLS")

	// GRPCCA is the CA to use if TLS is enabled
	GRPCCA = flag.String("grpc_ca", "", "server CA to use for gRPC connections, requires TLS, and enforces client certificate check")

	// GRPCCRL is the CRL (Certificate Revocation List) to use if TLS is enabled
	GRPCCRL = flag.String("grpc_crl", "", "path to a certificate revocation list in PEM format, client certificates will be further verified against this file during TLS handshake")

	GRPCEnableOptionalTLS = flag.Bool("grpc_enable_optional_tls", false, "enable optional TLS mode when a server accepts both TLS and plain-text connections on the same port")

	// GRPCServerCA if specified will combine server cert and server CA
	GRPCServerCA = flag.String("grpc_server_ca", "", "path to server CA in PEM format, which will be combine with server cert, return full certificate chain to clients")

	// GRPCAuth which auth plugin to use (at the moment now only static is supported)
	GRPCAuth = flag.String("grpc_auth_mode", "", "Which auth plugin implementation to use (eg: static)")

	// GRPCServer is the global server to serve gRPC.
	GRPCServer *grpc.Server

	// GRPCMaxConnectionAge is the maximum age of a client connection, before GoAway is sent.
	// This is useful for L4 loadbalancing to ensure rebalancing after scaling.
	GRPCMaxConnectionAge = flag.Duration("grpc_max_connection_age", time.Duration(math.MaxInt64), "Maximum age of a client connection before GoAway is sent.")

	// GRPCMaxConnectionAgeGrace is an additional grace period after GRPCMaxConnectionAge, after which
	// connections are forcibly closed.
	GRPCMaxConnectionAgeGrace = flag.Duration("grpc_max_connection_age_grace", time.Duration(math.MaxInt64), "Additional grace period after grpc_max_connection_age, after which connections are forcibly closed.")

	// GRPCInitialConnWindowSize ServerOption that sets window size for a connection.
	// The lower bound for window size is 64K and any value smaller than that will be ignored.
	GRPCInitialConnWindowSize = flag.Int("grpc_server_initial_conn_window_size", 0, "gRPC server initial connection window size")

	// GRPCInitialWindowSize ServerOption that sets window size for stream.
	// The lower bound for window size is 64K and any value smaller than that will be ignored.
	GRPCInitialWindowSize = flag.Int("grpc_server_initial_window_size", 0, "gRPC server initial window size")

	// EnforcementPolicy MinTime that sets the keepalive enforcement policy on the server.
	// This is the minimum amount of time a client should wait before sending a keepalive ping.
	GRPCKeepAliveEnforcementPolicyMinTime = flag.Duration("grpc_server_keepalive_enforcement_policy_min_time", 10*time.Second, "gRPC server minimum keepalive time")

	// EnforcementPolicy PermitWithoutStream - If true, server allows keepalive pings
	// even when there are no active streams (RPCs). If false, and client sends ping when
	// there are no active streams, server will send GOAWAY and close the connection.
	GRPCKeepAliveEnforcementPolicyPermitWithoutStream = flag.Bool("grpc_server_keepalive_enforcement_policy_permit_without_stream", false, "gRPC server permit client keepalive pings even when there are no active streams (RPCs)")

	authPlugin Authenticator
)

// isGRPCEnabled returns true if gRPC server is set
func isGRPCEnabled() bool {
	if GRPCPort != nil && *GRPCPort != 0 {
		return true
	}

	if SocketFile != nil && *SocketFile != "" {
		return true
	}

	return false
}

// createGRPCServer create the gRPC server we will be using.
// It has to be called after flags are parsed, but before
// services register themselves.
func createGRPCServer() {
	// skip if not registered
	if !isGRPCEnabled() {
		log.Infof("Skipping gRPC server creation")
		return
	}

	grpccommon.EnableTracingOpt()

	var opts []grpc.ServerOption
	if GRPCPort != nil && *GRPCCert != "" && *GRPCKey != "" {
		config, err := vttls.ServerConfig(*GRPCCert, *GRPCKey, *GRPCCA, *GRPCCRL, *GRPCServerCA, tls.VersionTLS12)
		if err != nil {
			log.Exitf("Failed to log gRPC cert/key/ca: %v", err)
		}

		// create the creds server options
		creds := credentials.NewTLS(config)
		if *GRPCEnableOptionalTLS {
			log.Warning("Optional TLS is active. Plain-text connections will be accepted")
			creds = grpcoptionaltls.New(creds)
		}
		opts = []grpc.ServerOption{grpc.Creds(creds)}
	}
	// Override the default max message size for both send and receive
	// (which is 4 MiB in gRPC 1.0.0).
	// Large messages can occur when users try to insert or fetch very big
	// rows. If they hit the limit, they'll see the following error:
	// grpc: received message length XXXXXXX exceeding the max size 4194304
	// Note: For gRPC 1.0.0 it's sufficient to set the limit on the server only
	// because it's not enforced on the client side.
	log.Infof("Setting grpc max message size to %d", *grpccommon.MaxMessageSize)
	opts = append(opts, grpc.MaxRecvMsgSize(*grpccommon.MaxMessageSize))
	opts = append(opts, grpc.MaxSendMsgSize(*grpccommon.MaxMessageSize))

	if *GRPCInitialConnWindowSize != 0 {
		log.Infof("Setting grpc server initial conn window size to %d", int32(*GRPCInitialConnWindowSize))
		opts = append(opts, grpc.InitialConnWindowSize(int32(*GRPCInitialConnWindowSize)))
	}

	if *GRPCInitialWindowSize != 0 {
		log.Infof("Setting grpc server initial window size to %d", int32(*GRPCInitialWindowSize))
		opts = append(opts, grpc.InitialWindowSize(int32(*GRPCInitialWindowSize)))
	}

	ep := keepalive.EnforcementPolicy{
		MinTime:             *GRPCKeepAliveEnforcementPolicyMinTime,
		PermitWithoutStream: *GRPCKeepAliveEnforcementPolicyPermitWithoutStream,
	}
	opts = append(opts, grpc.KeepaliveEnforcementPolicy(ep))

	if GRPCMaxConnectionAge != nil {
		ka := keepalive.ServerParameters{
			MaxConnectionAge: *GRPCMaxConnectionAge,
		}
		if GRPCMaxConnectionAgeGrace != nil {
			ka.MaxConnectionAgeGrace = *GRPCMaxConnectionAgeGrace
		}
		opts = append(opts, grpc.KeepaliveParams(ka))
	}

	opts = append(opts, interceptors()...)

	GRPCServer = grpc.NewServer(opts...)
}

// We can only set a ServerInterceptor once, so we chain multiple interceptors into one
func interceptors() []grpc.ServerOption {
	interceptors := &serverInterceptorBuilder{}

	if *GRPCAuth != "" {
		log.Infof("enabling auth plugin %v", *GRPCAuth)
		pluginInitializer := GetAuthenticator(*GRPCAuth)
		authPluginImpl, err := pluginInitializer()
		if err != nil {
			log.Fatalf("Failed to load auth plugin: %v", err)
		}
		authPlugin = authPluginImpl
		interceptors.Add(authenticatingStreamInterceptor, authenticatingUnaryInterceptor)
	}

	if *grpccommon.EnableGRPCPrometheus {
		interceptors.Add(grpc_prometheus.StreamServerInterceptor, grpc_prometheus.UnaryServerInterceptor)
	}

	trace.AddGrpcServerOptions(interceptors.Add)

	return interceptors.Build()
}

func serveGRPC() {
	if *grpccommon.EnableGRPCPrometheus {
		grpc_prometheus.Register(GRPCServer)
		grpc_prometheus.EnableHandlingTimeHistogram()
	}
	// skip if not registered
	if GRPCPort == nil || *GRPCPort == 0 {
		return
	}

	// register reflection to support list calls :)
	reflection.Register(GRPCServer)

	// register health service to support health checks
	healthServer := health.NewServer()
	healthpb.RegisterHealthServer(GRPCServer, healthServer)

	for service := range GRPCServer.GetServiceInfo() {
		healthServer.SetServingStatus(service, healthpb.HealthCheckResponse_SERVING)
	}

	// listen on the port
	log.Infof("Listening for gRPC calls on port %v", *GRPCPort)
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *GRPCPort))
	if err != nil {
		log.Exitf("Cannot listen on port %v for gRPC: %v", *GRPCPort, err)
	}

	// and serve on it
	// NOTE: Before we call Serve(), all services must have registered themselves
	//       with "GRPCServer". This is the case because go/vt/servenv/run.go
	//       runs all OnRun() hooks after createGRPCServer() and before
	//       serveGRPC(). If this was not the case, the binary would crash with
	//       the error "grpc: Server.RegisterService after Server.Serve".
	go func() {
		err := GRPCServer.Serve(listener)
		if err != nil {
			log.Exitf("Failed to start grpc server: %v", err)
		}
	}()

	OnTermSync(func() {
		log.Info("Initiated graceful stop of gRPC server")
		GRPCServer.GracefulStop()
		log.Info("gRPC server stopped")
	})
}

// GRPCCheckServiceMap returns if we should register a gRPC service
// (and also logs how to enable / disable it)
func GRPCCheckServiceMap(name string) bool {
	// Silently fail individual services if gRPC is not enabled in
	// the first place (either on a grpc port or on the socket file)
	if !isGRPCEnabled() {
		return false
	}

	// then check ServiceMap
	return CheckServiceMap("grpc", name)
}

func authenticatingStreamInterceptor(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	newCtx, err := authPlugin.Authenticate(stream.Context(), info.FullMethod)

	if err != nil {
		return err
	}

	wrapped := WrapServerStream(stream)
	wrapped.WrappedContext = newCtx
	return handler(srv, wrapped)
}

func authenticatingUnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	newCtx, err := authPlugin.Authenticate(ctx, info.FullMethod)
	if err != nil {
		return nil, err
	}

	return handler(newCtx, req)
}

// WrappedServerStream is based on the service stream wrapper from: https://github.com/grpc-ecosystem/go-grpc-middleware
type WrappedServerStream struct {
	grpc.ServerStream
	WrappedContext context.Context
}

// Context returns the wrapper's WrappedContext, overwriting the nested grpc.ServerStream.Context()
func (w *WrappedServerStream) Context() context.Context {
	return w.WrappedContext
}

// WrapServerStream returns a ServerStream that has the ability to overwrite context.
func WrapServerStream(stream grpc.ServerStream) *WrappedServerStream {
	if existing, ok := stream.(*WrappedServerStream); ok {
		return existing
	}
	return &WrappedServerStream{ServerStream: stream, WrappedContext: stream.Context()}
}

// serverInterceptorBuilder chains together multiple ServerInterceptors
type serverInterceptorBuilder struct {
	streamInterceptors []grpc.StreamServerInterceptor
	unaryInterceptors  []grpc.UnaryServerInterceptor
}

// Add adds interceptors to the builder
func (collector *serverInterceptorBuilder) Add(s grpc.StreamServerInterceptor, u grpc.UnaryServerInterceptor) {
	collector.streamInterceptors = append(collector.streamInterceptors, s)
	collector.unaryInterceptors = append(collector.unaryInterceptors, u)
}

// AddUnary adds a single unary interceptor to the builder
func (collector *serverInterceptorBuilder) AddUnary(u grpc.UnaryServerInterceptor) {
	collector.unaryInterceptors = append(collector.unaryInterceptors, u)
}

// Build returns DialOptions to add to the grpc.Dial call
func (collector *serverInterceptorBuilder) Build() []grpc.ServerOption {
	log.Infof("Building interceptors with %d unary interceptors and %d stream interceptors", len(collector.unaryInterceptors), len(collector.streamInterceptors))
	switch len(collector.unaryInterceptors) + len(collector.streamInterceptors) {
	case 0:
		return []grpc.ServerOption{}
	default:
		return []grpc.ServerOption{
			grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(collector.unaryInterceptors...)),
			grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(collector.streamInterceptors...)),
		}
	}
}
<|endoftext|>


<commit_before>package api

import (
	. "github.com/gocircuit/circuit/gocircuit.org/render"
)

func RenderProcessPage() string {
	figs := A{
		"FigMkProc": RenderFigurePngSvg("Process elements execute OS processes on behalf of the user.", "mkproc", "600px"),
	}
	return RenderHtml("Using processes", Render(processBody, figs))
}

const processBody = `

Using processes

You can start an OS process on any host in your cluster by creating a
new process element at an anchor of your choosing that is a descendant of the
host's server anchor. The created process element becomes your interface to the
underlying OS process.

Creating a process

Suppose the variable anchor holds an Anchor object,
corresponding to a path in the anchor hierarchy that has no element attached to it.
For instance, say we obtained anchor like this:

	anchor := root.Walk([]string{"Xe2ac4c8c83976ce6", "job", "demo"})

This anchor corresponds to the path /Xe2ac4c8c83976ce6/job/demo.
(Read more on navigating anchors here.)

To create a new process element and attach it to anchor,
we use the anchor's MakeProc method:

	MakeProc(cmd Cmd) (Proc, error)

MakeProc will start a new process on the host /Xe2ac4c8c83976ce6,
as specified by the command parameter cmd. If successful, it will create a
corresponding process element and attach it to the anchor. MakeProc returns the
newly created process element (of type Proc) as well as an
application error (of type error), or it panics if a
system error occurs.

An application error can occur in one of two cases. Either the anchor already has another element
attached to it, or the process execution was rejected by the host OS (due to a missing binary or
	insufficient permissions, for example).

MakeProc never blocks.

The command parameter, of type Cmd, specifies the standard POSIX-level execution
parameters and an additional parameter called Scrub:

type Cmd struct {
	Env []string
	Dir string
	Path string
	Args []string
	Scrub bool
}

If Scrub is set, the process element will automatically be detached from the anchor
and discarded, as soon as the underlying OS process exits. If Scrub is not set,
the process element will remain attached to the anchor even after the underlying OS process dies.
The latter regime is useful when one wants to start a job and return at a later time to check if
the job has already completed and what was its exit status. Furthermore, removing process elements
explicitly (rather than automatically) is a way of explicit accounting on the user's side. Thus
this regime is particularly well suited for applications that control circuit processes
programmatically (as opposed to manually).

Example

For instance, the following code executes the GNU list command:

	proc, err := a.MakeProc(
		cli.Cmd{
			Env: []string{"TERM=xterm"},
			Dir: "/",
			Path: "/bin/ls",
			Args: []string{"-l", "/"},
			Scrub: true,
		},
	)

The following picture tries to illustrate the relationship between the
process element and the underlying OS process itself.

{{.FigMkProc}}

Controlling the standard file descriptors of a process

After its invocation, MakeProc returns immediately,
while the underlying OS process is executing on the host machine.

After a successful execution the user is obligated, by the POSIX
standard, to take care of the standard input, output and error
streams of the underlying process. (For instance, if the standard
	input is not written to or closed, or if the output is not
	read from, some programs will pause in waiting.)

The standard streams of the executed process can be retrieved
with the following methods of the process element:

	Stdin() io.WriteCloser
	Stdout() io.ReadCloser
	Stderr() io.ReadCloser

It is allowed to close the standard output and error at any point
into the stream. This will result in discarding all remaining data
in the stream, without blocking the underlying process.

Eventually, the user is responsible for closing all standard streams
otherwise the underlying process will block and not exit.

Sending signals and killing processes

You can send a POSIX signal to the underlying process
at any point (asynchronously) using:

	Signal(sig string) error

The sig string must be one of the following recognized
signal names:
	ABRT,
	ALRM,
	BUS,
	CHLD,
	CONT,
	FPE,
	HUP,
	ILL,
	INT,
	IO,
	IOT,
	KILL,
	PIPE,
	PROF,
	QUIT,
	SEGV,
	STOP,
	SYS,
	TERM,
	TRAP,
	TSTP,
	TTIN,
	TTOU,
	URG,
	USR1,
	USR2,
	VTALRM,
	WINCH,
	XCPU,
	XFSZ.

Querying the status of a process asynchronously

You can query the status of a process asynchronously, using:

	Peek() ProcStat

The returned structure includes the command that started the process, a phase string describing the state of the
process and, in the event that the process has exited, an exit error value or nil on successful exit.

	type ProcStat struct {
		Cmd Cmd
		Exit error
		Phase string
	}

The phase string takes on one of the following values:
running,
exited,
stopped,
signaled,
 continued.

Waiting until a process exits

	Wait() (ProcStat, error)

        `
<commit_msg>proc doc complete<commit_after>package api

import (
	. "github.com/gocircuit/circuit/gocircuit.org/render"
)

func RenderProcessPage() string {
	figs := A{
		"FigMkProc": RenderFigurePngSvg("Process elements execute OS processes on behalf of the user.", "mkproc", "600px"),
	}
	return RenderHtml("Using processes", Render(processBody, figs))
}

const processBody = `

Using processes

You can start an OS process on any host in your cluster by creating a
new process element at an anchor of your choosing that is a descendant of the
host's server anchor. The created process element becomes your interface to the
underlying OS process.

Creating a process

Suppose the variable anchor holds an Anchor object,
corresponding to a path in the anchor hierarchy that has no element attached to it.
For instance, say we obtained anchor like this:

	anchor := root.Walk([]string{"Xe2ac4c8c83976ce6", "job", "demo"})

This anchor corresponds to the path /Xe2ac4c8c83976ce6/job/demo.
(Read more on navigating anchors here.)

To create a new process element and attach it to anchor,
we use the anchor's MakeProc method:

	MakeProc(cmd Cmd) (Proc, error)

MakeProc will start a new process on the host /Xe2ac4c8c83976ce6,
as specified by the command parameter cmd. If successful, it will create a
corresponding process element and attach it to the anchor. MakeProc returns the
newly created process element (of type Proc) as well as an
application error (of type error), or it panics if a
system error occurs.

An application error can occur in one of two cases. Either the anchor already has another element
attached to it, or the process execution was rejected by the host OS (due to a missing binary or
	insufficient permissions, for example).

MakeProc never blocks.

The command parameter, of type Cmd, specifies the standard POSIX-level execution
parameters and an additional parameter called Scrub:

type Cmd struct {
	Env []string
	Dir string
	Path string
	Args []string
	Scrub bool
}

If Scrub is set, the process element will automatically be detached from the anchor
and discarded, as soon as the underlying OS process exits. If Scrub is not set,
the process element will remain attached to the anchor even after the underlying OS process dies.
The latter regime is useful when one wants to start a job and return at a later time to check if
the job has already completed and what was its exit status. Furthermore, removing process elements
explicitly (rather than automatically) is a way of explicit accounting on the user's side. Thus
this regime is particularly well suited for applications that control circuit processes
programmatically (as opposed to manually).

Regardless of the setting of the Scrub parameter, the user can
use the Scrub method to discard the process element at any point:

	Scrub()

A call to Scrub will detach the process element from its anchor
and discard it, thereby freeing the anchor to attach other elements.
If the underlying OS process is still running, ‘scrubbing’ will not
terminate the process. (If OS process termination is desired, the user
must explicitly send a kill signal to the process, using a Signal
which is described later.)

Example

For instance, the following code executes the GNU list command:

	proc, err := a.MakeProc(
		cli.Cmd{
			Env: []string{"TERM=xterm"},
			Dir: "/",
			Path: "/bin/ls",
			Args: []string{"-l", "/"},
			Scrub: true,
		},
	)

The following picture tries to illustrate the relationship between the
process element and the underlying OS process itself.

{{.FigMkProc}}

Controlling the standard file descriptors of a process

After its invocation, MakeProc returns immediately,
while the underlying OS process is executing on the host machine.

After a successful execution the user is obligated, by the POSIX
standard, to take care of the standard input, output and error
streams of the underlying process. (For instance, if the standard
	input is not written to or closed, or if the output is not
	read from, some programs will pause in waiting.)

The standard streams of the executed process can be retrieved
with the following methods of the process element:

	Stdin() io.WriteCloser
	Stdout() io.ReadCloser
	Stderr() io.ReadCloser

It is allowed to close the standard output and error at any point
into the stream. This will result in discarding all remaining data
in the stream, without blocking the underlying process.

Eventually, the user is responsible for closing all standard streams
otherwise the underlying process will block and not exit.

Sending signals and killing processes

You can send a POSIX signal to the underlying process
at any point (asynchronously) using:

	Signal(sig string) error

The sig string must be one of the following recognized
signal names:
	ABRT,
	ALRM,
	BUS,
	CHLD,
	CONT,
	FPE,
	HUP,
	ILL,
	INT,
	IO,
	IOT,
	KILL,
	PIPE,
	PROF,
	QUIT,
	SEGV,
	STOP,
	SYS,
	TERM,
	TRAP,
	TSTP,
	TTIN,
	TTOU,
	URG,
	USR1,
	USR2,
	VTALRM,
	WINCH,
	XCPU,
	XFSZ.

Querying the status of a process asynchronously

You can query the status of a process asynchronously, using:

	Peek() ProcStat

The returned structure includes the command that started the process, a phase string describing the state of the
process and, in the event that the process has exited, an exit error value or nil on successful exit.

	type ProcStat struct {
		Cmd Cmd
		Exit error
		Phase string
	}

The phase string takes on one of the following values:
running,
exited,
stopped,
signaled,
 continued.

Waiting until a process exits

Finally, you can call Wait asynchronously to block until the process ends:

	Wait() (ProcStat, error)

If you call Wait before the process has exited, the invocation will block
until exit occurs. Otherwise, it will return immediately. In both cases, a process status
structure (described earlier) is returned, which captures the exit state (successful or not) of the
underlying OS process.

Wait can return an application error only in the event that it is interrupted
by a concurring call to Scrub.

        `
<|endoftext|>