<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Julien Vehent jvehent@mozilla.com [:ulfr]

package main

import (
	"bytes"
	"crypto/sha256"
	"net/http"
	"testing"
	"time"

	"go.mozilla.org/hawk"
)

func TestMissingAuthorization(t *testing.T) {
	t.Parallel()

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	_, _, err = ag.authorizeHeader(req)
	if err == nil {
		t.Errorf("expected auth to fail with missing authorization but succeeded")
	}
	if err.Error() != "missing Authorization header" {
		t.Errorf("expected auth to fail with missing authorization but got error: %v", err)
	}
}

func TestBogusAuthorization(t *testing.T) {
	t.Parallel()

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", `Hawk thisisbob="bob", andhereisamac="nVg5STp2fD+P7G3ELmUztb3hP/LQajwD+FDQM7rZvhw=", ts="1453681057"`)
	_, _, err = ag.authorizeHeader(req)
	if err == nil {
		t.Errorf("expected auth to fail with invalid authorization but succeeded")
	}
	if err.Error() != "hawk: invalid mac, missing or empty" {
		t.Errorf("expected auth to fail with no authorization but got error: %v", err)
	}
}

func TestBadPayload(t *testing.T) {
	t.Parallel()

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	authheader := getAuthHeader(req, ag.auths[conf.Authorizations[0].ID].ID, ag.auths[conf.Authorizations[0].ID].Key, sha256.New, id(), "application/json", []byte(`9247oldfjd18weohfa`))
	req.Header.Set("Authorization", authheader)
	_, err = ag.authorize(req, body)
	if err == nil {
		t.Errorf("expected auth to fail with payload validation failed but succeeded")
	}
	if err.Error() != "payload validation failed" {
		t.Errorf("expected auth to fail with payload validation failed but got error: %v", err)
	}
}

func TestExpiredAuth(t *testing.T) {
	t.Parallel()

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", `Hawk id="bob", mac="nVg5STp2fD+P7G3ELmUztb3hP/LQajwD+FDQM7rZvhw=", ts="1453681057", nonce="TKLzwtGS", hash="sL12YYG2CnALd5o5dqHRKjNO0AvgmPPeIqlfZQfszfo=", ext="59d2rtbmji6617pthvwa1h370"`)
	_, _, err = ag.authorizeHeader(req)
	if err == nil {
		t.Errorf("expected auth to fail with expired timestamp but succeeded")
	}
	if err.Error() != hawk.ErrTimestampSkew.Error() {
		t.Errorf("expected auth to fail with expired timestamp but got error: %v", err)
	}
}

func TestDuplicateNonce(t *testing.T) {
	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	authheader := getAuthHeader(req, ag.auths[conf.Authorizations[0].ID].ID, ag.auths[conf.Authorizations[0].ID].Key, sha256.New, id(), "application/json", body)
	req.Header.Set("Authorization", authheader)
	// run it once
	_, _, err = ag.authorizeHeader(req)
	// and run it twice
	_, _, err = ag.authorizeHeader(req)
	if err == nil {
		t.Errorf("expected auth to fail with duplicate nonce, but succeeded")
	}
	if err.Error() != hawk.ErrReplay.Error() {
		t.Errorf("expected auth to fail with duplicate nonces but got error: %v", err)
	}

}

func TestNonceFromLRU(t *testing.T) {
	t.Parallel()

	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", nil)
	if err != nil {
		t.Fatal(err)
	}

	auth1 := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   ag.auths[conf.Authorizations[0].ID].ID,
			Key:  ag.auths[conf.Authorizations[0].ID].Key,
			Hash: sha256.New},
		0)
	req.Header.Set("Authorization", auth1.RequestHeader())
	_, _, err = ag.authorizeHeader(req)
	if err != nil {
		t.Fatalf("error authorizing header for first request: %s", err)
	}

	auth2 := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   ag.auths[conf.Authorizations[0].ID].ID,
			Key:  ag.auths[conf.Authorizations[0].ID].Key,
			Hash: sha256.New},
		0)
	req.Header.Set("Authorization", auth2.RequestHeader())
	_, _, err = ag.authorizeHeader(req)
	if err != nil {
		t.Fatalf("error authorizing header for second request: %s", err)
	}

	if ag.nonces.Contains(auth1.Nonce) {
		t.Errorf("First nonce %q found in cache, should have been removed", auth1.Nonce)
		t.Logf("nonces: %+v", ag.nonces.Keys())
	}
	if !ag.nonces.Contains(auth2.Nonce) {
		t.Errorf("Second nonce %q not found in cache, should have been present", auth2.Nonce)
		t.Logf("nonces: %+v", ag.nonces.Keys())
	}
}

func TestSignerNotFound(t *testing.T) {
	t.Parallel()

	pos, err := ag.getSignerID(`unknown018qoegdxc`, `unkown093ytid`)
	if err == nil || pos != -1 {
		t.Errorf("expected to fail lookup up a signer but succeeded")
	}
}

func TestDefaultSignerNotFound(t *testing.T) {
	t.Parallel()

	pos, err := ag.getSignerID(`unknown018qoegdxc`, ``)
	if err == nil || pos != -1 {
		t.Errorf("expected to fail lookup up a signer but succeeded")
	}
}

// Two authorizations sharing the same ID should fail
func TestAddDuplicateAuthorization(t *testing.T) {
	t.Parallel()

	var authorizations = []authorization{
		{
			ID: "alice",
		},
		{
			ID: "alice",
		},
	}
	defer func() {
		if e := recover(); e != nil {
			if e != `authorization id 'alice' already defined, duplicates are not permitted` {
				t.Fatalf("expected authorization loading to fail with duplicate error but got: %v", e)
			}
		}
	}()
	tmpag := newAutographer(1)
	tmpag.addSigners(conf.Signers)
	tmpag.addAuthorizations(authorizations)
}

// set an authorization with a ts validity of 2 seconds, then sleep 5 seconds
// to trigger the hawk skew error
func TestHawkTimestampSkewFail(t *testing.T) {
	t.Parallel()

	tmpag := newAutographer(1)
	tmpag.addSigners(conf.Signers)
	tmpag.addAuthorizations([]authorization{
		{
			ID:                    "alice",
			Key:                   "1862300e9bd18eafab2eb8d6",
			HawkTimestampValidity: "2s",
			Signers:               []string{"appkey1"},
		}})

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	authheader := getAuthHeader(req, "alice", "1862300e9bd18eafab2eb8d6", sha256.New, id(), "application/json", body)
	req.Header.Set("Authorization", authheader)
	time.Sleep(5 * time.Second)
	_, _, err = tmpag.authorizeHeader(req)
	if err.Error() != hawk.ErrTimestampSkew.Error() {
		t.Errorf("expected auth to fail with skewed timestamp but got error: %v", err)
	}

}
<commit_msg>test: do not run TestNonceFromLRU in parallel<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Julien Vehent jvehent@mozilla.com [:ulfr]

package main

import (
	"bytes"
	"crypto/sha256"
	"net/http"
	"testing"
	"time"

	"go.mozilla.org/hawk"
)

func TestMissingAuthorization(t *testing.T) {
	t.Parallel()

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	_, _, err = ag.authorizeHeader(req)
	if err == nil {
		t.Errorf("expected auth to fail with missing authorization but succeeded")
	}
	if err.Error() != "missing Authorization header" {
		t.Errorf("expected auth to fail with missing authorization but got error: %v", err)
	}
}

func TestBogusAuthorization(t *testing.T) {
	t.Parallel()

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", `Hawk thisisbob="bob", andhereisamac="nVg5STp2fD+P7G3ELmUztb3hP/LQajwD+FDQM7rZvhw=", ts="1453681057"`)
	_, _, err = ag.authorizeHeader(req)
	if err == nil {
		t.Errorf("expected auth to fail with invalid authorization but succeeded")
	}
	if err.Error() != "hawk: invalid mac, missing or empty" {
		t.Errorf("expected auth to fail with no authorization but got error: %v", err)
	}
}

func TestBadPayload(t *testing.T) {
	t.Parallel()

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	authheader := getAuthHeader(req, ag.auths[conf.Authorizations[0].ID].ID, ag.auths[conf.Authorizations[0].ID].Key, sha256.New, id(), "application/json", []byte(`9247oldfjd18weohfa`))
	req.Header.Set("Authorization", authheader)
	_, err = ag.authorize(req, body)
	if err == nil {
		t.Errorf("expected auth to fail with payload validation failed but succeeded")
	}
	if err.Error() != "payload validation failed" {
		t.Errorf("expected auth to fail with payload validation failed but got error: %v", err)
	}
}

func TestExpiredAuth(t *testing.T) {
	t.Parallel()

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", `Hawk id="bob", mac="nVg5STp2fD+P7G3ELmUztb3hP/LQajwD+FDQM7rZvhw=", ts="1453681057", nonce="TKLzwtGS", hash="sL12YYG2CnALd5o5dqHRKjNO0AvgmPPeIqlfZQfszfo=", ext="59d2rtbmji6617pthvwa1h370"`)
	_, _, err = ag.authorizeHeader(req)
	if err == nil {
		t.Errorf("expected auth to fail with expired timestamp but succeeded")
	}
	if err.Error() != hawk.ErrTimestampSkew.Error() {
		t.Errorf("expected auth to fail with expired timestamp but got error: %v", err)
	}
}

func TestDuplicateNonce(t *testing.T) {
	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	authheader := getAuthHeader(req, ag.auths[conf.Authorizations[0].ID].ID, ag.auths[conf.Authorizations[0].ID].Key, sha256.New, id(), "application/json", body)
	req.Header.Set("Authorization", authheader)
	// run it once
	_, _, err = ag.authorizeHeader(req)
	// and run it twice
	_, _, err = ag.authorizeHeader(req)
	if err == nil {
		t.Errorf("expected auth to fail with duplicate nonce, but succeeded")
	}
	if err.Error() != hawk.ErrReplay.Error() {
		t.Errorf("expected auth to fail with duplicate nonces but got error: %v", err)
	}

}

func TestNonceFromLRU(t *testing.T) {
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", nil)
	if err != nil {
		t.Fatal(err)
	}

	auth1 := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   ag.auths[conf.Authorizations[0].ID].ID,
			Key:  ag.auths[conf.Authorizations[0].ID].Key,
			Hash: sha256.New},
		0)
	req.Header.Set("Authorization", auth1.RequestHeader())
	_, _, err = ag.authorizeHeader(req)
	if err != nil {
		t.Fatalf("error authorizing header for first request: %s", err)
	}

	auth2 := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   ag.auths[conf.Authorizations[0].ID].ID,
			Key:  ag.auths[conf.Authorizations[0].ID].Key,
			Hash: sha256.New},
		0)
	req.Header.Set("Authorization", auth2.RequestHeader())
	_, _, err = ag.authorizeHeader(req)
	if err != nil {
		t.Fatalf("error authorizing header for second request: %s", err)
	}

	if ag.nonces.Contains(auth1.Nonce) {
		t.Errorf("First nonce %q found in cache, should have been removed", auth1.Nonce)
		t.Logf("nonces: %+v", ag.nonces.Keys())
	}
	if !ag.nonces.Contains(auth2.Nonce) {
		t.Errorf("Second nonce %q not found in cache, should have been present", auth2.Nonce)
		t.Logf("nonces: %+v", ag.nonces.Keys())
	}
}

func TestSignerNotFound(t *testing.T) {
	t.Parallel()

	pos, err := ag.getSignerID(`unknown018qoegdxc`, `unkown093ytid`)
	if err == nil || pos != -1 {
		t.Errorf("expected to fail lookup up a signer but succeeded")
	}
}

func TestDefaultSignerNotFound(t *testing.T) {
	t.Parallel()

	pos, err := ag.getSignerID(`unknown018qoegdxc`, ``)
	if err == nil || pos != -1 {
		t.Errorf("expected to fail lookup up a signer but succeeded")
	}
}

// Two authorizations sharing the same ID should fail
func TestAddDuplicateAuthorization(t *testing.T) {
	t.Parallel()

	var authorizations = []authorization{
		{
			ID: "alice",
		},
		{
			ID: "alice",
		},
	}
	defer func() {
		if e := recover(); e != nil {
			if e != `authorization id 'alice' already defined, duplicates are not permitted` {
				t.Fatalf("expected authorization loading to fail with duplicate error but got: %v", e)
			}
		}
	}()
	tmpag := newAutographer(1)
	tmpag.addSigners(conf.Signers)
	tmpag.addAuthorizations(authorizations)
}

// set an authorization with a ts validity of 2 seconds, then sleep 5 seconds
// to trigger the hawk skew error
func TestHawkTimestampSkewFail(t *testing.T) {
	t.Parallel()

	tmpag := newAutographer(1)
	tmpag.addSigners(conf.Signers)
	tmpag.addAuthorizations([]authorization{
		{
			ID:                    "alice",
			Key:                   "1862300e9bd18eafab2eb8d6",
			HawkTimestampValidity: "2s",
			Signers:               []string{"appkey1"},
		}})

	body := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	bodyrdr := bytes.NewReader(body)
	req, err := http.NewRequest("POST", "http://foo.bar/sign/data", bodyrdr)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	authheader := getAuthHeader(req, "alice", "1862300e9bd18eafab2eb8d6", sha256.New, id(), "application/json", body)
	req.Header.Set("Authorization", authheader)
	time.Sleep(5 * time.Second)
	_, _, err = tmpag.authorizeHeader(req)
	if err.Error() != hawk.ErrTimestampSkew.Error() {
		t.Errorf("expected auth to fail with skewed timestamp but got error: %v", err)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/dghubble/sling"
	"net/http"
	"strings"
	"io/ioutil"
)

type SpectoLabSimulation struct {
	Version     string `json:"version"`
	Name        string `json:"name"`
	Description string `json:"description"`
}

type SpectoLab struct {
	Host   string
	Port   string
	ApiKey string
}

func (s *SpectoLab) CreateSimulation(simulationName Simulation) (error) {
	simulation := SpectoLabSimulation{Version: simulationName.Version, Name: simulationName.Name, Description: "A description could go here"}

	url := s.buildUrl("/api/v1/simulations")
	request, err := sling.New().Post(url).BodyJSON(simulation).Add("Authorization", s.buildAuthorizationHeaderValue()).Request()

	if err != nil {
		return err
	}

	response, err := http.DefaultClient.Do(request)
	defer response.Body.Close()

	if err != nil {
		return err
	}

	return nil
}

func (s *SpectoLab) UploadSimulation(simulation Simulation, data []byte) (bool, error) {
	err := s.CreateSimulation(simulation)

	if err != nil {
		return false, err
	}

	url := s.buildUrl(fmt.Sprintf("/api/v1/users/%v/simulations/%v/versions/%v/data", simulation.Vendor,  simulation.Name, simulation.Version))

	request, err := sling.New().Put(url).Add("Authorization", s.buildAuthorizationHeaderValue()).Add("Content-Type", "application/json").Body(strings.NewReader(string(data))).Request()

	if err != nil {
		return false, err
	}

	response, err := http.DefaultClient.Do(request)

	if err != nil {
		return false, err
	}

	defer response.Body.Close()

	return response.StatusCode >= 200 && response.StatusCode <= 299, nil
}

func (s *SpectoLab) GetSimulation(simulation Simulation, overrideHost string) []byte {
	var url string
	if len(overrideHost) > 0 {
		url = s.buildUrl(fmt.Sprintf("/api/v1/users/%v/simulations/%v/versions/%v/data?override-host=%v", simulation.Vendor, simulation.Name, simulation.Version, overrideHost))
	} else {
		url = s.buildUrl(fmt.Sprintf("/api/v1/users/%v/simulations/%v/versions/%v/data", simulation.Vendor, simulation.Name, simulation.Version))
	}

	request, _ := sling.New().Get(url).Add("Authorization", s.buildAuthorizationHeaderValue()).Add("Content-Type", "application/json").Request()
	response, _ := http.DefaultClient.Do(request)
	defer response.Body.Close()

	body, _ := ioutil.ReadAll(response.Body)

	return body
}

func (s *SpectoLab) buildUrl(endpoint string) string {
	return fmt.Sprintf("%v%v", s.buildBaseUrl(), endpoint)
}

func (s *SpectoLab) buildBaseUrl() string {
	if len(s.Port) > 0 {
		return fmt.Sprintf("http://%v:%v", s.Host, s.Port)
	} else {
		return fmt.Sprintf("http://%v", s.Host)
	}
}

func (s *SpectoLab) buildAuthorizationHeaderValue() string {
	return fmt.Sprintf("Bearer %v", s.ApiKey)
}<commit_msg>Moving some stuff around<commit_after>package main

import (
	"fmt"
	"github.com/dghubble/sling"
	"net/http"
	"strings"
	"io/ioutil"
)

type SpectoLab struct {
	Host   string
	Port   string
	ApiKey string
}

type SpectoLabSimulation struct {
	Version     string `json:"version"`
	Name        string `json:"name"`
	Description string `json:"description"`
}

func (s *SpectoLab) CreateSimulation(simulationName Simulation) (error) {
	simulation := SpectoLabSimulation{Version: simulationName.Version, Name: simulationName.Name, Description: "A description could go here"}

	url := s.buildUrl("/api/v1/simulations")
	request, err := sling.New().Post(url).BodyJSON(simulation).Add("Authorization", s.buildAuthorizationHeaderValue()).Request()

	if err != nil {
		return err
	}

	response, err := http.DefaultClient.Do(request)
	defer response.Body.Close()

	if err != nil {
		return err
	}

	return nil
}

func (s *SpectoLab) UploadSimulation(simulation Simulation, data []byte) (bool, error) {
	err := s.CreateSimulation(simulation)

	if err != nil {
		return false, err
	}

	url := s.buildUrl(fmt.Sprintf("/api/v1/users/%v/simulations/%v/versions/%v/data", simulation.Vendor,  simulation.Name, simulation.Version))

	request, err := sling.New().Put(url).Add("Authorization", s.buildAuthorizationHeaderValue()).Add("Content-Type", "application/json").Body(strings.NewReader(string(data))).Request()

	if err != nil {
		return false, err
	}

	response, err := http.DefaultClient.Do(request)

	if err != nil {
		return false, err
	}

	defer response.Body.Close()

	return response.StatusCode >= 200 && response.StatusCode <= 299, nil
}

func (s *SpectoLab) GetSimulation(simulation Simulation, overrideHost string) []byte {
	var url string
	if len(overrideHost) > 0 {
		url = s.buildUrl(fmt.Sprintf("/api/v1/users/%v/simulations/%v/versions/%v/data?override-host=%v", simulation.Vendor, simulation.Name, simulation.Version, overrideHost))
	} else {
		url = s.buildUrl(fmt.Sprintf("/api/v1/users/%v/simulations/%v/versions/%v/data", simulation.Vendor, simulation.Name, simulation.Version))
	}

	request, _ := sling.New().Get(url).Add("Authorization", s.buildAuthorizationHeaderValue()).Add("Content-Type", "application/json").Request()
	response, _ := http.DefaultClient.Do(request)
	defer response.Body.Close()

	body, _ := ioutil.ReadAll(response.Body)

	return body
}

func (s *SpectoLab) buildUrl(endpoint string) string {
	return fmt.Sprintf("%v%v", s.buildBaseUrl(), endpoint)
}

func (s *SpectoLab) buildBaseUrl() string {
	if len(s.Port) > 0 {
		return fmt.Sprintf("http://%v:%v", s.Host, s.Port)
	} else {
		return fmt.Sprintf("http://%v", s.Host)
	}
}

func (s *SpectoLab) buildAuthorizationHeaderValue() string {
	return fmt.Sprintf("Bearer %v", s.ApiKey)
}<|endoftext|>


<commit_before>// Copyright 2020 The LUCI 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 span

import (
	"context"
	"sync"
	"time"

	"cloud.google.com/go/spanner"
	"google.golang.org/grpc/codes"

	"go.chromium.org/luci/common/errors"
)

// Transaction is a common interface of spanner.ReadOnlyTransaction and
// spanner.ReadWriteTransaction.
type Transaction interface {
	// ReadRow reads a single row from the database.
	ReadRow(ctx context.Context, table string, key spanner.Key, columns []string) (*spanner.Row, error)

	// Read reads multiple rows from the database.
	Read(ctx context.Context, table string, keys spanner.KeySet, columns []string) *spanner.RowIterator

	// ReadWithOptions reads multiple rows from the database, and allows
	// customizing options.
	ReadWithOptions(ctx context.Context, table string, keys spanner.KeySet, columns []string, opts *spanner.ReadOptions) *spanner.RowIterator

	// Query reads multiple rows returned by SQL statement.
	Query(ctx context.Context, statement spanner.Statement) *spanner.RowIterator

	// QueryWithOptions reads multiple rows returned by SQL statement, and allows
	// customizing options.
	QueryWithOptions(ctx context.Context, statement spanner.Statement, opts spanner.QueryOptions) *spanner.RowIterator
}

// UseClient installs a Spanner client into the context.
//
// Primarily used by the module initialization code. May be useful in tests as
// well.
func UseClient(ctx context.Context, client *spanner.Client) context.Context {
	return context.WithValue(ctx, &clientContextKey, client)
}

// Apply applies a list of mutations atomically to the database.
//
// Panics if called from inside a read-write transaction. Use BufferWrite to
// apply mutations there.
func Apply(ctx context.Context, ms []*spanner.Mutation, opts ...spanner.ApplyOption) (commitTimestamp time.Time, err error) {
	panicOnNestedRW(ctx)
	return client(ctx).Apply(ctx, ms, opts...)
}

// Single returns a derived context with a single-use read-only transaction.
//
// It provides a read-only snapshot transaction optimized for the case where
// only a single read or query is needed. This is more efficient than using
// ReadOnlyTransaction() for a single read or query.
//
// The transaction object can be obtained through RO(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, etc. wrappers.
//
// Panics if `ctx` already holds a transaction (either read-write or read-only).
func Single(ctx context.Context) context.Context {
	panicOnNestedRO(ctx)
	return setTxnState(ctx, &txnState{ro: client(ctx).Single()})
}

// ReadOnlyTransaction returns a derived context with a read-only transaction.
//
// It can be used for multiple reads from the database. To avoid leaking
// resources on the server this context *must* be canceled as soon as all reads
// are done.
//
// The transaction object can be obtained through RO(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, etc. wrappers.
//
// Panics if `ctx` already holds a transaction (either read-write or read-only).
func ReadOnlyTransaction(ctx context.Context) (context.Context, context.CancelFunc) {
	panicOnNestedRO(ctx)
	txn := client(ctx).ReadOnlyTransaction()
	ctx, cancel := context.WithCancel(setTxnState(ctx, &txnState{ro: txn}))
	return ctx, func() { cancel(); txn.Close() }
}

// ReadWriteTransaction executes a read-write transaction, with retries as
// necessary.
//
// The callback may be called multiple times if Spanner client decides to retry
// the transaction. In particular this happens if the callback returns (perhaps
// wrapped) ABORTED error. This error is returned by Spanner client methods if
// they encounter a stale transaction.
//
// See https://godoc.org/cloud.google.com/go/spanner#ReadWriteTransaction for
// more details.
//
// The callback can access the transaction object via RW(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, BufferWrite, etc. wrappers.
//
// Panics if `ctx` already holds a read-write transaction. Starting a read-write
// transaction from a read-only transaction is OK though, but beware that they
// are completely separate unrelated transactions.
func ReadWriteTransaction(ctx context.Context, f func(ctx context.Context) error) (commitTimestamp time.Time, err error) {
	panicOnNestedRW(ctx)

	var state *txnState

	cts, err := client(ctx).ReadWriteTransaction(ctx, func(ctx context.Context, rw *spanner.ReadWriteTransaction) error {
		state = &txnState{rw: rw}
		err := f(setTxnState(ctx, state))
		if unwrapped := errors.Unwrap(err); spanner.ErrCode(unwrapped) == codes.Aborted {
			err = unwrapped
		}
		return err
	})

	if err == nil {
		state.execCBs(ctx)
	}

	return cts, err
}

// RO returns the current read-only transaction in the context of nil if it's
// not a read-only transactional context.
func RO(ctx context.Context) *spanner.ReadOnlyTransaction {
	if s := getTxnState(ctx); s != nil {
		return s.ro
	}
	return nil
}

// MustRO is like RO except it panics if `ctx` is not read-only transactional.
func MustRO(ctx context.Context) *spanner.ReadOnlyTransaction {
	if ro := RO(ctx); ro != nil {
		return ro
	}
	panic("not a read-only Spanner transactional context")
}

// RW returns the current read-write transaction in the context or nil if it's
// not a read-write transactional context.
func RW(ctx context.Context) *spanner.ReadWriteTransaction {
	if s := getTxnState(ctx); s != nil {
		return s.rw
	}
	return nil
}

// MustRW is like RW except it panics if `ctx` is not read-write transactional.
func MustRW(ctx context.Context) *spanner.ReadWriteTransaction {
	if rw := RW(ctx); rw != nil {
		return rw
	}
	panic("not a read-write Spanner transactional context")
}

// Txn returns an interface that can be used to read data in the current
// read-only or read-write transaction.
//
// Returns nil if `ctx` is not a transactional context.
func Txn(ctx context.Context) Transaction {
	switch s := getTxnState(ctx); {
	case s == nil:
		return nil
	case s.ro != nil:
		return s.ro
	default:
		return s.rw
	}
}

// MustTxn is like Txn except it panics if `ctx` is not transactional.
func MustTxn(ctx context.Context) Transaction {
	if txn := Txn(ctx); txn != nil {
		return txn
	}
	panic("not a transactional Spanner context")
}

// WithoutTxn returns a copy of the context without the transaction in it.
//
// This can be used to spawn separate independent transactions from within
// a transaction.
func WithoutTxn(ctx context.Context) context.Context {
	if getTxnState(ctx) == nil {
		return ctx
	}
	return setTxnState(ctx, nil)
}

// Defer schedules `cb` for execution when the current read-write transaction
// successfully lands.
//
// Intended for a best-effort non-transactional follow up to a successful
// transaction. Note that in presence of failures there's no guarantee the
// callback will be called. For example, the callback won't ever be called if
// the process crashes right after landing the transaction. Or if the
// transaction really landed, but ReadWriteTransaction finished with
// "deadline exceeded" (or some similar) error.
//
// Callbacks are executed sequentially in the reverse order they were deferred.
// They receive the non-transactional version of the context initially passed to
// ReadWriteTransaction.
//
// Panics if the given context is not transactional.
func Defer(ctx context.Context, cb func(context.Context)) {
	state := getTxnState(ctx)
	if state == nil || state.rw == nil {
		panic("not a read-write Spanner transactional context")
	}
	state.deferCB(cb)
}

// ReadRow reads a single row from the database.
//
// It is a shortcut for MustTxn(ctx).ReadRow(ctx, ...). Panics if the context
// is not transactional.
func ReadRow(ctx context.Context, table string, key spanner.Key, columns []string) (*spanner.Row, error) {
	return MustTxn(ctx).ReadRow(ctx, table, key, columns)
}

// Read reads multiple rows from the database.
//
// It is a shortcut for MustTxn(ctx).Read(ctx, ...). Panics if the context
// is not transactional.
func Read(ctx context.Context, table string, keys spanner.KeySet, columns []string) *spanner.RowIterator {
	return MustTxn(ctx).Read(ctx, table, keys, columns)
}

// ReadWithOptions reads multiple rows from the database, and allows customizing
// options.
//
// It is a shortcut for MustTxn(ctx).ReadWithOptions(ctx, ...). Panics if the
// context is not transactional.
func ReadWithOptions(ctx context.Context, table string, keys spanner.KeySet, columns []string, opts *spanner.ReadOptions) *spanner.RowIterator {
	return MustTxn(ctx).ReadWithOptions(ctx, table, keys, columns, opts)
}

// Query reads multiple rows returned by SQL statement.
//
// It is a shortcut for MustTxn(ctx).Query(ctx, ...). Panics if the context is
// not transactional.
func Query(ctx context.Context, statement spanner.Statement) *spanner.RowIterator {
	return MustTxn(ctx).Query(ctx, statement)
}

// BufferWrite adds a list of mutations to the set of updates that will be
// applied when the transaction is committed.
//
// It does not actually apply the write until the transaction is committed, so
// the operation does not block. The effects of the write won't be visible to
// any reads (including reads done in the same transaction) until the
// transaction commits.
//
// It is a wrapper over MustRW(ctx).BufferWrite(...). Panics if the context is
// not read-write transactional.
func BufferWrite(ctx context.Context, ms ...*spanner.Mutation) {
	// BufferWrite just appends mutation to an internal buffer. It fails only if
	// the transaction has already landed. We are OK to panic in this case:
	// calling BufferWrite outside of a transaction is a programming error.
	if err := MustRW(ctx).BufferWrite(ms); err != nil {
		panic(err)
	}
}

// Update executes a DML statement against the database. It returns the number
// of affected rows. Update returns an error if the statement is a query.
// However, the query is executed, and any data read will be validated upon
// commit.
//
// It is a shortcut for MustRW(ctx).Update(...). Panics if the context is not
// read-write transactional.
func Update(ctx context.Context, stmt spanner.Statement) (rowCount int64, err error) {
	return MustRW(ctx).Update(ctx, stmt)
}

// UpdateWithOptions executes a DML statement against the database. It returns
// the number of affected rows. The sql query execution will be optimized based
// on the given query options.
//
// It is a shortcut for MustRW(ctx).Update(...). Panics if the context is not
// read-write transactional.
func UpdateWithOptions(ctx context.Context, stmt spanner.Statement, opts spanner.QueryOptions) (rowCount int64, err error) {
	return MustRW(ctx).UpdateWithOptions(ctx, stmt, opts)
}

////////////////////////////////////////////////////////////////////////////////

var (
	clientContextKey = "go.chromium.org/luci/server/span:client"
	txnContextKey    = "go.chromium.org/luci/server/span:txn"
)

// client returns a Spanner client installed in the context.
//
// Panics if it is not there.
//
// Intentionally private to force all callers to go through package's functions
// like ReadWriteTransaction, ReadOnlyTransaction, Single, etc. since they
// generally add additional functionality on top of the raw Spanner client that
// other LUCI packages assume to be present. Using the Spanner client directly
// may violate such assumptions leading to undefined behavior when multiple
// packages are used together.
func client(ctx context.Context) *spanner.Client {
	cl, _ := ctx.Value(&clientContextKey).(*spanner.Client)
	if cl == nil {
		panic("no Spanner client in the context")
	}
	return cl
}

type txnState struct {
	ro *spanner.ReadOnlyTransaction
	rw *spanner.ReadWriteTransaction

	m   sync.Mutex
	cbs []func(context.Context)
}

func (s *txnState) deferCB(cb func(context.Context)) {
	s.m.Lock()
	s.cbs = append(s.cbs, cb)
	s.m.Unlock()
}

func (s *txnState) execCBs(ctx context.Context) {
	// Note: execCBs happens after ReadWriteTransaction has finished. If it
	// spawned any goroutines, they must have been finished already too (calling
	// Defer from a goroutine that outlives a transaction is rightfully a race).
	// Thus all writes to `s.cbs` are finished already and we also passed some
	// synchronization barrier that waited for the goroutines to join. It's fine
	// to avoid locking s.m in this case saving 200ns on hot code path.
	for i := len(s.cbs) - 1; i >= 0; i-- {
		s.cbs[i](ctx)
	}
}

func setTxnState(ctx context.Context, s *txnState) context.Context {
	return context.WithValue(ctx, &txnContextKey, s)
}

func getTxnState(ctx context.Context) *txnState {
	s, _ := ctx.Value(&txnContextKey).(*txnState)
	return s
}

func panicOnNestedRW(ctx context.Context) {
	if RW(ctx) != nil {
		panic("nested Spanner write transactions are not allowed")
	}
}

func panicOnNestedRO(ctx context.Context) {
	if getTxnState(ctx) != nil {
		panic("nested Spanner read transactions are not allowed")
	}
}
<commit_msg>[span] fix typo<commit_after>// Copyright 2020 The LUCI 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 span

import (
	"context"
	"sync"
	"time"

	"cloud.google.com/go/spanner"
	"google.golang.org/grpc/codes"

	"go.chromium.org/luci/common/errors"
)

// Transaction is a common interface of spanner.ReadOnlyTransaction and
// spanner.ReadWriteTransaction.
type Transaction interface {
	// ReadRow reads a single row from the database.
	ReadRow(ctx context.Context, table string, key spanner.Key, columns []string) (*spanner.Row, error)

	// Read reads multiple rows from the database.
	Read(ctx context.Context, table string, keys spanner.KeySet, columns []string) *spanner.RowIterator

	// ReadWithOptions reads multiple rows from the database, and allows
	// customizing options.
	ReadWithOptions(ctx context.Context, table string, keys spanner.KeySet, columns []string, opts *spanner.ReadOptions) *spanner.RowIterator

	// Query reads multiple rows returned by SQL statement.
	Query(ctx context.Context, statement spanner.Statement) *spanner.RowIterator

	// QueryWithOptions reads multiple rows returned by SQL statement, and allows
	// customizing options.
	QueryWithOptions(ctx context.Context, statement spanner.Statement, opts spanner.QueryOptions) *spanner.RowIterator
}

// UseClient installs a Spanner client into the context.
//
// Primarily used by the module initialization code. May be useful in tests as
// well.
func UseClient(ctx context.Context, client *spanner.Client) context.Context {
	return context.WithValue(ctx, &clientContextKey, client)
}

// Apply applies a list of mutations atomically to the database.
//
// Panics if called from inside a read-write transaction. Use BufferWrite to
// apply mutations there.
func Apply(ctx context.Context, ms []*spanner.Mutation, opts ...spanner.ApplyOption) (commitTimestamp time.Time, err error) {
	panicOnNestedRW(ctx)
	return client(ctx).Apply(ctx, ms, opts...)
}

// Single returns a derived context with a single-use read-only transaction.
//
// It provides a read-only snapshot transaction optimized for the case where
// only a single read or query is needed. This is more efficient than using
// ReadOnlyTransaction() for a single read or query.
//
// The transaction object can be obtained through RO(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, etc. wrappers.
//
// Panics if `ctx` already holds a transaction (either read-write or read-only).
func Single(ctx context.Context) context.Context {
	panicOnNestedRO(ctx)
	return setTxnState(ctx, &txnState{ro: client(ctx).Single()})
}

// ReadOnlyTransaction returns a derived context with a read-only transaction.
//
// It can be used for multiple reads from the database. To avoid leaking
// resources on the server this context *must* be canceled as soon as all reads
// are done.
//
// The transaction object can be obtained through RO(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, etc. wrappers.
//
// Panics if `ctx` already holds a transaction (either read-write or read-only).
func ReadOnlyTransaction(ctx context.Context) (context.Context, context.CancelFunc) {
	panicOnNestedRO(ctx)
	txn := client(ctx).ReadOnlyTransaction()
	ctx, cancel := context.WithCancel(setTxnState(ctx, &txnState{ro: txn}))
	return ctx, func() { cancel(); txn.Close() }
}

// ReadWriteTransaction executes a read-write transaction, with retries as
// necessary.
//
// The callback may be called multiple times if Spanner client decides to retry
// the transaction. In particular this happens if the callback returns (perhaps
// wrapped) ABORTED error. This error is returned by Spanner client methods if
// they encounter a stale transaction.
//
// See https://godoc.org/cloud.google.com/go/spanner#ReadWriteTransaction for
// more details.
//
// The callback can access the transaction object via RW(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, BufferWrite, etc. wrappers.
//
// Panics if `ctx` already holds a read-write transaction. Starting a read-write
// transaction from a read-only transaction is OK though, but beware that they
// are completely separate unrelated transactions.
func ReadWriteTransaction(ctx context.Context, f func(ctx context.Context) error) (commitTimestamp time.Time, err error) {
	panicOnNestedRW(ctx)

	var state *txnState

	cts, err := client(ctx).ReadWriteTransaction(ctx, func(ctx context.Context, rw *spanner.ReadWriteTransaction) error {
		state = &txnState{rw: rw}
		err := f(setTxnState(ctx, state))
		if unwrapped := errors.Unwrap(err); spanner.ErrCode(unwrapped) == codes.Aborted {
			err = unwrapped
		}
		return err
	})

	if err == nil {
		state.execCBs(ctx)
	}

	return cts, err
}

// RO returns the current read-only transaction in the context or nil if it's
// not a read-only transactional context.
func RO(ctx context.Context) *spanner.ReadOnlyTransaction {
	if s := getTxnState(ctx); s != nil {
		return s.ro
	}
	return nil
}

// MustRO is like RO except it panics if `ctx` is not read-only transactional.
func MustRO(ctx context.Context) *spanner.ReadOnlyTransaction {
	if ro := RO(ctx); ro != nil {
		return ro
	}
	panic("not a read-only Spanner transactional context")
}

// RW returns the current read-write transaction in the context or nil if it's
// not a read-write transactional context.
func RW(ctx context.Context) *spanner.ReadWriteTransaction {
	if s := getTxnState(ctx); s != nil {
		return s.rw
	}
	return nil
}

// MustRW is like RW except it panics if `ctx` is not read-write transactional.
func MustRW(ctx context.Context) *spanner.ReadWriteTransaction {
	if rw := RW(ctx); rw != nil {
		return rw
	}
	panic("not a read-write Spanner transactional context")
}

// Txn returns an interface that can be used to read data in the current
// read-only or read-write transaction.
//
// Returns nil if `ctx` is not a transactional context.
func Txn(ctx context.Context) Transaction {
	switch s := getTxnState(ctx); {
	case s == nil:
		return nil
	case s.ro != nil:
		return s.ro
	default:
		return s.rw
	}
}

// MustTxn is like Txn except it panics if `ctx` is not transactional.
func MustTxn(ctx context.Context) Transaction {
	if txn := Txn(ctx); txn != nil {
		return txn
	}
	panic("not a transactional Spanner context")
}

// WithoutTxn returns a copy of the context without the transaction in it.
//
// This can be used to spawn separate independent transactions from within
// a transaction.
func WithoutTxn(ctx context.Context) context.Context {
	if getTxnState(ctx) == nil {
		return ctx
	}
	return setTxnState(ctx, nil)
}

// Defer schedules `cb` for execution when the current read-write transaction
// successfully lands.
//
// Intended for a best-effort non-transactional follow up to a successful
// transaction. Note that in presence of failures there's no guarantee the
// callback will be called. For example, the callback won't ever be called if
// the process crashes right after landing the transaction. Or if the
// transaction really landed, but ReadWriteTransaction finished with
// "deadline exceeded" (or some similar) error.
//
// Callbacks are executed sequentially in the reverse order they were deferred.
// They receive the non-transactional version of the context initially passed to
// ReadWriteTransaction.
//
// Panics if the given context is not transactional.
func Defer(ctx context.Context, cb func(context.Context)) {
	state := getTxnState(ctx)
	if state == nil || state.rw == nil {
		panic("not a read-write Spanner transactional context")
	}
	state.deferCB(cb)
}

// ReadRow reads a single row from the database.
//
// It is a shortcut for MustTxn(ctx).ReadRow(ctx, ...). Panics if the context
// is not transactional.
func ReadRow(ctx context.Context, table string, key spanner.Key, columns []string) (*spanner.Row, error) {
	return MustTxn(ctx).ReadRow(ctx, table, key, columns)
}

// Read reads multiple rows from the database.
//
// It is a shortcut for MustTxn(ctx).Read(ctx, ...). Panics if the context
// is not transactional.
func Read(ctx context.Context, table string, keys spanner.KeySet, columns []string) *spanner.RowIterator {
	return MustTxn(ctx).Read(ctx, table, keys, columns)
}

// ReadWithOptions reads multiple rows from the database, and allows customizing
// options.
//
// It is a shortcut for MustTxn(ctx).ReadWithOptions(ctx, ...). Panics if the
// context is not transactional.
func ReadWithOptions(ctx context.Context, table string, keys spanner.KeySet, columns []string, opts *spanner.ReadOptions) *spanner.RowIterator {
	return MustTxn(ctx).ReadWithOptions(ctx, table, keys, columns, opts)
}

// Query reads multiple rows returned by SQL statement.
//
// It is a shortcut for MustTxn(ctx).Query(ctx, ...). Panics if the context is
// not transactional.
func Query(ctx context.Context, statement spanner.Statement) *spanner.RowIterator {
	return MustTxn(ctx).Query(ctx, statement)
}

// BufferWrite adds a list of mutations to the set of updates that will be
// applied when the transaction is committed.
//
// It does not actually apply the write until the transaction is committed, so
// the operation does not block. The effects of the write won't be visible to
// any reads (including reads done in the same transaction) until the
// transaction commits.
//
// It is a wrapper over MustRW(ctx).BufferWrite(...). Panics if the context is
// not read-write transactional.
func BufferWrite(ctx context.Context, ms ...*spanner.Mutation) {
	// BufferWrite just appends mutation to an internal buffer. It fails only if
	// the transaction has already landed. We are OK to panic in this case:
	// calling BufferWrite outside of a transaction is a programming error.
	if err := MustRW(ctx).BufferWrite(ms); err != nil {
		panic(err)
	}
}

// Update executes a DML statement against the database. It returns the number
// of affected rows. Update returns an error if the statement is a query.
// However, the query is executed, and any data read will be validated upon
// commit.
//
// It is a shortcut for MustRW(ctx).Update(...). Panics if the context is not
// read-write transactional.
func Update(ctx context.Context, stmt spanner.Statement) (rowCount int64, err error) {
	return MustRW(ctx).Update(ctx, stmt)
}

// UpdateWithOptions executes a DML statement against the database. It returns
// the number of affected rows. The sql query execution will be optimized based
// on the given query options.
//
// It is a shortcut for MustRW(ctx).Update(...). Panics if the context is not
// read-write transactional.
func UpdateWithOptions(ctx context.Context, stmt spanner.Statement, opts spanner.QueryOptions) (rowCount int64, err error) {
	return MustRW(ctx).UpdateWithOptions(ctx, stmt, opts)
}

////////////////////////////////////////////////////////////////////////////////

var (
	clientContextKey = "go.chromium.org/luci/server/span:client"
	txnContextKey    = "go.chromium.org/luci/server/span:txn"
)

// client returns a Spanner client installed in the context.
//
// Panics if it is not there.
//
// Intentionally private to force all callers to go through package's functions
// like ReadWriteTransaction, ReadOnlyTransaction, Single, etc. since they
// generally add additional functionality on top of the raw Spanner client that
// other LUCI packages assume to be present. Using the Spanner client directly
// may violate such assumptions leading to undefined behavior when multiple
// packages are used together.
func client(ctx context.Context) *spanner.Client {
	cl, _ := ctx.Value(&clientContextKey).(*spanner.Client)
	if cl == nil {
		panic("no Spanner client in the context")
	}
	return cl
}

type txnState struct {
	ro *spanner.ReadOnlyTransaction
	rw *spanner.ReadWriteTransaction

	m   sync.Mutex
	cbs []func(context.Context)
}

func (s *txnState) deferCB(cb func(context.Context)) {
	s.m.Lock()
	s.cbs = append(s.cbs, cb)
	s.m.Unlock()
}

func (s *txnState) execCBs(ctx context.Context) {
	// Note: execCBs happens after ReadWriteTransaction has finished. If it
	// spawned any goroutines, they must have been finished already too (calling
	// Defer from a goroutine that outlives a transaction is rightfully a race).
	// Thus all writes to `s.cbs` are finished already and we also passed some
	// synchronization barrier that waited for the goroutines to join. It's fine
	// to avoid locking s.m in this case saving 200ns on hot code path.
	for i := len(s.cbs) - 1; i >= 0; i-- {
		s.cbs[i](ctx)
	}
}

func setTxnState(ctx context.Context, s *txnState) context.Context {
	return context.WithValue(ctx, &txnContextKey, s)
}

func getTxnState(ctx context.Context) *txnState {
	s, _ := ctx.Value(&txnContextKey).(*txnState)
	return s
}

func panicOnNestedRW(ctx context.Context) {
	if RW(ctx) != nil {
		panic("nested Spanner write transactions are not allowed")
	}
}

func panicOnNestedRO(ctx context.Context) {
	if getTxnState(ctx) != nil {
		panic("nested Spanner read transactions are not allowed")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/Netflix-Skunkworks/go-jira/jira/cli"
	"github.com/docopt/docopt-go"
	"github.com/op/go-logging"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"strings"
)

var log = logging.MustGetLogger("jira")
var format = "%{color}%{time:2006-01-02T15:04:05.000Z07:00} %{level:-5s} [%{shortfile}]%{color:reset} %{message}"

func main() {
	user := os.Getenv("USER")
	home := os.Getenv("HOME")
	usage := fmt.Sprintf(`
Usage:
  jira [-v ...] [-u USER] [-e URI] [-t FILE] (ls|list) ( [-q JQL] | [-p PROJECT] [-c COMPONENT] [-a ASSIGNEE] [-i ISSUETYPE] [-w WATCHER] [-r REPORTER])
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] view ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] edit ISSUE [--noedit] [-m COMMENT] [-o KEY=VAL]...
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] create [--noedit] [-p PROJECT] [-i ISSUETYPE] [-o KEY=VAL]...
  jira [-v ...] [-u USER] [-e URI] [-b] DUPLICATE dups ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] BLOCKER blocks ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] watch ISSUE [-w WATCHER]
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] (trans|transition) TRANSITION ISSUE [-m COMMENT] [-o KEY=VAL] [--noedit]
  jira [-v ...] [-u USER] [-e URI] [-b] ack ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] close ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] resolve ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] reopen ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] start ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] stop ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] comment ISSUE [-m COMMENT]
  jira [-v ...] [-u USER] [-e URI] [-b] take ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] (assign|give) ISSUE ASSIGNEE
  jira [-v ...] [-u USER] [-e URI] [-t FILE] fields
  jira [-v ...] [-u USER] [-e URI] [-t FILE] issuelinktypes
  jira [-v ...] [-u USER] [-e URI] [-b][-t FILE] transmeta ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] editmeta ISSUE
  jira [-v ...] [-u USER] [-e URI] [-t FILE] issuetypes [-p PROJECT]
  jira [-v ...] [-u USER] [-e URI] [-t FILE] createmeta [-p PROJECT] [-i ISSUETYPE]
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] transitions ISSUE
  jira [-v ...] export-templates [-d DIR]
  jira [-v ...] [-u USER] [-e URI] (b|browse) ISSUE
  jira [-v ...] [-u USER] [-e URI] [-t FILE] login
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] ISSUE

General Options:
  -e --endpoint=URI   URI to use for jira
  -h --help           Show this usage
  -t --template=FILE  Template file to use for output/editing
  -u --user=USER      Username to use for authenticaion (default: %s)
  -v --verbose        Increase output logging
  --version           Show this version

Command Options:
  -a --assignee=USER        Username assigned the issue
  -b --browse               Open your browser to the Jira issue
  -c --component=COMPONENT  Component to Search for
  -d --directory=DIR        Directory to export templates to (default: %s)
  -f --queryfields          Fields that are used in "list" template: (default: summary)
  -i --issuetype=ISSUETYPE  Jira Issue Type (default: Bug)
  -m --comment=COMMENT      Comment message for transition
  -o --override=KEY:VAL     Set custom key/value pairs
  -p --project=PROJECT      Project to Search for
  -q --query=JQL            Jira Query Language expression for the search
  -r --reporter=USER        Reporter to search for
  -w --watcher=USER         Watcher to add to issue (default: %s)
                            or Watcher to search for
`, user, fmt.Sprintf("%s/.jira.d/templates", home), user)

	args, err := docopt.Parse(usage, nil, true, "0.0.2", false, false)
	if err != nil {
		log.Error("Failed to parse options: %s", err)
		os.Exit(1)
	}
	logBackend := logging.NewLogBackend(os.Stderr, "", 0)
	logging.SetBackend(
		logging.NewBackendFormatter(
			logBackend,
			logging.MustStringFormatter(format),
		),
	)
	logging.SetLevel(logging.NOTICE, "")
	if verbose, ok := args["--verbose"]; ok {
		if verbose.(int) > 1 {
			logging.SetLevel(logging.DEBUG, "")
		} else if verbose.(int) > 0 {
			logging.SetLevel(logging.INFO, "")
		}
	}

	log.Info("Args: %v", args)

	opts := make(map[string]string)
	loadConfigs(opts)

	// strip the "--" off the command line options
	// and populate the opts that we pass to the cli ctor
	for key, val := range args {
		if val != nil && strings.HasPrefix(key, "--") {
			opt := key[2:]
			if opt == "override" {
				for _, v := range val.([]string) {
					if strings.Contains(v, "=") {
						kv := strings.SplitN(v, "=", 2)
						opts[kv[0]] = kv[1]
					} else {
						log.Error("Malformed override, expected KEY=VALUE, got %s", v)
						os.Exit(1)
					}
				}
			} else {
				switch v := val.(type) {
				case string:
					opts[opt] = v
				case int:
					opts[opt] = fmt.Sprintf("%d", v)
				case bool:
					opts[opt] = fmt.Sprintf("%t", v)
				}
			}
		}
	}

	// cant use proper [default:x] syntax in docopt
	// because only want to default if the option is not
	// already specified in some .jira.d/config.yml file
	if _, ok := opts["user"]; !ok {
		opts["user"] = user
	}
	if _, ok := opts["queryfields"]; !ok {
		opts["queryfields"] = "summary"
	}
	if _, ok := opts["directory"]; !ok {
		opts["directory"] = fmt.Sprintf("%s/.jira.d/templates", home)
	}

	if _, ok := opts["endpoint"]; !ok {
		log.Error("endpoint option required.  Either use --endpoint or set a enpoint option in your ~/.jira.d/config.yml file")
		os.Exit(1)
	}

	c := cli.New(opts)

	log.Debug("opts: %s", opts)

	validCommand := func(cmd string) bool {
		if val, ok := args[cmd]; ok && val.(bool) {
			return true
		}
		return false
	}

	validOpt := func(opt string, dflt interface{}) interface{} {
		if val, ok := opts[opt]; ok {
			return val
		}
		if dflt == nil {
			log.Error("Missing required option --%s or \"%s\" property override in the config file", opt, opt)
			os.Exit(1)
		}
		return dflt
	}

	setEditing := func(dflt bool) {
		if dflt {
			if val, ok := opts["noedit"]; ok && val == "true" {
				opts["edit"] = "false"
			} else {
				opts["edit"] = "true"
			}
		} else {
			if val, ok := opts["edit"]; ok && val == "true" {
				opts["edit"] = "true"
			} else {
				opts["edit"] = "false"
			}
		}
	}

	if validCommand("login") {
		err = c.CmdLogin()
	} else if validCommand("fields") {
		err = c.CmdFields()
	} else if validCommand("ls") || validCommand("list") {
		err = c.CmdList()
	} else if validCommand("edit") {
		setEditing(true)
		err = c.CmdEdit(args["ISSUE"].(string))
	} else if validCommand("editmeta") {
		err = c.CmdEditMeta(args["ISSUE"].(string))
	} else if validCommand("transmeta") {
		err = c.CmdTransitionMeta(args["ISSUE"].(string))
	} else if validCommand("issuelinktypes") {
		err = c.CmdIssueLinkTypes()
	} else if validCommand("issuetypes") {
		err = c.CmdIssueTypes(validOpt("project", nil).(string))
	} else if validCommand("createmeta") {
		err = c.CmdCreateMeta(
			validOpt("project", nil).(string),
			validOpt("issuetype", "Bug").(string),
		)
	} else if validCommand("create") {
		setEditing(true)
		err = c.CmdCreate(
			validOpt("project", nil).(string),
			validOpt("issuetype", "Bug").(string),
		)
	} else if validCommand("transitions") {
		err = c.CmdTransitions(args["ISSUE"].(string))
	} else if validCommand("blocks") {
		err = c.CmdBlocks(
			args["BLOCKER"].(string),
			args["ISSUE"].(string),
		)
	} else if validCommand("dups") {
		err = c.CmdDups(
			args["DUPLICATE"].(string),
			args["ISSUE"].(string),
		)
	} else if validCommand("watch") {
		err = c.CmdWatch(
			args["ISSUE"].(string),
			validOpt("watcher", user).(string),
		)
	} else if validCommand("trans") || validCommand("transition") {
		setEditing(true)
		err = c.CmdTransition(
			args["ISSUE"].(string),
			args["TRANSITION"].(string),
		)
	} else if validCommand("close") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "close")
	} else if validCommand("ack") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "acknowledge")
	} else if validCommand("reopen") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "reopen")
	} else if validCommand("resolve") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "resolve")
	} else if validCommand("start") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "start")
	} else if validCommand("stop") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "stop")
	} else if validCommand("comment") {
		setEditing(true)
		err = c.CmdComment(args["ISSUE"].(string))
	} else if validCommand("take") {
		err = c.CmdAssign(args["ISSUE"].(string), user)
	} else if validCommand("browse") || validCommand("b") {
		opts["browse"] = "true"
		err = c.Browse(args["ISSUE"].(string))
	} else if validCommand("export-templates") {
		err = c.CmdExportTemplates()
	} else if validCommand("assign") || validCommand("give") {
		err = c.CmdAssign(
			args["ISSUE"].(string),
			args["ASSIGNEE"].(string),
		)
	} else if val, ok := args["ISSUE"]; ok {
		err = c.CmdView(val.(string))
	}

	if err != nil {
		os.Exit(1)
	}
	os.Exit(0)
}

func parseYaml(file string, opts map[string]string) {
	if fh, err := ioutil.ReadFile(file); err == nil {
		log.Debug("Found Config file: %s", file)
		yaml.Unmarshal(fh, &opts)
	}
}

func loadConfigs(opts map[string]string) {
	paths := cli.FindParentPaths(".jira.d/config.yml")
	// prepend
	paths = append([]string{"/etc/jira-cli.yml"}, paths...)

	for _, file := range paths {
		parseYaml(file, opts)
	}
}
<commit_msg>bump version<commit_after>package main

import (
	"fmt"
	"github.com/Netflix-Skunkworks/go-jira/jira/cli"
	"github.com/docopt/docopt-go"
	"github.com/op/go-logging"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"strings"
)

var log = logging.MustGetLogger("jira")
var format = "%{color}%{time:2006-01-02T15:04:05.000Z07:00} %{level:-5s} [%{shortfile}]%{color:reset} %{message}"

func main() {
	user := os.Getenv("USER")
	home := os.Getenv("HOME")
	usage := fmt.Sprintf(`
Usage:
  jira [-v ...] [-u USER] [-e URI] [-t FILE] (ls|list) ( [-q JQL] | [-p PROJECT] [-c COMPONENT] [-a ASSIGNEE] [-i ISSUETYPE] [-w WATCHER] [-r REPORTER])
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] view ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] edit ISSUE [--noedit] [-m COMMENT] [-o KEY=VAL]...
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] create [--noedit] [-p PROJECT] [-i ISSUETYPE] [-o KEY=VAL]...
  jira [-v ...] [-u USER] [-e URI] [-b] DUPLICATE dups ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] BLOCKER blocks ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] watch ISSUE [-w WATCHER]
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] (trans|transition) TRANSITION ISSUE [-m COMMENT] [-o KEY=VAL] [--noedit]
  jira [-v ...] [-u USER] [-e URI] [-b] ack ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] close ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] resolve ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] reopen ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] start ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] stop ISSUE [-m COMMENT] [-o KEY=VAL] [--edit]
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] comment ISSUE [-m COMMENT]
  jira [-v ...] [-u USER] [-e URI] [-b] take ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] (assign|give) ISSUE ASSIGNEE
  jira [-v ...] [-u USER] [-e URI] [-t FILE] fields
  jira [-v ...] [-u USER] [-e URI] [-t FILE] issuelinktypes
  jira [-v ...] [-u USER] [-e URI] [-b][-t FILE] transmeta ISSUE
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] editmeta ISSUE
  jira [-v ...] [-u USER] [-e URI] [-t FILE] issuetypes [-p PROJECT]
  jira [-v ...] [-u USER] [-e URI] [-t FILE] createmeta [-p PROJECT] [-i ISSUETYPE]
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] transitions ISSUE
  jira [-v ...] export-templates [-d DIR]
  jira [-v ...] [-u USER] [-e URI] (b|browse) ISSUE
  jira [-v ...] [-u USER] [-e URI] [-t FILE] login
  jira [-v ...] [-u USER] [-e URI] [-b] [-t FILE] ISSUE

General Options:
  -e --endpoint=URI   URI to use for jira
  -h --help           Show this usage
  -t --template=FILE  Template file to use for output/editing
  -u --user=USER      Username to use for authenticaion (default: %s)
  -v --verbose        Increase output logging
  --version           Show this version

Command Options:
  -a --assignee=USER        Username assigned the issue
  -b --browse               Open your browser to the Jira issue
  -c --component=COMPONENT  Component to Search for
  -d --directory=DIR        Directory to export templates to (default: %s)
  -f --queryfields          Fields that are used in "list" template: (default: summary)
  -i --issuetype=ISSUETYPE  Jira Issue Type (default: Bug)
  -m --comment=COMMENT      Comment message for transition
  -o --override=KEY:VAL     Set custom key/value pairs
  -p --project=PROJECT      Project to Search for
  -q --query=JQL            Jira Query Language expression for the search
  -r --reporter=USER        Reporter to search for
  -w --watcher=USER         Watcher to add to issue (default: %s)
                            or Watcher to search for
`, user, fmt.Sprintf("%s/.jira.d/templates", home), user)

	args, err := docopt.Parse(usage, nil, true, "0.0.3", false, false)
	if err != nil {
		log.Error("Failed to parse options: %s", err)
		os.Exit(1)
	}
	logBackend := logging.NewLogBackend(os.Stderr, "", 0)
	logging.SetBackend(
		logging.NewBackendFormatter(
			logBackend,
			logging.MustStringFormatter(format),
		),
	)
	logging.SetLevel(logging.NOTICE, "")
	if verbose, ok := args["--verbose"]; ok {
		if verbose.(int) > 1 {
			logging.SetLevel(logging.DEBUG, "")
		} else if verbose.(int) > 0 {
			logging.SetLevel(logging.INFO, "")
		}
	}

	log.Info("Args: %v", args)

	opts := make(map[string]string)
	loadConfigs(opts)

	// strip the "--" off the command line options
	// and populate the opts that we pass to the cli ctor
	for key, val := range args {
		if val != nil && strings.HasPrefix(key, "--") {
			opt := key[2:]
			if opt == "override" {
				for _, v := range val.([]string) {
					if strings.Contains(v, "=") {
						kv := strings.SplitN(v, "=", 2)
						opts[kv[0]] = kv[1]
					} else {
						log.Error("Malformed override, expected KEY=VALUE, got %s", v)
						os.Exit(1)
					}
				}
			} else {
				switch v := val.(type) {
				case string:
					opts[opt] = v
				case int:
					opts[opt] = fmt.Sprintf("%d", v)
				case bool:
					opts[opt] = fmt.Sprintf("%t", v)
				}
			}
		}
	}

	// cant use proper [default:x] syntax in docopt
	// because only want to default if the option is not
	// already specified in some .jira.d/config.yml file
	if _, ok := opts["user"]; !ok {
		opts["user"] = user
	}
	if _, ok := opts["queryfields"]; !ok {
		opts["queryfields"] = "summary"
	}
	if _, ok := opts["directory"]; !ok {
		opts["directory"] = fmt.Sprintf("%s/.jira.d/templates", home)
	}

	if _, ok := opts["endpoint"]; !ok {
		log.Error("endpoint option required.  Either use --endpoint or set a enpoint option in your ~/.jira.d/config.yml file")
		os.Exit(1)
	}

	c := cli.New(opts)

	log.Debug("opts: %s", opts)

	validCommand := func(cmd string) bool {
		if val, ok := args[cmd]; ok && val.(bool) {
			return true
		}
		return false
	}

	validOpt := func(opt string, dflt interface{}) interface{} {
		if val, ok := opts[opt]; ok {
			return val
		}
		if dflt == nil {
			log.Error("Missing required option --%s or \"%s\" property override in the config file", opt, opt)
			os.Exit(1)
		}
		return dflt
	}

	setEditing := func(dflt bool) {
		if dflt {
			if val, ok := opts["noedit"]; ok && val == "true" {
				opts["edit"] = "false"
			} else {
				opts["edit"] = "true"
			}
		} else {
			if val, ok := opts["edit"]; ok && val == "true" {
				opts["edit"] = "true"
			} else {
				opts["edit"] = "false"
			}
		}
	}

	if validCommand("login") {
		err = c.CmdLogin()
	} else if validCommand("fields") {
		err = c.CmdFields()
	} else if validCommand("ls") || validCommand("list") {
		err = c.CmdList()
	} else if validCommand("edit") {
		setEditing(true)
		err = c.CmdEdit(args["ISSUE"].(string))
	} else if validCommand("editmeta") {
		err = c.CmdEditMeta(args["ISSUE"].(string))
	} else if validCommand("transmeta") {
		err = c.CmdTransitionMeta(args["ISSUE"].(string))
	} else if validCommand("issuelinktypes") {
		err = c.CmdIssueLinkTypes()
	} else if validCommand("issuetypes") {
		err = c.CmdIssueTypes(validOpt("project", nil).(string))
	} else if validCommand("createmeta") {
		err = c.CmdCreateMeta(
			validOpt("project", nil).(string),
			validOpt("issuetype", "Bug").(string),
		)
	} else if validCommand("create") {
		setEditing(true)
		err = c.CmdCreate(
			validOpt("project", nil).(string),
			validOpt("issuetype", "Bug").(string),
		)
	} else if validCommand("transitions") {
		err = c.CmdTransitions(args["ISSUE"].(string))
	} else if validCommand("blocks") {
		err = c.CmdBlocks(
			args["BLOCKER"].(string),
			args["ISSUE"].(string),
		)
	} else if validCommand("dups") {
		err = c.CmdDups(
			args["DUPLICATE"].(string),
			args["ISSUE"].(string),
		)
	} else if validCommand("watch") {
		err = c.CmdWatch(
			args["ISSUE"].(string),
			validOpt("watcher", user).(string),
		)
	} else if validCommand("trans") || validCommand("transition") {
		setEditing(true)
		err = c.CmdTransition(
			args["ISSUE"].(string),
			args["TRANSITION"].(string),
		)
	} else if validCommand("close") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "close")
	} else if validCommand("ack") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "acknowledge")
	} else if validCommand("reopen") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "reopen")
	} else if validCommand("resolve") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "resolve")
	} else if validCommand("start") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "start")
	} else if validCommand("stop") {
		setEditing(false)
		err = c.CmdTransition(args["ISSUE"].(string), "stop")
	} else if validCommand("comment") {
		setEditing(true)
		err = c.CmdComment(args["ISSUE"].(string))
	} else if validCommand("take") {
		err = c.CmdAssign(args["ISSUE"].(string), user)
	} else if validCommand("browse") || validCommand("b") {
		opts["browse"] = "true"
		err = c.Browse(args["ISSUE"].(string))
	} else if validCommand("export-templates") {
		err = c.CmdExportTemplates()
	} else if validCommand("assign") || validCommand("give") {
		err = c.CmdAssign(
			args["ISSUE"].(string),
			args["ASSIGNEE"].(string),
		)
	} else if val, ok := args["ISSUE"]; ok {
		err = c.CmdView(val.(string))
	}

	if err != nil {
		os.Exit(1)
	}
	os.Exit(0)
}

func parseYaml(file string, opts map[string]string) {
	if fh, err := ioutil.ReadFile(file); err == nil {
		log.Debug("Found Config file: %s", file)
		yaml.Unmarshal(fh, &opts)
	}
}

func loadConfigs(opts map[string]string) {
	paths := cli.FindParentPaths(".jira.d/config.yml")
	// prepend
	paths = append([]string{"/etc/jira-cli.yml"}, paths...)

	for _, file := range paths {
		parseYaml(file, opts)
	}
}
<|endoftext|>


<commit_before>package job

import (
	"sort"
)

type JobOffer struct {
	Job Job

	// MachineIDs represents a set of machines for which this offer is valid.
	// If nil or len == 0 then all machines. Must be sorted.
	// MachineIDs started life as MachineBootIds in the datastore.
	// It cannot be changed without a migration
	MachineIDs []string `json:"MachineBootIds"`
}

func NewOfferFromJob(j Job, machineIDs []string) *JobOffer {
	return &JobOffer{
		Job:            j,
		MachineIDs: machineIDs,
	}
}

// OfferedTo returns true if job is being offered to specified machine.
func (jo *JobOffer) OfferedTo(machineID string) bool {
	// for backward compatibility: if not populated, assume all machines are considered
	if len(jo.MachineIDs) == 0 {
		return true
	}

	k := sort.SearchStrings(jo.MachineIDs, machineID)
	return k < len(jo.MachineIDs) && jo.MachineIDs[k] == machineID
}

type JobBid struct {
	JobName string

	// MachineID started life as MachineBootId in the datastore.
	// It cannot be changed without a migration.
	MachineID string `json:"MachineBootId"`
}

func NewBid(jobName string, machID string) *JobBid {
	return &JobBid{jobName, machID}
}
<commit_msg>refactor(offer): Rename MachineBootIds to MachineIDs<commit_after>package job

import (
	"sort"
)

type JobOffer struct {
	Job Job

	// MachineIDs represents a set of machines for which this offer is valid.
	// If nil or len == 0 then all machines. Must be sorted.
	MachineIDs []string
}

func NewOfferFromJob(j Job, machineIDs []string) *JobOffer {
	return &JobOffer{
		Job:            j,
		MachineIDs: machineIDs,
	}
}

// OfferedTo returns true if job is being offered to specified machine.
func (jo *JobOffer) OfferedTo(machineID string) bool {
	// for backward compatibility: if not populated, assume all machines are considered
	if len(jo.MachineIDs) == 0 {
		return true
	}

	k := sort.SearchStrings(jo.MachineIDs, machineID)
	return k < len(jo.MachineIDs) && jo.MachineIDs[k] == machineID
}

type JobBid struct {
	JobName string

	// MachineID started life as MachineBootId in the datastore.
	// It cannot be changed without a migration.
	MachineID string `json:"MachineBootId"`
}

func NewBid(jobName string, machID string) *JobBid {
	return &JobBid{jobName, machID}
}
<|endoftext|>


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

package starttls

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"
)

func dumpAuthTLSFromFTP(dialer Dialer, address string, config *tls.Config) (*tls.ConnectionState, error) {
	c, err := dialer.Dial("tcp", address)
	if err != nil {
		return nil, err
	}

	conn := c.(*net.TCPConn)
	status, err := readFTP(conn)
	if err != nil {
		return nil, err
	}
	if status != "220" {
		return nil, fmt.Errorf("FTP server responded with status %s, was expecting 220", status)
	}

	fmt.Fprintf(conn, "AUTH TLS\r\n")
	status, err = readFTP(conn)
	if err != nil {
		return nil, err
	}
	if status != "234" {
		return nil, fmt.Errorf("FTP server responded with status %s, was expecting 234", status)
	}

	tlsConn := tls.Client(conn, config)
	err = tlsConn.Handshake()
	if err != nil {
		return nil, err
	}

	state := tlsConn.ConnectionState()
	return &state, nil
}

func readFTP(conn *net.TCPConn) (string, error) {
	reader := bufio.NewReader(conn)
	response, err := reader.ReadString('\n')
	if err != nil {
		return "", err
	}
	status := strings.Split(response, " ")
	if len(status) == 0 {
		return "", errors.New("garbled response from FTP server after AUTH TLS command")
	}
	return status[0], nil
}
<commit_msg>More robust parsing of FTP protocol<commit_after>/*-
 * Copyright 2017 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package starttls

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net"
	"strconv"
)

func dumpAuthTLSFromFTP(dialer Dialer, address string, config *tls.Config) (*tls.ConnectionState, error) {
	c, err := dialer.Dial("tcp", address)
	if err != nil {
		return nil, err
	}

	conn := c.(*net.TCPConn)
	status, err := readFTP(conn)
	if err != nil {
		return nil, err
	}
	if status != 220 {
		return nil, fmt.Errorf("FTP server responded with status %d, was expecting 220", status)
	}

	fmt.Fprintf(conn, "AUTH TLS\r\n")
	status, err = readFTP(conn)
	if err != nil {
		return nil, err
	}
	if status != 234 {
		return nil, fmt.Errorf("FTP server responded with status %d, was expecting 234", status)
	}

	tlsConn := tls.Client(conn, config)
	err = tlsConn.Handshake()
	if err != nil {
		return nil, err
	}

	state := tlsConn.ConnectionState()
	return &state, nil
}

func readFTP(conn *net.TCPConn) (int, error) {
	reader := bufio.NewReader(conn)
	response, err := reader.ReadString('\n')
	if err != nil {
		return 0, err
	}
	if len(response) <= 3 {
		return 0, fmt.Errorf("Error parsing ftp protocol: Status code too short: '%s'", response)
	}
	return strconv.Atoi(response[:3])
}
<|endoftext|>


<commit_before>package main

import (
	"ApacheLog2DB/core"
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"os"
)

func usage() {
	fmt.Println("USAGE: log2db [OPTION]... SOURCE DEST")
	flag.PrintDefaults()
}

func main() {
	flag.Usage = func() { usage() }
	export := flag.Bool("e", false, "Export from SOURCE to DEST")
	flag.Parse()

	args := flag.Args()

	if len(args) != 2 {
		usage()
		os.Exit(1)
	}

	var db *sql.DB
	var err error
	reader := os.Stdin
	writer := os.Stdout

	if !*export {
		if !(args[0] == "-" || args[0] == "--") {
			reader, err = os.Open(args[0])
			defer reader.Close()
		}
	} else {
		if args[0] == "-" || args[0] == "--" {
			fmt.Fprintf(os.Stderr, "Input file must be a db string")
			os.Exit(1)
		}
		db, err = sql.Open("sqlite3", args[0])
		defer db.Close()
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	if !*export {
		if args[1] == "-" || args[1] == "--" {
			fmt.Fprintf(os.Stderr, "Output file must be a db string")
			os.Exit(1)
		}
		db, err = sql.Open("sqlite3", args[1])
		defer db.Close()
	} else {
		if !(args[1] == "-" || args[1] == "--") {
			writer, err = os.Create(args[1])
			defer writer.Close()
		}
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	if *export {
		core.ExportLog(db, writer)
	} else {
		core.ImportLog(reader, db)
	}
	os.Exit(0)
}
<commit_msg>REMOVED: import no longer needed<commit_after>package main

import (
	"ApacheLog2DB/core"
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"os"
)

func usage() {
	fmt.Println("USAGE: log2db [OPTION]... SOURCE DEST")
	flag.PrintDefaults()
}

func main() {
	flag.Usage = func() { usage() }
	export := flag.Bool("e", false, "Export from SOURCE to DEST")
	flag.Parse()

	args := flag.Args()

	if len(args) != 2 {
		usage()
		os.Exit(1)
	}

	var db *sql.DB
	var err error
	reader := os.Stdin
	writer := os.Stdout

	if !*export {
		if !(args[0] == "-" || args[0] == "--") {
			reader, err = os.Open(args[0])
			defer reader.Close()
		}
	} else {
		if args[0] == "-" || args[0] == "--" {
			fmt.Fprintf(os.Stderr, "Input file must be a db string")
			os.Exit(1)
		}
		db, err = sql.Open("sqlite3", args[0])
		defer db.Close()
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	if !*export {
		if args[1] == "-" || args[1] == "--" {
			fmt.Fprintf(os.Stderr, "Output file must be a db string")
			os.Exit(1)
		}
		db, err = sql.Open("sqlite3", args[1],)
		defer db.Close()
	} else {
		if !(args[1] == "-" || args[1] == "--") {
			writer, err = os.Create(args[1])
			defer writer.Close()
		}
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	if *export {
		core.ExportLog(db, writer)
	} else {
		core.ImportLog(reader, db)
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>package middleware

import (
	"bufio"
	"compress/flate"
	"compress/gzip"
	"errors"
	"io"
	"net"
	"net/http"
	"regexp"
	"sort"
	"strings"
)

var encoders = map[string]EncoderFunc{}

var acceptEncodingAlgorithmsRe = regexp.MustCompile(`([a-z]{2,}|\*)`)

func init() {
	// TODO:
	// lzma: Opera.
	// sdch: Chrome, Android. Gzip output + dictionary header.
	// br:   Brotli.

	// TODO: Exception for old MSIE browsers that can't handle non-HTML?
	// https://zoompf.com/blog/2012/02/lose-the-wait-http-compression
	SetEncoder("gzip", encoderGzip)

	// HTTP 1.1 "deflate" (RFC 2616) stands for DEFLATE data (RFC 1951)
	// wrapped with zlib (RFC 1950). The zlib wrapper uses Adler-32
	// checksum compared to CRC-32 used in "gzip" and thus is faster.
	//
	// But.. some old browsers (MSIE, Safari 5.1) incorrectly expect
	// raw DEFLATE data only, without the mentioned zlib wrapper.
	// Because of this major confusion, most modern browsers try it
	// both ways, first looking for zlib headers.
	// Quote by Mark Adler: http://stackoverflow.com/a/9186091/385548
	//
	// The list of browsers having problems is quite big, see:
	// http://zoompf.com/blog/2012/02/lose-the-wait-http-compression
	// https://web.archive.org/web/20120321182910/http://www.vervestudios.co/projects/compression-tests/results
	//
	// That's why we prefer gzip over deflate. It's just more reliable
	// and not significantly slower than gzip.
	SetEncoder("deflate", encoderDeflate)

	// NOTE: Not implemented, intentionally:
	// case "compress": // LZW. Deprecated.
	// case "bzip2":    // Too slow on-the-fly.
	// case "zopfli":   // Too slow on-the-fly.
	// case "xz":       // Too slow on-the-fly.
}

// An EncoderFunc is a function that wraps the provided ResponseWriter with a
// streaming compression algorithm and returns it.
//
// In case of failure, the function should return nil.
type EncoderFunc func(w http.ResponseWriter, level int) io.Writer

// SetEncoder can be used to set the implementation of a compression algorithm.
//
// The encoding should be a standardised identifier. See:
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
//
// For example, add the Brotli algortithm:
//
//  import brotli_enc "gopkg.in/kothar/brotli-go.v0/enc"
//
//  middleware.SetEncoder("br", func(w http.ResponseWriter, level int) io.Writer {
//    params := brotli_enc.NewBrotliParams()
//    params.SetQuality(level)
//    return brotli_enc.NewBrotliWriter(params, w)
//  })
func SetEncoder(encoding string, fn EncoderFunc) {
	if encoding == "" {
		panic("the encoding can not be empty")
	}
	if fn == nil {
		panic("attempted to set a nil encoder function")
	}
	encoders[encoding] = fn
}

var defaultContentTypes = map[string]struct{}{
	"text/html":                {},
	"text/css":                 {},
	"text/plain":               {},
	"text/javascript":          {},
	"application/javascript":   {},
	"application/x-javascript": {},
	"application/json":         {},
	"application/atom+xml":     {},
	"application/rss+xml":      {},
	"image/svg+xml":            {},
}

// DefaultCompress is a middleware that compresses response
// body of predefined content types to a data format based
// on Accept-Encoding request header. It uses a default
// compression level.
func DefaultCompress(next http.Handler) http.Handler {
	return Compress(flate.DefaultCompression)(next)
}

// Compress is a middleware that compresses response
// body of a given content types to a data format based
// on Accept-Encoding request header. It uses a given
// compression level.
func Compress(level int, types ...string) func(next http.Handler) http.Handler {
	contentTypes := defaultContentTypes
	if len(types) > 0 {
		contentTypes = make(map[string]struct{}, len(types))
		for _, t := range types {
			contentTypes[t] = struct{}{}
		}
	}

	return func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			encoder, encoding := selectEncoder(r.Header)
			mcw := &maybeCompressResponseWriter{
				ResponseWriter: w,
				w:              w,
				contentTypes:   contentTypes,
				encoder:        encoder,
				encoding:       encoding,
				level:          level,
			}
			defer mcw.Close()

			next.ServeHTTP(mcw, r)
		}

		return http.HandlerFunc(fn)
	}
}

func selectEncoder(h http.Header) (EncoderFunc, string) {
	header := h.Get("Accept-Encoding")

	// Parse the names of all accepted algorithms from the header.
	var accepted []string
	for _, m := range acceptEncodingAlgorithmsRe.FindAllStringSubmatch(header, -1) {
		accepted = append(accepted, m[1])
	}

	sort.Sort(byPerformance(accepted))

	// Select the first mutually supported algorithm.
	for _, name := range accepted {
		if fn, ok := encoders[name]; ok {
			return fn, name
		}
	}
	return nil, ""
}

type byPerformance []string

func (l byPerformance) Len() int      { return len(l) }
func (l byPerformance) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l byPerformance) Less(i, j int) bool {
	// Higher number = higher preference. This causes unknown names, which map
	// to 0, to always be less prefered.
	scores := map[string]int{
		"br":      3,
		"gzip":    2,
		"deflate": 1,
	}
	return scores[l[i]] > scores[l[j]]
}

type maybeCompressResponseWriter struct {
	http.ResponseWriter
	w            io.Writer
	encoder      EncoderFunc
	encoding     string
	contentTypes map[string]struct{}
	level        int
	wroteHeader  bool
}

func (w *maybeCompressResponseWriter) WriteHeader(code int) {
	if w.wroteHeader {
		return
	}
	w.wroteHeader = true
	defer w.ResponseWriter.WriteHeader(code)

	// Already compressed data?
	if w.ResponseWriter.Header().Get("Content-Encoding") != "" {
		return
	}

	// Parse the first part of the Content-Type response header.
	contentType := ""
	parts := strings.Split(w.ResponseWriter.Header().Get("Content-Type"), ";")
	if len(parts) > 0 {
		contentType = parts[0]
	}

	// Is the content type compressable?
	if _, ok := w.contentTypes[contentType]; !ok {
		return
	}

	if w.encoder != nil && w.encoding != "" {
		if wr := w.encoder(w.ResponseWriter, w.level); wr != nil {
			w.w = wr
			w.Header().Set("Content-Encoding", w.encoding)
			// The content-length after compression is unknown
			w.Header().Del("Content-Length")
		}
	}
}

func (w *maybeCompressResponseWriter) Write(p []byte) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}

	return w.w.Write(p)
}

func (w *maybeCompressResponseWriter) Flush() {
	if f, ok := w.w.(http.Flusher); ok {
		f.Flush()
	}
}

func (w *maybeCompressResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := w.w.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, errors.New("chi/middleware: http.Hijacker is unavailable on the writer")
}

func (w *maybeCompressResponseWriter) CloseNotify() <-chan bool {
	if cn, ok := w.w.(http.CloseNotifier); ok {
		return cn.CloseNotify()
	}

	// If the underlying writer does not implement http.CloseNotifier, return
	// a channel that never receives a value. The semantics here is that the
	// client never disconnnects before the request is processed by the
	// http.Handler, which is close enough to the default behavior (when
	// CloseNotify() is not even called).
	return make(chan bool, 1)
}

func (w *maybeCompressResponseWriter) Close() error {
	if c, ok := w.w.(io.WriteCloser); ok {
		return c.Close()
	}
	return errors.New("chi/middleware: io.WriteCloser is unavailable on the writer")
}

func encoderGzip(w http.ResponseWriter, level int) io.Writer {
	gw, err := gzip.NewWriterLevel(w, level)
	if err != nil {
		return nil
	}
	return gw
}

func encoderDeflate(w http.ResponseWriter, level int) io.Writer {
	dw, err := flate.NewWriter(w, level)
	if err != nil {
		return nil
	}
	return dw
}
<commit_msg>tiny fix: formatting of previous PR<commit_after>package middleware

import (
	"bufio"
	"compress/flate"
	"compress/gzip"
	"errors"
	"io"
	"net"
	"net/http"
	"regexp"
	"sort"
	"strings"
)

var encoders = map[string]EncoderFunc{}

var acceptEncodingAlgorithmsRe = regexp.MustCompile(`([a-z]{2,}|\*)`)

func init() {
	// TODO:
	// lzma: Opera.
	// sdch: Chrome, Android. Gzip output + dictionary header.
	// br:   Brotli.

	// TODO: Exception for old MSIE browsers that can't handle non-HTML?
	// https://zoompf.com/blog/2012/02/lose-the-wait-http-compression
	SetEncoder("gzip", encoderGzip)

	// HTTP 1.1 "deflate" (RFC 2616) stands for DEFLATE data (RFC 1951)
	// wrapped with zlib (RFC 1950). The zlib wrapper uses Adler-32
	// checksum compared to CRC-32 used in "gzip" and thus is faster.
	//
	// But.. some old browsers (MSIE, Safari 5.1) incorrectly expect
	// raw DEFLATE data only, without the mentioned zlib wrapper.
	// Because of this major confusion, most modern browsers try it
	// both ways, first looking for zlib headers.
	// Quote by Mark Adler: http://stackoverflow.com/a/9186091/385548
	//
	// The list of browsers having problems is quite big, see:
	// http://zoompf.com/blog/2012/02/lose-the-wait-http-compression
	// https://web.archive.org/web/20120321182910/http://www.vervestudios.co/projects/compression-tests/results
	//
	// That's why we prefer gzip over deflate. It's just more reliable
	// and not significantly slower than gzip.
	SetEncoder("deflate", encoderDeflate)

	// NOTE: Not implemented, intentionally:
	// case "compress": // LZW. Deprecated.
	// case "bzip2":    // Too slow on-the-fly.
	// case "zopfli":   // Too slow on-the-fly.
	// case "xz":       // Too slow on-the-fly.
}

// An EncoderFunc is a function that wraps the provided ResponseWriter with a
// streaming compression algorithm and returns it.
//
// In case of failure, the function should return nil.
type EncoderFunc func(w http.ResponseWriter, level int) io.Writer

// SetEncoder can be used to set the implementation of a compression algorithm.
//
// The encoding should be a standardised identifier. See:
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
//
// For example, add the Brotli algortithm:
//
//  import brotli_enc "gopkg.in/kothar/brotli-go.v0/enc"
//
//  middleware.SetEncoder("br", func(w http.ResponseWriter, level int) io.Writer {
//    params := brotli_enc.NewBrotliParams()
//    params.SetQuality(level)
//    return brotli_enc.NewBrotliWriter(params, w)
//  })
func SetEncoder(encoding string, fn EncoderFunc) {
	if encoding == "" {
		panic("the encoding can not be empty")
	}
	if fn == nil {
		panic("attempted to set a nil encoder function")
	}
	encoders[encoding] = fn
}

var defaultContentTypes = map[string]struct{}{
	"text/html":                {},
	"text/css":                 {},
	"text/plain":               {},
	"text/javascript":          {},
	"application/javascript":   {},
	"application/x-javascript": {},
	"application/json":         {},
	"application/atom+xml":     {},
	"application/rss+xml":      {},
	"image/svg+xml":            {},
}

// DefaultCompress is a middleware that compresses response
// body of predefined content types to a data format based
// on Accept-Encoding request header. It uses a default
// compression level.
func DefaultCompress(next http.Handler) http.Handler {
	return Compress(flate.DefaultCompression)(next)
}

// Compress is a middleware that compresses response
// body of a given content types to a data format based
// on Accept-Encoding request header. It uses a given
// compression level.
func Compress(level int, types ...string) func(next http.Handler) http.Handler {
	contentTypes := defaultContentTypes
	if len(types) > 0 {
		contentTypes = make(map[string]struct{}, len(types))
		for _, t := range types {
			contentTypes[t] = struct{}{}
		}
	}

	return func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			encoder, encoding := selectEncoder(r.Header)
			mcw := &maybeCompressResponseWriter{
				ResponseWriter: w,
				w:              w,
				contentTypes:   contentTypes,
				encoder:        encoder,
				encoding:       encoding,
				level:          level,
			}
			defer mcw.Close()

			next.ServeHTTP(mcw, r)
		}

		return http.HandlerFunc(fn)
	}
}

func selectEncoder(h http.Header) (EncoderFunc, string) {
	header := h.Get("Accept-Encoding")

	// Parse the names of all accepted algorithms from the header.
	var accepted []string
	for _, m := range acceptEncodingAlgorithmsRe.FindAllStringSubmatch(header, -1) {
		accepted = append(accepted, m[1])
	}

	sort.Sort(byPerformance(accepted))

	// Select the first mutually supported algorithm.
	for _, name := range accepted {
		if fn, ok := encoders[name]; ok {
			return fn, name
		}
	}
	return nil, ""
}

type byPerformance []string

func (l byPerformance) Len() int      { return len(l) }
func (l byPerformance) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l byPerformance) Less(i, j int) bool {
	// Higher number = higher preference. This causes unknown names, which map
	// to 0, to always be less prefered.
	scores := map[string]int{
		"br":      3,
		"gzip":    2,
		"deflate": 1,
	}
	return scores[l[i]] > scores[l[j]]
}

type maybeCompressResponseWriter struct {
	http.ResponseWriter
	w            io.Writer
	encoder      EncoderFunc
	encoding     string
	contentTypes map[string]struct{}
	level        int
	wroteHeader  bool
}

func (w *maybeCompressResponseWriter) WriteHeader(code int) {
	if w.wroteHeader {
		return
	}
	w.wroteHeader = true
	defer w.ResponseWriter.WriteHeader(code)

	// Already compressed data?
	if w.ResponseWriter.Header().Get("Content-Encoding") != "" {
		return
	}

	// Parse the first part of the Content-Type response header.
	contentType := ""
	parts := strings.Split(w.ResponseWriter.Header().Get("Content-Type"), ";")
	if len(parts) > 0 {
		contentType = parts[0]
	}

	// Is the content type compressable?
	if _, ok := w.contentTypes[contentType]; !ok {
		return
	}

	if w.encoder != nil && w.encoding != "" {
		if wr := w.encoder(w.ResponseWriter, w.level); wr != nil {
			w.w = wr
			w.Header().Set("Content-Encoding", w.encoding)
			// The content-length after compression is unknown
			w.Header().Del("Content-Length")
		}
	}
}

func (w *maybeCompressResponseWriter) Write(p []byte) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}

	return w.w.Write(p)
}

func (w *maybeCompressResponseWriter) Flush() {
	if f, ok := w.w.(http.Flusher); ok {
		f.Flush()
	}
}

func (w *maybeCompressResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := w.w.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, errors.New("chi/middleware: http.Hijacker is unavailable on the writer")
}

func (w *maybeCompressResponseWriter) CloseNotify() <-chan bool {
	if cn, ok := w.w.(http.CloseNotifier); ok {
		return cn.CloseNotify()
	}

	// If the underlying writer does not implement http.CloseNotifier, return
	// a channel that never receives a value. The semantics here is that the
	// client never disconnnects before the request is processed by the
	// http.Handler, which is close enough to the default behavior (when
	// CloseNotify() is not even called).
	return make(chan bool, 1)
}

func (w *maybeCompressResponseWriter) Close() error {
	if c, ok := w.w.(io.WriteCloser); ok {
		return c.Close()
	}
	return errors.New("chi/middleware: io.WriteCloser is unavailable on the writer")
}

func encoderGzip(w http.ResponseWriter, level int) io.Writer {
	gw, err := gzip.NewWriterLevel(w, level)
	if err != nil {
		return nil
	}
	return gw
}

func encoderDeflate(w http.ResponseWriter, level int) io.Writer {
	dw, err := flate.NewWriter(w, level)
	if err != nil {
		return nil
	}
	return dw
}
<|endoftext|>


<commit_before>// This binary is the canonical way to use autoroute. It supports most use cases and should
// be able to provide most needs. Ideally the binary should merely be a wrapper around the core library
// and should be the canonical example of how to use it. There should be nothing which this binary can
// do which isn't possible in the core library.
package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"

	"github.com/AutoRoute/node"
	"github.com/AutoRoute/node/types"
	"github.com/AutoRoute/tuntap"
)

var listen = flag.String("listen", "[::]:34321",
	"The address to listen to incoming connections on")
var connect = flag.String("connect", "",
	"Comma separated list of addresses to connect to")
var autodiscover = flag.Bool("auto", false,
	"Whether we should try and find neighboring routers")
var dev_names = flag.String("devs", "",
	"Comma separated list of interfaces to discover on")
var keyfile = flag.String("keyfile", "",
	"The keyfile we should check for a key and write our current key to")
var btc_host = flag.String("btc_host", "localhost:8333",
	"The bitcoin daemon we should connect to")
var btc_user = flag.String("btc_user", "user",
	"The bitcoin daemon username")
var btc_pass = flag.String("btc_pass", "password",
	"The bitcoin daemon password")
var fake_money = flag.Bool("fake_money", false, "Enables a money system which is purely fake")
var status = flag.String("status", "[::1]:12345", "The port to expose status information on")
var unix = flag.String("unix", "", "The path to accept / receive packets as unix packets from")
var tcptun = flag.String("tcptun", "", "Address to try and tcp tunnel to")

func main() {
	log.Print(os.Args)
	flag.Parse()

	// Capture all signals to the quit channel
	quit := make(chan os.Signal)
	signal.Notify(quit)

	// Figure out and load what key we are using for our identity
	var key node.Key
	var err error
	if len(*keyfile) > 0 {
		key, err = node.LoadKey(*keyfile)
		if err != nil {
			log.Fatalf("Error loading key: %v", err)
		}
	} else {
		log.Print("Generating key")
		key, err = node.NewKey()
		if err != nil {
			log.Fatalf("Error Generating Key: %v", err)
		}
	}
	log.Printf("Key is %x", key)

	// Connect to our money server
	money := node.FakeMoney()
	if !*fake_money {
		log.Print("Connecting to bitcoin daemon")
		rpc, err := node.NewRPCMoney(*btc_host, *btc_user, *btc_pass)
		if err != nil {
			log.Fatalf("Error connection to bitcoin daemon: %v", err)
		}
		money = rpc
	}
	log.Printf("Connected")
	n := node.NewServer(key, money)

	log.Printf("Starting to listen on %s", *listen)
	err = n.Listen(*listen)
	if err != nil {
		log.Printf("Error listening: %v", err)
	}

	if *autodiscover {
		devs := make([]net.Interface, 0)

		if *dev_names == "" {
			devs, err = net.Interfaces()
			if err != nil {
				log.Fatal(err)
			}
		} else {
			for _, dev_name := range strings.Split(*dev_names, ",") {
				dev, err := net.InterfaceByName(dev_name)
				if err != nil {
					log.Fatal(err)
				}
				devs = append(devs, *dev)
			}
		}

		for _, dev := range devs {
			go func(dev net.Interface) {
				err := n.Probe(dev)
				if err != nil {
					log.Fatal(err)
				}
			}(dev)
		}
	}

	for _, ip := range strings.Split(*connect, ",") {
		if len(ip) == 0 {
			continue
		}
		err := n.Connect(ip)
		if err != nil {
			log.Printf("Error connecting to %s: %v", ip, err)
		}
	}

	go func() {
		log.Fatal(http.ListenAndServe(*status, nil))
	}()

	if len(*tcptun) > 0 {
		log.Printf("Establishing tcp tunnel to %v", *tcptun)
		i, err := tuntap.Open("tun%d", tuntap.DevTun)
		if err != nil {
			log.Fatal(err)
		}
		dest := ""
		_, err = fmt.Sscanf(*tcptun, "%x", &dest)
		if err != nil {
			log.Fatal(err)
		}
		t := node.NewTCPTunnel(i, n.Node(), types.NodeAddress(dest), 10000)
		defer t.Close()
		<-quit
		return
	}

	if len(*unix) > 0 {
		log.Printf("Establishing unix interface %s", *unix)
		c, err := node.NewUnixSocket(*unix, n.Node())
		if err != nil {
			log.Fatal(err)
		}
		defer os.Remove(*unix)
		defer c.Close()
		<-quit
		return
	}

	<-quit
}
<commit_msg>Give autoroute the ability to statically assign an IP when you turn on tunneling mode.<commit_after>// This binary is the canonical way to use autoroute. It supports most use cases and should
// be able to provide most needs. Ideally the binary should merely be a wrapper around the core library
// and should be the canonical example of how to use it. There should be nothing which this binary can
// do which isn't possible in the core library.
package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strings"

	"github.com/AutoRoute/node"
	"github.com/AutoRoute/node/types"
	"github.com/AutoRoute/tuntap"
)

var listen = flag.String("listen", "[::]:34321",
	"The address to listen to incoming connections on")
var connect = flag.String("connect", "",
	"Comma separated list of addresses to connect to")
var autodiscover = flag.Bool("auto", false,
	"Whether we should try and find neighboring routers")
var dev_names = flag.String("devs", "",
	"Comma separated list of interfaces to discover on")
var keyfile = flag.String("keyfile", "",
	"The keyfile we should check for a key and write our current key to")
var btc_host = flag.String("btc_host", "localhost:8333",
	"The bitcoin daemon we should connect to")
var btc_user = flag.String("btc_user", "user",
	"The bitcoin daemon username")
var btc_pass = flag.String("btc_pass", "password",
	"The bitcoin daemon password")
var fake_money = flag.Bool("fake_money", false, "Enables a money system which is purely fake")
var status = flag.String("status", "[::1]:12345", "The port to expose status information on")
var unix = flag.String("unix", "", "The path to accept / receive packets as unix packets from")
var tcptun = flag.String("tcptun", "", "Address to try and tcp tunnel to")
var tcpaddress = flag.String("tcp_address", "", "IP address to assign to the tcp tunnel")

func main() {
	log.SetFlags(log.LstdFlags | log.Lmicroseconds | log.Lshortfile)
	log.Print(os.Args)
	flag.Parse()

	// Capture all signals to the quit channel
	quit := make(chan os.Signal)
	signal.Notify(quit)

	// Figure out and load what key we are using for our identity
	var key node.Key
	var err error
	if len(*keyfile) > 0 {
		key, err = node.LoadKey(*keyfile)
		if err != nil {
			log.Fatalf("Error loading key: %v", err)
		}
	} else {
		log.Print("Generating key")
		key, err = node.NewKey()
		if err != nil {
			log.Fatalf("Error Generating Key: %v", err)
		}
	}
	log.Printf("Key is %x", key)

	// Connect to our money server
	money := node.FakeMoney()
	if !*fake_money {
		log.Print("Connecting to bitcoin daemon")
		rpc, err := node.NewRPCMoney(*btc_host, *btc_user, *btc_pass)
		if err != nil {
			log.Fatalf("Error connection to bitcoin daemon: %v", err)
		}
		money = rpc
	}
	log.Printf("Connected")
	n := node.NewServer(key, money)

	log.Printf("Starting to listen on %s", *listen)
	err = n.Listen(*listen)
	if err != nil {
		log.Printf("Error listening: %v", err)
	}

	if *autodiscover {
		devs := make([]net.Interface, 0)

		if *dev_names == "" {
			devs, err = net.Interfaces()
			if err != nil {
				log.Fatal(err)
			}
		} else {
			for _, dev_name := range strings.Split(*dev_names, ",") {
				dev, err := net.InterfaceByName(dev_name)
				if err != nil {
					log.Fatal(err)
				}
				devs = append(devs, *dev)
			}
		}

		for _, dev := range devs {
			go func(dev net.Interface) {
				err := n.Probe(dev)
				if err != nil {
					log.Fatal(err)
				}
			}(dev)
		}
	}

	for _, ip := range strings.Split(*connect, ",") {
		if len(ip) == 0 {
			continue
		}
		err := n.Connect(ip)
		if err != nil {
			log.Printf("Error connecting to %s: %v", ip, err)
		}
	}

	go func() {
		log.Fatal(http.ListenAndServe(*status, nil))
	}()

	if len(*tcptun) > 0 {
		log.Printf("Establishing tcp tunnel to %v", *tcptun)
		i, err := tuntap.Open("tun%d", tuntap.DevTun)
		if err != nil {
			log.Fatal(err)
		}
		dest := ""
		_, err = fmt.Sscanf(*tcptun, "%x", &dest)
		if err != nil {
			log.Fatal(err)
		}
		t := node.NewTCPTunnel(i, n.Node(), types.NodeAddress(dest), 10000)
		defer t.Close()

		if len(*tcpaddress) > 0 {
			ip, _, err := net.ParseCIDR(*tcpaddress)
			if err != nil {
				log.Fatal(err)
			}
			err = exec.Command("ip", "link", "set", "dev", strings.TrimRight(i.Name(), "\x00"), "up").Run()
			if err != nil {
				log.Fatal(err)
				quit := make(chan int)
				<-quit
			}
			err := exec.Command("ip", "addr", "add", *tcpaddress, "dev", strings.TrimRight(i.Name(), "\x00")).Run()
			if err != nil {
				log.Fatal(err)
			}
			err = exec.Command("ip", "route", "add", "0/1", "via", ip.String(), "dev", strings.TrimRight(i.Name(), "\x00")).Run()
			if err != nil {
				log.Fatal(err)
			}
			err = exec.Command("ip", "route", "add", "128/1", "via", ip.String(), "dev", strings.TrimRight(i.Name(), "\x00")).Run()
			if err != nil {
				log.Fatal(err)
			}

		}
		<-quit
		return
	}

	if len(*unix) > 0 {
		log.Printf("Establishing unix interface %s", *unix)
		c, err := node.NewUnixSocket(*unix, n.Node())
		if err != nil {
			log.Fatal(err)
		}
		defer os.Remove(*unix)
		defer c.Close()
		<-quit
		return
	}

	<-quit
}
<|endoftext|>