<commit_before>//go:generate swagger generate spec
// Package main CloudTab.
//
// the purpose of this application is to provide an CMDB application
// that will store information in mongodb backend
//
// Terms Of Service:
//
// there are no TOS at this moment, use at your own risk we take no responsibility
//
//     Schemes: http
//     Host: localhost
//     BasePath: /
//     Version: 0.0.1
//     License: MIT http://opensource.org/licenses/MIT
//     Contact: Julien SENON <julien.senon@gmail.com>
package main

import (
	"api"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"net/http"
	"web"
)

// TO FIX

// Struct JSON
// Server{
// 	CMDBName
// 	Function
// 	SerialNumber
// 	AssetCode
// 	HardwareDefinition {
// 		Model
// 		CPU
// 		RAM
// 	}
// 	Localisation {
// 		Room
// 		Building
// 		Rack
// 	}
// 	Networks {
// 		IpAddr
// 		PatchPanel
// 		ServerPort
// 		Switch
// 		Vlan
// 		MAC
// 	}
// 	Remarks
//  Status
// }

func main() {

	r := mux.NewRouter()

	headersOk := handlers.AllowedHeaders([]string{"X-Requested-With"})
	originsOk := handlers.AllowedOrigins([]string{"*"})
	methodsOk := handlers.AllowedMethods([]string{"GET", "HEAD", "POST", "PUT", "OPTIONS", "PATCH"})

	r.HandleFunc("/index", web.Index)
	r.HandleFunc("/send", web.Send)
	r.HandleFunc("/delete", web.Delete)
	r.HandleFunc("/details/{id}", web.Details)
	r.HandleFunc("/update/{id}", web.Update)
	r.HandleFunc("/api", web.ApiHelp)

	// Login
	r.HandleFunc("/login", web.Login)

	// API Part
	r.HandleFunc("/api/servers", api.GetAllItems).Methods("GET")
	r.HandleFunc("/api/servers", api.PostItem).Methods("POST")
	r.HandleFunc("/api/servers/{id}", api.DeleteItem).Methods("DELETE")
	r.HandleFunc("/api/servers/{id}", api.GetItem).Methods("GET")
	r.HandleFunc("/api/servers/{id}", api.UpdateItem).Methods("PATCH")

	// Static dir
	r.PathPrefix("/").Handler(http.StripPrefix("/", http.FileServer(http.Dir("templates/static/"))))

	http.ListenAndServe(":9010", handlers.CORS(originsOk, headersOk, methodsOk)(r))
}
<commit_msg>remove cors protection for swagger<commit_after>//go:generate swagger generate spec
// Package main CloudTab.
//
// the purpose of this application is to provide an CMDB application
// that will store information in mongodb backend
//
// Terms Of Service:
//
// there are no TOS at this moment, use at your own risk we take no responsibility
//
//     Schemes: http
//     Host: localhost
//     BasePath: /
//     Version: 0.0.1
//     License: MIT http://opensource.org/licenses/MIT
//     Contact: Julien SENON <julien.senon@gmail.com>
package main

import (
	"api"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"net/http"
	"web"
)

// TO FIX

// Struct JSON
// Server{
// 	CMDBName
// 	Function
// 	SerialNumber
// 	AssetCode
// 	HardwareDefinition {
// 		Model
// 		CPU
// 		RAM
// 	}
// 	Localisation {
// 		Room
// 		Building
// 		Rack
// 	}
// 	Networks {
// 		IpAddr
// 		PatchPanel
// 		ServerPort
// 		Switch
// 		Vlan
// 		MAC
// 	}
// 	Remarks
//  Status
// }

func main() {

	r := mux.NewRouter()

	// Remove CORS Header check to allow swagger and application on same host and port
	headersOk := handlers.AllowedHeaders([]string{"X-Requested-With", "Content-Type"})
	// To be changed
	originsOk := handlers.AllowedOrigins([]string{"*"})
	methodsOk := handlers.AllowedMethods([]string{"GET", "HEAD", "POST", "PUT", "OPTIONS", "PATCH"})

	r.HandleFunc("/index", web.Index)
	r.HandleFunc("/send", web.Send)
	r.HandleFunc("/delete", web.Delete)
	r.HandleFunc("/details/{id}", web.Details)
	r.HandleFunc("/update/{id}", web.Update)
	r.HandleFunc("/api", web.ApiHelp)

	// Login
	r.HandleFunc("/login", web.Login)

	// API Part
	r.HandleFunc("/api/servers", api.GetAllItems).Methods("GET")
	r.HandleFunc("/api/servers", api.PostItem).Methods("POST")
	r.HandleFunc("/api/servers/{id}", api.DeleteItem).Methods("DELETE")
	r.HandleFunc("/api/servers/{id}", api.GetItem).Methods("GET")
	r.HandleFunc("/api/servers/{id}", api.UpdateItem).Methods("PATCH")

	// Static dir
	r.PathPrefix("/").Handler(http.StripPrefix("/", http.FileServer(http.Dir("templates/static/"))))

	http.ListenAndServe(":9010", handlers.CORS(originsOk, headersOk, methodsOk)(r))
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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 internal_test

import (
	"crypto/rand"
	"crypto/rsa"
	"fmt"
	"math/big"
	"testing"

	"github.com/google/tink/go/signature/internal"
	"github.com/google/tink/go/subtle/random"
	"github.com/google/tink/go/subtle"
	"github.com/google/tink/go/testutil"
)

func TestRSASSAPKCS1SignVerify(t *testing.T) {
	data := random.GetRandomBytes(20)
	hash := "SHA256"
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa.GenerateKey(rand.Reader, 2048) err = %v, want nil", err)
	}
	signer, err := internal.New_RSA_SSA_PKCS1_Signer(hash, privKey)
	if err != nil {
		t.Fatalf("New_RSA_SSA_PKCS1_Signer() err = %v, want nil", err)
	}
	verifier, err := internal.New_RSA_SSA_PKCS1_Verifier(hash, &privKey.PublicKey)
	if err != nil {
		t.Fatalf("New_RSA_SSA_PKCS1_Verifier() err = %v, want nil", err)
	}
	signature, err := signer.Sign(data)
	if err != nil {
		t.Fatalf("Sign() err = %v, want nil", err)
	}
	if err := verifier.Verify(signature, data); err != nil {
		t.Fatalf("Verify() err = %v, want nil", err)
	}
}

func TestRSASSAPKCS1ModifySignatureFails(t *testing.T) {
	data := random.GetRandomBytes(20)
	hash := "SHA256"
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa.GenerateKey(rand.Reader, 2048) err = %v, want nil", err)
	}
	signer, err := internal.New_RSA_SSA_PKCS1_Signer(hash, privKey)
	if err != nil {
		t.Fatalf("New_RSA_SSA_PKCS1_Signer() err = %v, want nil", err)
	}
	signature, err := signer.Sign(data)
	if err != nil {
		t.Fatalf("Sign() err = %v, want nil", err)
	}
	verifier, err := internal.New_RSA_SSA_PKCS1_Verifier(hash, &privKey.PublicKey)
	if err != nil {
		t.Fatalf("New_RSA_SSA_PKCS1_Verifier() err = %v, want nil", err)
	}
	appendSig := append(signature, 0x01)
	if err := verifier.Verify(appendSig, data); err == nil {
		t.Fatalf("Verify() err = nil, want error")
	}
	truncSig := signature[:len(signature)-2]
	if err := verifier.Verify(truncSig, data); err == nil {
		t.Fatalf("Verify() err = nil, want error")
	}
	signature[0] <<= 1
	if err := verifier.Verify(truncSig, data); err == nil {
		t.Fatalf("Verify() err = nil, want error")
	}
}

func TestNewRSASSAPKCS1SignerVerifierInvalidInput(t *testing.T) {
	validPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa.GenerateKey(rand.Reader, 2048) err = %v, want nil", err)
	}
	rsaShortModulusKey, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		t.Fatalf("decoding rsa short modulus: %v", err)
	}
	testCases := []struct {
		name    string
		hash    string
		privKey *rsa.PrivateKey
	}{
		{
			name:    "weak signature hash algorithm",
			hash:    "SHA1",
			privKey: validPrivKey,
		},
		{
			name: "invalid public key exponent",
			hash: "SHA256",
			privKey: &rsa.PrivateKey{
				D:           validPrivKey.D,
				Primes:      validPrivKey.Primes,
				Precomputed: validPrivKey.Precomputed,
				PublicKey: rsa.PublicKey{
					N: validPrivKey.PublicKey.N,
					E: 3,
				},
			},
		},
		{
			name:    "small modulus size",
			hash:    "SHA256",
			privKey: rsaShortModulusKey,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := internal.New_RSA_SSA_PKCS1_Signer(tc.hash, tc.privKey); err == nil {
				t.Errorf("New_RSA_SSA_PKCS1_Signer() err = nil, want error")
			}
			if _, err := internal.New_RSA_SSA_PKCS1_Verifier(tc.hash, &tc.privKey.PublicKey); err == nil {
				t.Errorf("New_RSA_SSA_PKCS1_Verifier() err = nil, want error")
			}
		})
	}
}

type rsaSSAPKCS1Suite struct {
	testutil.WycheproofSuite
	TestGroups []*rsaSSAPKCS1Group `json:"testGroups"`
}

type rsaSSAPKCS1Group struct {
	testutil.WycheproofGroup
	SHA   string             `json:"sha"`
	E     testutil.HexBytes  `json:"e"`
	N     testutil.HexBytes  `json:"n"`
	Type  string             `json:"type"`
	Tests []*rsaSSAPKCS1Case `json:"tests"`
}

type rsaSSAPKCS1Case struct {
	testutil.WycheproofCase
	Message   testutil.HexBytes `json:"msg"`
	Signature testutil.HexBytes `json:"sig"`
}

func TestRSASSAPKCS1WycheproofCases(t *testing.T) {
	testsRan := 0
	for _, v := range []string{
		"rsa_signature_2048_sha256_test.json",
		"rsa_signature_3072_sha512_test.json",
		"rsa_signature_4096_sha512_test.json",
	} {
		suite := &rsaSSAPKCS1Suite{}
		if err := testutil.PopulateSuite(suite, v); err != nil {
			t.Fatalf("testutil.PopulateSuite() err = %v, want nil", err)
		}
		for _, group := range suite.TestGroups {
			hash := subtle.ConvertHashName(group.SHA)
			if hash == "" {
				t.Fatalf("invalid hash name")
			}
			publicKey := &rsa.PublicKey{
				E: int(new(big.Int).SetBytes(group.E).Uint64()),
				N: new(big.Int).SetBytes(group.N),
			}
			if publicKey.E != 65537 {
				// golang "crypto/rsa" only supports 65537 as an exponent.
				if _, err := internal.New_RSA_SSA_PKCS1_Verifier(hash, publicKey); err == nil {
					t.Errorf("NewRSASSAPKCS1Verifier() err = nil, want error")
				}
				continue
			}
			verifier, err := internal.New_RSA_SSA_PKCS1_Verifier(hash, publicKey)
			if err != nil {
				t.Fatalf("NewRSASSAPKCS1Verifier() err = %v, want nil", err)
			}
			for _, test := range group.Tests {
				caseName := fmt.Sprintf("%s: %s-%s:Case-%d", v, group.Type, group.SHA, test.CaseID)
				t.Run(caseName, func(t *testing.T) {
					testsRan++
					err := verifier.Verify(test.Signature, test.Message)
					switch test.Result {
					case "valid":
						if err != nil {
							t.Errorf("Verify() err = %v, want nil", err)
						}
					case "invalid":
						if err == nil {
							t.Errorf("Verify() err = nil, want error")
						}
					case "acceptable":
						// TODO(b/230489047): Inspect flags to appropriately handle acceptable test cases.
					default:
						t.Errorf("unsupported test result: %q", test.Result)
					}
				})
			}
		}
	}
	if testsRan != 716 {
		t.Errorf("testsRan = %d, want = %d", testsRan, 716)
	}
}
<commit_msg>Skip Wycheproof tests when external files aren't available.<commit_after>// Copyright 2022 Google LLC
//
// 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 internal_test

import (
	"crypto/rand"
	"crypto/rsa"
	"fmt"
	"math/big"
	"testing"

	"github.com/google/tink/go/signature/internal"
	"github.com/google/tink/go/subtle/random"
	"github.com/google/tink/go/subtle"
	"github.com/google/tink/go/testutil"
)

func TestRSASSAPKCS1SignVerify(t *testing.T) {
	data := random.GetRandomBytes(20)
	hash := "SHA256"
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa.GenerateKey(rand.Reader, 2048) err = %v, want nil", err)
	}
	signer, err := internal.New_RSA_SSA_PKCS1_Signer(hash, privKey)
	if err != nil {
		t.Fatalf("New_RSA_SSA_PKCS1_Signer() err = %v, want nil", err)
	}
	verifier, err := internal.New_RSA_SSA_PKCS1_Verifier(hash, &privKey.PublicKey)
	if err != nil {
		t.Fatalf("New_RSA_SSA_PKCS1_Verifier() err = %v, want nil", err)
	}
	signature, err := signer.Sign(data)
	if err != nil {
		t.Fatalf("Sign() err = %v, want nil", err)
	}
	if err := verifier.Verify(signature, data); err != nil {
		t.Fatalf("Verify() err = %v, want nil", err)
	}
}

func TestRSASSAPKCS1ModifySignatureFails(t *testing.T) {
	data := random.GetRandomBytes(20)
	hash := "SHA256"
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa.GenerateKey(rand.Reader, 2048) err = %v, want nil", err)
	}
	signer, err := internal.New_RSA_SSA_PKCS1_Signer(hash, privKey)
	if err != nil {
		t.Fatalf("New_RSA_SSA_PKCS1_Signer() err = %v, want nil", err)
	}
	signature, err := signer.Sign(data)
	if err != nil {
		t.Fatalf("Sign() err = %v, want nil", err)
	}
	verifier, err := internal.New_RSA_SSA_PKCS1_Verifier(hash, &privKey.PublicKey)
	if err != nil {
		t.Fatalf("New_RSA_SSA_PKCS1_Verifier() err = %v, want nil", err)
	}
	appendSig := append(signature, 0x01)
	if err := verifier.Verify(appendSig, data); err == nil {
		t.Fatalf("Verify() err = nil, want error")
	}
	truncSig := signature[:len(signature)-2]
	if err := verifier.Verify(truncSig, data); err == nil {
		t.Fatalf("Verify() err = nil, want error")
	}
	signature[0] <<= 1
	if err := verifier.Verify(truncSig, data); err == nil {
		t.Fatalf("Verify() err = nil, want error")
	}
}

func TestNewRSASSAPKCS1SignerVerifierInvalidInput(t *testing.T) {
	validPrivKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa.GenerateKey(rand.Reader, 2048) err = %v, want nil", err)
	}
	rsaShortModulusKey, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		t.Fatalf("decoding rsa short modulus: %v", err)
	}
	testCases := []struct {
		name    string
		hash    string
		privKey *rsa.PrivateKey
	}{
		{
			name:    "weak signature hash algorithm",
			hash:    "SHA1",
			privKey: validPrivKey,
		},
		{
			name: "invalid public key exponent",
			hash: "SHA256",
			privKey: &rsa.PrivateKey{
				D:           validPrivKey.D,
				Primes:      validPrivKey.Primes,
				Precomputed: validPrivKey.Precomputed,
				PublicKey: rsa.PublicKey{
					N: validPrivKey.PublicKey.N,
					E: 3,
				},
			},
		},
		{
			name:    "small modulus size",
			hash:    "SHA256",
			privKey: rsaShortModulusKey,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := internal.New_RSA_SSA_PKCS1_Signer(tc.hash, tc.privKey); err == nil {
				t.Errorf("New_RSA_SSA_PKCS1_Signer() err = nil, want error")
			}
			if _, err := internal.New_RSA_SSA_PKCS1_Verifier(tc.hash, &tc.privKey.PublicKey); err == nil {
				t.Errorf("New_RSA_SSA_PKCS1_Verifier() err = nil, want error")
			}
		})
	}
}

type rsaSSAPKCS1Suite struct {
	testutil.WycheproofSuite
	TestGroups []*rsaSSAPKCS1Group `json:"testGroups"`
}

type rsaSSAPKCS1Group struct {
	testutil.WycheproofGroup
	SHA   string             `json:"sha"`
	E     testutil.HexBytes  `json:"e"`
	N     testutil.HexBytes  `json:"n"`
	Type  string             `json:"type"`
	Tests []*rsaSSAPKCS1Case `json:"tests"`
}

type rsaSSAPKCS1Case struct {
	testutil.WycheproofCase
	Message   testutil.HexBytes `json:"msg"`
	Signature testutil.HexBytes `json:"sig"`
}

func TestRSASSAPKCS1WycheproofCases(t *testing.T) {
	testutil.SkipTestIfTestSrcDirIsNotSet(t)

	testsRan := 0
	for _, v := range []string{
		"rsa_signature_2048_sha256_test.json",
		"rsa_signature_3072_sha512_test.json",
		"rsa_signature_4096_sha512_test.json",
	} {
		suite := &rsaSSAPKCS1Suite{}
		if err := testutil.PopulateSuite(suite, v); err != nil {
			t.Fatalf("testutil.PopulateSuite() err = %v, want nil", err)
		}
		for _, group := range suite.TestGroups {
			hash := subtle.ConvertHashName(group.SHA)
			if hash == "" {
				t.Fatalf("invalid hash name")
			}
			publicKey := &rsa.PublicKey{
				E: int(new(big.Int).SetBytes(group.E).Uint64()),
				N: new(big.Int).SetBytes(group.N),
			}
			if publicKey.E != 65537 {
				// golang "crypto/rsa" only supports 65537 as an exponent.
				if _, err := internal.New_RSA_SSA_PKCS1_Verifier(hash, publicKey); err == nil {
					t.Errorf("NewRSASSAPKCS1Verifier() err = nil, want error")
				}
				continue
			}
			verifier, err := internal.New_RSA_SSA_PKCS1_Verifier(hash, publicKey)
			if err != nil {
				t.Fatalf("NewRSASSAPKCS1Verifier() err = %v, want nil", err)
			}
			for _, test := range group.Tests {
				caseName := fmt.Sprintf("%s: %s-%s:Case-%d", v, group.Type, group.SHA, test.CaseID)
				t.Run(caseName, func(t *testing.T) {
					testsRan++
					err := verifier.Verify(test.Signature, test.Message)
					switch test.Result {
					case "valid":
						if err != nil {
							t.Errorf("Verify() err = %v, want nil", err)
						}
					case "invalid":
						if err == nil {
							t.Errorf("Verify() err = nil, want error")
						}
					case "acceptable":
						// TODO(b/230489047): Inspect flags to appropriately handle acceptable test cases.
					default:
						t.Errorf("unsupported test result: %q", test.Result)
					}
				})
			}
		}
	}
	if testsRan != 716 {
		t.Errorf("testsRan = %d, want = %d", testsRan, 716)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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

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

Unless required by applicable law or agreedto 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 txserializer provides the vttablet hot row protection.
// See the TxSerializer struct for details.
package txserializer

import (
	"fmt"
	"net/http"
	"sync"
	"time"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/acl"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/streamlog"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/vterrors"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

var (
	// waits stores how many times a transaction was queued because another
	// transaction was already in flight for the same row (range).
	// The key of the map is the table name of the query.
	waits = stats.NewCountersWithSingleLabel(
		"TxSerializerWaits",
		"Number of times a transaction was queued because another transaction was already in flight for the same row range",
		"table_name")
	// waitsDryRun is similar as "waits": In dry-run mode it records how many
	// transactions would have been queued.
	// The key of the map is the table name of the query.
	waitsDryRun = stats.NewCountersWithSingleLabel(
		"TxSerializerWaitsDryRun",
		"Dry run number of transactions that would've been queued",
		"table_name")

	// queueExceeded counts per table how many transactions were rejected because
	// the max queue size per row (range) was exceeded.
	queueExceeded = stats.NewCountersWithSingleLabel(
		"TxSerializerQueueExceeded",
		"Number of transactions that were rejected because the max queue size per row range was exceeded",
		"table_name")
	// queueExceededDryRun counts in dry-run mode how many transactions would have
	// been rejected due to exceeding the max queue size per row (range).
	queueExceededDryRun = stats.NewCountersWithSingleLabel(
		"TxSerializerQueueExceededDryRun",
		"Dry-run Number of transactions that were rejcted because the max queue size was exceeded",
		"table_name")

	// globalQueueExceeded is the same as queueExceeded but for the global queue.
	globalQueueExceeded = stats.NewCounter(
		"TxSerializerGlobalQueueExceeded",
		"Number of transactions that were rejected on the global queue because of exceeding the max queue size per row range")
	globalQueueExceededDryRun = stats.NewCounter(
		"TxSerializerGlobalQueueExceededDryRun",
		"Dry-run stats for TxSerializerGlobalQueueExceeded")
)

// TxSerializer serializes incoming transactions which target the same row range
// i.e. table name and WHERE clause are identical.
// Additional transactions are queued and woken up in arrival order.
//
// This implementation has some parallels to the sync2.Consolidator class.
// However, there are many substantial differences:
// - Results are not shared between queued transactions.
// - Only one waiting transaction and not all are notified when the current one
//   has finished.
// - Waiting transactions are woken up in FIFO order.
// - Waiting transactions are unblocked if their context is done.
// - Both the local queue (per row range) and global queue (whole process) are
//   limited to avoid that queued transactions can consume the full capacity
//   of vttablet. This is important if the capaciy is finite. For example, the
//   number of RPCs in flight could be limited by the RPC subsystem.
type TxSerializer struct {
	*sync2.ConsolidatorCache

	// Immutable fields.
	dryRun                 bool
	maxQueueSize           int
	maxGlobalQueueSize     int
	concurrentTransactions int

	log                          *logutil.ThrottledLogger
	logDryRun                    *logutil.ThrottledLogger
	logWaitsDryRun               *logutil.ThrottledLogger
	logQueueExceededDryRun       *logutil.ThrottledLogger
	logGlobalQueueExceededDryRun *logutil.ThrottledLogger

	mu         sync.Mutex
	queues     map[string]*queue
	globalSize int
}

// New returns a TxSerializer object.
func New(dryRun bool, maxQueueSize, maxGlobalQueueSize, concurrentTransactions int) *TxSerializer {
	return &TxSerializer{
		ConsolidatorCache:            sync2.NewConsolidatorCache(1000),
		dryRun:                       dryRun,
		maxQueueSize:                 maxQueueSize,
		maxGlobalQueueSize:           maxGlobalQueueSize,
		concurrentTransactions:       concurrentTransactions,
		log:                          logutil.NewThrottledLogger("HotRowProtection", 5*time.Second),
		logDryRun:                    logutil.NewThrottledLogger("HotRowProtection DryRun", 5*time.Second),
		logWaitsDryRun:               logutil.NewThrottledLogger("HotRowProtection Waits DryRun", 5*time.Second),
		logQueueExceededDryRun:       logutil.NewThrottledLogger("HotRowProtection QueueExceeded DryRun", 5*time.Second),
		logGlobalQueueExceededDryRun: logutil.NewThrottledLogger("HotRowProtection GlobalQueueExceeded DryRun", 5*time.Second),
		queues:                       make(map[string]*queue),
	}
}

// DoneFunc is returned by Wait() and must be called by the caller.
type DoneFunc func()

// Wait blocks if another transaction for the same range is already in flight.
// It returns when this transaction has its turn.
// "done" is != nil if err == nil and must be called once the transaction is
// done and the next waiting transaction can be unblocked.
// "waited" is true if Wait() had to wait for other transactions.
// "err" is not nil if a) the context is done or b) a queue limit was reached.
func (t *TxSerializer) Wait(ctx context.Context, key, table string) (done DoneFunc, waited bool, err error) {
	t.mu.Lock()
	defer t.mu.Unlock()

	waited, err = t.lockLocked(ctx, key, table)
	if err != nil {
		if waited {
			// Waiting failed early e.g. due a canceled context and we did NOT get the
			// slot. Call "done" now because we don't return it to the caller.
			t.unlockLocked(key, false /* returnSlot */)
		}
		return nil, waited, err
	}
	return func() { t.unlock(key) }, waited, nil
}

// lockLocked queues this transaction. It will unblock immediately if this
// transaction is the first in the queue or when it acquired a slot.
// The method has the suffix "Locked" to clarify that "t.mu" must be locked.
func (t *TxSerializer) lockLocked(ctx context.Context, key, table string) (bool, error) {
	q, ok := t.queues[key]
	if !ok {
		// First transaction in the queue i.e. we don't wait and return immediately.
		t.queues[key] = newQueueForFirstTransaction(t.concurrentTransactions)
		t.globalSize++
		return false, nil
	}

	if t.globalSize >= t.maxGlobalQueueSize {
		if t.dryRun {
			globalQueueExceededDryRun.Add(1)
			t.logGlobalQueueExceededDryRun.Warningf("Would have rejected BeginExecute RPC because there are too many queued transactions (%d >= %d)", t.globalSize, t.maxGlobalQueueSize)
		} else {
			globalQueueExceeded.Add(1)
			return false, vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED,
				"hot row protection: too many queued transactions (%d >= %d)", t.globalSize, t.maxGlobalQueueSize)
		}
	}

	if q.size >= t.maxQueueSize {
		if t.dryRun {
			queueExceededDryRun.Add(table, 1)
			t.logQueueExceededDryRun.Warningf("Would have rejected BeginExecute RPC because there are too many queued transactions (%d >= %d) for the same row (table + WHERE clause: '%v')", q.size, t.maxQueueSize, key)
		} else {
			queueExceeded.Add(table, 1)
			return false, vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED,
				"hot row protection: too many queued transactions (%d >= %d) for the same row (table + WHERE clause: '%v')", q.size, t.maxQueueSize, key)
		}
	}

	if q.availableSlots == nil {
		// Hot row detected: A second, concurrent transaction is seen for the
		// first time.

		// As an optimization, we deferred the creation of the channel until now.
		q.availableSlots = make(chan struct{}, t.concurrentTransactions)
		q.availableSlots <- struct{}{}

		// Include first transaction in the count at /debug/hotrows. (It was not
		// recorded on purpose because it did not wait.)
		t.Record(key)
	}

	t.globalSize++
	q.size++
	q.count++
	if q.size > q.max {
		q.max = q.size
	}
	// Publish the number of waits at /debug/hotrows.
	t.Record(key)

	if t.dryRun {
		waitsDryRun.Add(table, 1)
		t.logWaitsDryRun.Warningf("Would have queued BeginExecute RPC for row (range): '%v' because another transaction to the same range is already in progress.", key)
		return false, nil
	}

	// Unlock before the wait and relock before returning because our caller
	// Wait() holds the lock and assumes it still has it.
	t.mu.Unlock()
	defer t.mu.Lock()

	// Non-blocking write attempt to get a slot.
	select {
	case q.availableSlots <- struct{}{}:
		// Return waited=false because a slot was immediately available.
		return false, nil
	default:
	}

	// Blocking wait for the next available slot.
	waits.Add(table, 1)
	select {
	case q.availableSlots <- struct{}{}:
		return true, nil
	case <-ctx.Done():
		return true, ctx.Err()
	}
}

func (t *TxSerializer) unlock(key string) {
	t.mu.Lock()
	defer t.mu.Unlock()

	t.unlockLocked(key, true)
}

func (t *TxSerializer) unlockLocked(key string, returnSlot bool) {
	q := t.queues[key]
	q.size--
	t.globalSize--

	if q.size == 0 {
		// This is the last transaction in flight.
		delete(t.queues, key)

		if q.max > 1 {
			if t.dryRun {
				t.logDryRun.Infof("%v simultaneous transactions (%v in total) for the same row range (%v) would have been queued.", q.max, q.count, key)
			} else {
				t.log.Infof("%v simultaneous transactions (%v in total) for the same row range (%v) were queued.", q.max, q.count, key)
			}
		}

		// Return early because the queue "q" for this "key" will not be used any
		// more.
		// We intentionally skip returning the last slot and closing the
		// "availableSlots" channel because it is not required by Go.
		return
	}

	// Give up slot by removing ourselves from the channel.
	// Wakes up the next queued transaction.

	if t.dryRun {
		// Dry-run did not acquire a slot in the first place.
		return
	}

	if !returnSlot {
		// We did not acquire a slot in the first place e.g. due to a canceled context.
		return
	}

	// This should never block.
	<-q.availableSlots
}

// Pending returns the number of queued transactions (including the ones which
// are currently in flight.)
func (t *TxSerializer) Pending(key string) int {
	t.mu.Lock()
	defer t.mu.Unlock()

	q, ok := t.queues[key]
	if !ok {
		return 0
	}
	return q.size
}

// ServeHTTP lists the most recent, cached queries and their count.
func (t *TxSerializer) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	if *streamlog.RedactDebugUIQueries {
		response.Write([]byte(`
	<!DOCTYPE html>

Redacted

/debug/hotrows has been redacted for your protection

		`))
		return
	}

	if err := acl.CheckAccessHTTP(request, acl.DEBUGGING); err != nil {
		acl.SendError(response, err)
		return
	}
	items := t.Items()
	response.Header().Set("Content-Type", "text/plain")
	if items == nil {
		response.Write([]byte("empty\n"))
		return
	}
	response.Write([]byte(fmt.Sprintf("Length: %d\n", len(items))))
	for _, v := range items {
		response.Write([]byte(fmt.Sprintf("%v: %s\n", v.Count, v.Query)))
	}
}

// queue represents the local queue for a particular row (range).
//
// Note that we don't use a dedicated queue structure for all waiting
// transactions. Instead, we leverage that Go routines waiting for a channel
// are woken up in the order they are queued up. The "availableSlots" field is
// said channel which has n free slots (for the number of concurrent
// transactions which can access the tx pool). All queued transactions are
// competing for these slots and try to add themselves to the channel.
type queue struct {
	// NOTE: The following fields are guarded by TxSerializer.mu.
	// size counts how many transactions are currently queued/in flight (includes
	// the transactions which are not waiting.)
	size int
	// count is the same as "size", but never gets decremented.
	count int
	// max is the max of "size", i.e. the maximum number of transactions which
	// were simultaneously queued for the same row range.
	max int

	// availableSlots limits the number of concurrent transactions *per*
	// hot row (range). It holds one element for each allowed pending
	// transaction i.e. consumed tx pool slot. Consequently, if the channel
	// is full, subsequent transactions have to wait until they can place
	// their entry here.
	// NOTE: As an optimization, we defer the creation of the channel until
	// a second transaction for the same hot row is running.
	availableSlots chan struct{}
}

func newQueueForFirstTransaction(concurrentTransactions int) *queue {
	return &queue{
		size:  1,
		count: 1,
		max:   1,
	}
}
<commit_msg>cleanup: word correction in tx_serializer.go<commit_after>/*
Copyright 2017 Google Inc.

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

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

Unless required by applicable law or agreedto 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 txserializer provides the vttablet hot row protection.
// See the TxSerializer struct for details.
package txserializer

import (
	"fmt"
	"net/http"
	"sync"
	"time"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/acl"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/streamlog"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/vterrors"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

var (
	// waits stores how many times a transaction was queued because another
	// transaction was already in flight for the same row (range).
	// The key of the map is the table name of the query.
	waits = stats.NewCountersWithSingleLabel(
		"TxSerializerWaits",
		"Number of times a transaction was queued because another transaction was already in flight for the same row range",
		"table_name")
	// waitsDryRun is similar as "waits": In dry-run mode it records how many
	// transactions would have been queued.
	// The key of the map is the table name of the query.
	waitsDryRun = stats.NewCountersWithSingleLabel(
		"TxSerializerWaitsDryRun",
		"Dry run number of transactions that would've been queued",
		"table_name")

	// queueExceeded counts per table how many transactions were rejected because
	// the max queue size per row (range) was exceeded.
	queueExceeded = stats.NewCountersWithSingleLabel(
		"TxSerializerQueueExceeded",
		"Number of transactions that were rejected because the max queue size per row range was exceeded",
		"table_name")
	// queueExceededDryRun counts in dry-run mode how many transactions would have
	// been rejected due to exceeding the max queue size per row (range).
	queueExceededDryRun = stats.NewCountersWithSingleLabel(
		"TxSerializerQueueExceededDryRun",
		"Dry-run Number of transactions that were rejected because the max queue size was exceeded",
		"table_name")

	// globalQueueExceeded is the same as queueExceeded but for the global queue.
	globalQueueExceeded = stats.NewCounter(
		"TxSerializerGlobalQueueExceeded",
		"Number of transactions that were rejected on the global queue because of exceeding the max queue size per row range")
	globalQueueExceededDryRun = stats.NewCounter(
		"TxSerializerGlobalQueueExceededDryRun",
		"Dry-run stats for TxSerializerGlobalQueueExceeded")
)

// TxSerializer serializes incoming transactions which target the same row range
// i.e. table name and WHERE clause are identical.
// Additional transactions are queued and woken up in arrival order.
//
// This implementation has some parallels to the sync2.Consolidator class.
// However, there are many substantial differences:
// - Results are not shared between queued transactions.
// - Only one waiting transaction and not all are notified when the current one
//   has finished.
// - Waiting transactions are woken up in FIFO order.
// - Waiting transactions are unblocked if their context is done.
// - Both the local queue (per row range) and global queue (whole process) are
//   limited to avoid that queued transactions can consume the full capacity
//   of vttablet. This is important if the capaciy is finite. For example, the
//   number of RPCs in flight could be limited by the RPC subsystem.
type TxSerializer struct {
	*sync2.ConsolidatorCache

	// Immutable fields.
	dryRun                 bool
	maxQueueSize           int
	maxGlobalQueueSize     int
	concurrentTransactions int

	log                          *logutil.ThrottledLogger
	logDryRun                    *logutil.ThrottledLogger
	logWaitsDryRun               *logutil.ThrottledLogger
	logQueueExceededDryRun       *logutil.ThrottledLogger
	logGlobalQueueExceededDryRun *logutil.ThrottledLogger

	mu         sync.Mutex
	queues     map[string]*queue
	globalSize int
}

// New returns a TxSerializer object.
func New(dryRun bool, maxQueueSize, maxGlobalQueueSize, concurrentTransactions int) *TxSerializer {
	return &TxSerializer{
		ConsolidatorCache:            sync2.NewConsolidatorCache(1000),
		dryRun:                       dryRun,
		maxQueueSize:                 maxQueueSize,
		maxGlobalQueueSize:           maxGlobalQueueSize,
		concurrentTransactions:       concurrentTransactions,
		log:                          logutil.NewThrottledLogger("HotRowProtection", 5*time.Second),
		logDryRun:                    logutil.NewThrottledLogger("HotRowProtection DryRun", 5*time.Second),
		logWaitsDryRun:               logutil.NewThrottledLogger("HotRowProtection Waits DryRun", 5*time.Second),
		logQueueExceededDryRun:       logutil.NewThrottledLogger("HotRowProtection QueueExceeded DryRun", 5*time.Second),
		logGlobalQueueExceededDryRun: logutil.NewThrottledLogger("HotRowProtection GlobalQueueExceeded DryRun", 5*time.Second),
		queues:                       make(map[string]*queue),
	}
}

// DoneFunc is returned by Wait() and must be called by the caller.
type DoneFunc func()

// Wait blocks if another transaction for the same range is already in flight.
// It returns when this transaction has its turn.
// "done" is != nil if err == nil and must be called once the transaction is
// done and the next waiting transaction can be unblocked.
// "waited" is true if Wait() had to wait for other transactions.
// "err" is not nil if a) the context is done or b) a queue limit was reached.
func (t *TxSerializer) Wait(ctx context.Context, key, table string) (done DoneFunc, waited bool, err error) {
	t.mu.Lock()
	defer t.mu.Unlock()

	waited, err = t.lockLocked(ctx, key, table)
	if err != nil {
		if waited {
			// Waiting failed early e.g. due a canceled context and we did NOT get the
			// slot. Call "done" now because we don't return it to the caller.
			t.unlockLocked(key, false /* returnSlot */)
		}
		return nil, waited, err
	}
	return func() { t.unlock(key) }, waited, nil
}

// lockLocked queues this transaction. It will unblock immediately if this
// transaction is the first in the queue or when it acquired a slot.
// The method has the suffix "Locked" to clarify that "t.mu" must be locked.
func (t *TxSerializer) lockLocked(ctx context.Context, key, table string) (bool, error) {
	q, ok := t.queues[key]
	if !ok {
		// First transaction in the queue i.e. we don't wait and return immediately.
		t.queues[key] = newQueueForFirstTransaction(t.concurrentTransactions)
		t.globalSize++
		return false, nil
	}

	if t.globalSize >= t.maxGlobalQueueSize {
		if t.dryRun {
			globalQueueExceededDryRun.Add(1)
			t.logGlobalQueueExceededDryRun.Warningf("Would have rejected BeginExecute RPC because there are too many queued transactions (%d >= %d)", t.globalSize, t.maxGlobalQueueSize)
		} else {
			globalQueueExceeded.Add(1)
			return false, vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED,
				"hot row protection: too many queued transactions (%d >= %d)", t.globalSize, t.maxGlobalQueueSize)
		}
	}

	if q.size >= t.maxQueueSize {
		if t.dryRun {
			queueExceededDryRun.Add(table, 1)
			t.logQueueExceededDryRun.Warningf("Would have rejected BeginExecute RPC because there are too many queued transactions (%d >= %d) for the same row (table + WHERE clause: '%v')", q.size, t.maxQueueSize, key)
		} else {
			queueExceeded.Add(table, 1)
			return false, vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED,
				"hot row protection: too many queued transactions (%d >= %d) for the same row (table + WHERE clause: '%v')", q.size, t.maxQueueSize, key)
		}
	}

	if q.availableSlots == nil {
		// Hot row detected: A second, concurrent transaction is seen for the
		// first time.

		// As an optimization, we deferred the creation of the channel until now.
		q.availableSlots = make(chan struct{}, t.concurrentTransactions)
		q.availableSlots <- struct{}{}

		// Include first transaction in the count at /debug/hotrows. (It was not
		// recorded on purpose because it did not wait.)
		t.Record(key)
	}

	t.globalSize++
	q.size++
	q.count++
	if q.size > q.max {
		q.max = q.size
	}
	// Publish the number of waits at /debug/hotrows.
	t.Record(key)

	if t.dryRun {
		waitsDryRun.Add(table, 1)
		t.logWaitsDryRun.Warningf("Would have queued BeginExecute RPC for row (range): '%v' because another transaction to the same range is already in progress.", key)
		return false, nil
	}

	// Unlock before the wait and relock before returning because our caller
	// Wait() holds the lock and assumes it still has it.
	t.mu.Unlock()
	defer t.mu.Lock()

	// Non-blocking write attempt to get a slot.
	select {
	case q.availableSlots <- struct{}{}:
		// Return waited=false because a slot was immediately available.
		return false, nil
	default:
	}

	// Blocking wait for the next available slot.
	waits.Add(table, 1)
	select {
	case q.availableSlots <- struct{}{}:
		return true, nil
	case <-ctx.Done():
		return true, ctx.Err()
	}
}

func (t *TxSerializer) unlock(key string) {
	t.mu.Lock()
	defer t.mu.Unlock()

	t.unlockLocked(key, true)
}

func (t *TxSerializer) unlockLocked(key string, returnSlot bool) {
	q := t.queues[key]
	q.size--
	t.globalSize--

	if q.size == 0 {
		// This is the last transaction in flight.
		delete(t.queues, key)

		if q.max > 1 {
			if t.dryRun {
				t.logDryRun.Infof("%v simultaneous transactions (%v in total) for the same row range (%v) would have been queued.", q.max, q.count, key)
			} else {
				t.log.Infof("%v simultaneous transactions (%v in total) for the same row range (%v) were queued.", q.max, q.count, key)
			}
		}

		// Return early because the queue "q" for this "key" will not be used any
		// more.
		// We intentionally skip returning the last slot and closing the
		// "availableSlots" channel because it is not required by Go.
		return
	}

	// Give up slot by removing ourselves from the channel.
	// Wakes up the next queued transaction.

	if t.dryRun {
		// Dry-run did not acquire a slot in the first place.
		return
	}

	if !returnSlot {
		// We did not acquire a slot in the first place e.g. due to a canceled context.
		return
	}

	// This should never block.
	<-q.availableSlots
}

// Pending returns the number of queued transactions (including the ones which
// are currently in flight.)
func (t *TxSerializer) Pending(key string) int {
	t.mu.Lock()
	defer t.mu.Unlock()

	q, ok := t.queues[key]
	if !ok {
		return 0
	}
	return q.size
}

// ServeHTTP lists the most recent, cached queries and their count.
func (t *TxSerializer) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	if *streamlog.RedactDebugUIQueries {
		response.Write([]byte(`
	<!DOCTYPE html>

Redacted

/debug/hotrows has been redacted for your protection

		`))
		return
	}

	if err := acl.CheckAccessHTTP(request, acl.DEBUGGING); err != nil {
		acl.SendError(response, err)
		return
	}
	items := t.Items()
	response.Header().Set("Content-Type", "text/plain")
	if items == nil {
		response.Write([]byte("empty\n"))
		return
	}
	response.Write([]byte(fmt.Sprintf("Length: %d\n", len(items))))
	for _, v := range items {
		response.Write([]byte(fmt.Sprintf("%v: %s\n", v.Count, v.Query)))
	}
}

// queue represents the local queue for a particular row (range).
//
// Note that we don't use a dedicated queue structure for all waiting
// transactions. Instead, we leverage that Go routines waiting for a channel
// are woken up in the order they are queued up. The "availableSlots" field is
// said channel which has n free slots (for the number of concurrent
// transactions which can access the tx pool). All queued transactions are
// competing for these slots and try to add themselves to the channel.
type queue struct {
	// NOTE: The following fields are guarded by TxSerializer.mu.
	// size counts how many transactions are currently queued/in flight (includes
	// the transactions which are not waiting.)
	size int
	// count is the same as "size", but never gets decremented.
	count int
	// max is the max of "size", i.e. the maximum number of transactions which
	// were simultaneously queued for the same row range.
	max int

	// availableSlots limits the number of concurrent transactions *per*
	// hot row (range). It holds one element for each allowed pending
	// transaction i.e. consumed tx pool slot. Consequently, if the channel
	// is full, subsequent transactions have to wait until they can place
	// their entry here.
	// NOTE: As an optimization, we defer the creation of the channel until
	// a second transaction for the same hot row is running.
	availableSlots chan struct{}
}

func newQueueForFirstTransaction(concurrentTransactions int) *queue {
	return &queue{
		size:  1,
		count: 1,
		max:   1,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	"github.com/lib/pq"
)

func connectToDatabase() {
	var err error
	var rows *sql.Rows
	var schema string
	var schemaVer int64

	driver := "postgres"

	connect := fmt.Sprintf("dbname='%s' user='%s' password='%s' host='%s' port='%s' sslmode='%s' connect_timeout='%s'",
		SomaCfg.Database.Name,
		SomaCfg.Database.User,
		SomaCfg.Database.Pass,
		SomaCfg.Database.Host,
		SomaCfg.Database.Port,
		SomaCfg.Database.TlsMode,
		SomaCfg.Database.Timeout,
	)

	// enable handling of infinity timestamps
	pq.EnableInfinityTs(NegTimeInf, PosTimeInf)

	conn, err = sql.Open(driver, connect)
	if err != nil {
		log.Fatal(err)
	}
	if err = conn.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Print("Connected to database")
	if _, err = conn.Exec(`SET TIME ZONE 'UTC';`); err != nil {
		log.Fatal(err)
	}

	// required schema versions
	required := map[string]int64{
		"inventory": 201605060001,
		"root":      201605160001,
		"auth":      201605190001,
		"soma":      201606160001,
	}

	if rows, err = conn.Query(`
SELECT schema,
       MAX(version) AS version
FROM   public.schema_versions
GROUP  BY schema;`); err != nil {
		log.Fatal("Query db schema versions: ", err)
	}

	for rows.Next() {
		if err = rows.Scan(
			&schema,
			&schemaVer,
		); err != nil {
			log.Fatal("Schema check: ", err)
		}
		if rsv, ok := required[schema]; ok {
			if rsv != schemaVer {
				log.Fatalf("Incompatible schema %s: %d != %d", schema, rsv, schemaVer)
			} else {
				log.Printf("DB Schema %s, version: %d", schema, schemaVer)
				delete(required, schema)
			}
		} else {
			log.Fatal("Unknown schema: ", schema)
		}
	}
	if err = rows.Err(); err != nil {
		log.Fatal("Schema check: ", err)
	}
	if len(required) != 0 {
		for s, _ := range required {
			log.Printf("Missing schema: %s", s)
		}
		log.Fatal("FATAL - database incomplete")
	}
}

func pingDatabase() {
	ticker := time.NewTicker(time.Second).C

	for {
		<-ticker
		err := conn.Ping()
		if err != nil {
			log.Print(err)
		}
	}
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Require schema soma/201606210001<commit_after>package main

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	"github.com/lib/pq"
)

func connectToDatabase() {
	var err error
	var rows *sql.Rows
	var schema string
	var schemaVer int64

	driver := "postgres"

	connect := fmt.Sprintf("dbname='%s' user='%s' password='%s' host='%s' port='%s' sslmode='%s' connect_timeout='%s'",
		SomaCfg.Database.Name,
		SomaCfg.Database.User,
		SomaCfg.Database.Pass,
		SomaCfg.Database.Host,
		SomaCfg.Database.Port,
		SomaCfg.Database.TlsMode,
		SomaCfg.Database.Timeout,
	)

	// enable handling of infinity timestamps
	pq.EnableInfinityTs(NegTimeInf, PosTimeInf)

	conn, err = sql.Open(driver, connect)
	if err != nil {
		log.Fatal(err)
	}
	if err = conn.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Print("Connected to database")
	if _, err = conn.Exec(`SET TIME ZONE 'UTC';`); err != nil {
		log.Fatal(err)
	}

	// required schema versions
	required := map[string]int64{
		"inventory": 201605060001,
		"root":      201605160001,
		"auth":      201605190001,
		"soma":      201606210001,
	}

	if rows, err = conn.Query(`
SELECT schema,
       MAX(version) AS version
FROM   public.schema_versions
GROUP  BY schema;`); err != nil {
		log.Fatal("Query db schema versions: ", err)
	}

	for rows.Next() {
		if err = rows.Scan(
			&schema,
			&schemaVer,
		); err != nil {
			log.Fatal("Schema check: ", err)
		}
		if rsv, ok := required[schema]; ok {
			if rsv != schemaVer {
				log.Fatalf("Incompatible schema %s: %d != %d", schema, rsv, schemaVer)
			} else {
				log.Printf("DB Schema %s, version: %d", schema, schemaVer)
				delete(required, schema)
			}
		} else {
			log.Fatal("Unknown schema: ", schema)
		}
	}
	if err = rows.Err(); err != nil {
		log.Fatal("Schema check: ", err)
	}
	if len(required) != 0 {
		for s, _ := range required {
			log.Printf("Missing schema: %s", s)
		}
		log.Fatal("FATAL - database incomplete")
	}
}

func pingDatabase() {
	ticker := time.NewTicker(time.Second).C

	for {
		<-ticker
		err := conn.Ping()
		if err != nil {
			log.Print(err)
		}
	}
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package db

import (
	"github.com/boltdb/bolt"
	"io/ioutil"
	"os"
	"testing"
)

// Helper functions

func getTempDb() string {
	tmpDirPath := os.TempDir()
	f, err := ioutil.TempFile(tmpDirPath, "carton_dbTest")
	if err != nil {
		return ""
	}
	f.Close()
	return f.Name()
}

// BoltManager tests

func TestNewBoltManager(t *testing.T) {
	tempDb := getTempDb()
	if tempDb == "" {
		t.Skip("Cannot create temp file")
	}

	m, err := NewBoltManager(tempDb)
	if err != nil {
		t.Errorf("Error when calling NewBoltManager: %v", err)
	}

	createdUsersBucket := false
	m.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("users"))
		if b != nil {
			createdUsersBucket = true
		}
		return nil
	})

	if !createdUsersBucket {
		t.Errorf("'users' bucket does not exist")
	}
}

func TestUserRegistration(t *testing.T) {
	tempDb := getTempDb()
	if tempDb == "" {
		t.Skip("Cannot create temp file")
	}

	m, _ := NewBoltManager(tempDb)

	if m.IsUser("fake user") {
		t.Errorf("IsUsers returns true for 'fake user'")
	}
	if m.GetPwdHash("fake user") != nil {
		t.Errorf("GetPwdHash returns non nil value for 'fake user'")
	}

	m.RegisterUser("test user", []byte("test hash"))
	if !m.IsUser("test user") ||
		string(m.GetPwdHash("test user")) != "test hash" {
		t.Errorf("failed to register 'test user'")
	}

}
<commit_msg>Tests cleanup tempDb when finished<commit_after>package db

import (
	"github.com/boltdb/bolt"
	"io/ioutil"
	"os"
	"testing"
)

// Helper functions

func getTempDb() string {
	tmpDirPath := os.TempDir()
	f, err := ioutil.TempFile(tmpDirPath, "carton_dbTest")
	if err != nil {
		return ""
	}
	f.Close()
	return f.Name()
}

// BoltManager tests

func TestNewBoltManager(t *testing.T) {
	tempDb := getTempDb()
	if tempDb == "" {
		t.Skip("Cannot create temp file")
	}

	m, err := NewBoltManager(tempDb)
	if err != nil {
		t.Errorf("Error when calling NewBoltManager: %v", err)
	}

	createdUsersBucket := false
	m.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("users"))
		if b != nil {
			createdUsersBucket = true
		}
		return nil
	})

	if !createdUsersBucket {
		t.Errorf("'users' bucket does not exist")
	}

	m.Close()
	os.Remove(tempDb)
}

func TestUserRegistration(t *testing.T) {
	tempDb := getTempDb()
	if tempDb == "" {
		t.Skip("Cannot create temp file")
	}

	m, _ := NewBoltManager(tempDb)

	if m.IsUser("fake user") {
		t.Errorf("IsUsers returns true for 'fake user'")
	}
	if m.GetPwdHash("fake user") != nil {
		t.Errorf("GetPwdHash returns non nil value for 'fake user'")
	}

	m.RegisterUser("test user", []byte("test hash"))
	if !m.IsUser("test user") ||
		string(m.GetPwdHash("test user")) != "test hash" {
		t.Errorf("failed to register 'test user'")
	}

	m.Close()
	os.Remove(tempDb)
}
<|endoftext|>


<commit_before>package slackbot

import (
	"encoding/json"
	"fmt"
	"html"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/justinian/dice"
)

type Command struct {
	Name string
	Desc string
	Func func(*Instance, *Message, []string) error
}

var COMMANDS []*Command

func init() {
	COMMANDS = []*Command{
		&Command{"apollo", "Tell a random quote from Apollo's source code.", cmd_apollo},
		&Command{"catfact", "Tell a random cat fact.", cmd_catfact},
		&Command{"channels", "Show a list of known channels.", cmd_channels},
		&Command{"duck", "Quack.", cmd_duck},
		&Command{"goon", "Tell a random quote from Goon's source code.", cmd_goon},
		&Command{"help", "Show a list of commands.", cmd_help},
		&Command{"hi", "Start a conversation.", cmd_start_talking},
		&Command{"mute", "Mute my messages to this channel, for a while.", cmd_mute},
		&Command{"pun", "Tell a random pun.", cmd_pun},
		&Command{"roll", "Throw a dice roll.", cmd_roll},
		&Command{"shutup", "Ooor be silent.", cmd_stop_talking},
		&Command{"status", "Show my current status.", cmd_status},
		&Command{"users", "Show a list of known users.", cmd_users},
		&Command{"vote", "start/stop a vote or vote yes/no during a vote.", cmd_vote},
		&Command{"wiki", "Quote a page from our SS13 wiki.", cmd_wiki},
	}
}

func cmd_help(i *Instance, m *Message, args []string) error {
	tmp := "Available commands:\n"
	for _, c := range COMMANDS {
		tmp += fmt.Sprintf("- `%s\t%s`\n", c.Name, c.Desc)
	}
	i.UserMsg(m.User, tmp)
	return nil
}

func cmd_duck(i *Instance, m *Message, args []string) error {
	var quotes = []string{
		"Quack",
		"Wenk",
		"O RLY?",
		"Kraaawk!",
		"Bwaak",
		"Chirp",
	}
	i.ChannelMsg(m.Channel, quotes[rand.Intn(len(quotes))])
	return nil
}

func cmd_vote(i *Instance, m *Message, args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("Missing argument.")
	}

	if args[0] == "start" {
		i.StartVote(m.Channel)
		i.ChannelMsg(m.Channel, fmt.Sprintf("@%s has started a new vote! Tell me if you would like to `vote yes` or `vote no` on it.", i.Users[m.User].Name))
	} else if args[0] == "stop" {
		votes := i.StopVote(m.Channel)
		var result string
		if votes > 0 {
			result = "*YES* has won"
		} else if votes < 0 {
			result = "*NO* has won"
		} else {
			result = "It's a tie! No one won"
		}
		i.ChannelMsg(m.Channel, fmt.Sprintf("@%s has stopped the vote! The result is...\n%s (score: %d)!", i.Users[m.User].Name, result, votes))
	} else if args[0] == "yes" {
		i.Vote(+1, m.Channel)
	} else if args[0] == "no" {
		i.Vote(-1, m.Channel)
	}
	return nil
}

func cmd_mute(i *Instance, m *Message, args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("Missing argument.")
	}

	dur, e := time.ParseDuration(args[0])
	if e != nil {
		return fmt.Errorf("Couldn't parse a duration.")
	}
	if dur.Minutes() < 1 {
		dur = time.Duration(5) * time.Minute
	}
	if dur.Minutes() > 60 {
		dur = time.Duration(60) * time.Minute
	}

	i.ChannelMsg(m.Channel, fmt.Sprintf("I will now shut up for ~%.0f minutes.", dur.Minutes()))
	go func() {
		time.Sleep(5)
		i.Mute(m.Channel, dur)
	}()
	return nil
}

func cmd_status(i *Instance, m *Message, args []string) error {
	tmp := "My current status:\n"
	tmp += fmt.Sprintf("`Running: %v`\n", i.running)
	tmp += fmt.Sprintf("`Version: v%v`\n", VERSION)
	tmp += fmt.Sprintf("`Debug mode: %v`\n", i.Debug)
	tmp += fmt.Sprintf("`Latency: %v`\n", i.Latency)
	tmp += fmt.Sprintf("`My User ID: @%v`\n", i.BotId)
	tmp += fmt.Sprintf("`Tracked channels: %v`\n", len(i.Channels))
	tmp += fmt.Sprintf("`Tracked users: %v`\n", len(i.Users))
	tmp += fmt.Sprintf("`No. of Goon quotes: %v`\n", len(i.goon))
	tmp += fmt.Sprintf("`No. of Apollo quotes: %v`\n", len(i.apollo))
	tmp += fmt.Sprintf("\nI am currently muted on channels:\n")
	if len(i.mutes) < 1 {
		tmp += "- `None`\n"
	} else {
		for k, v := range i.mutes {
			c := i.Channels[k].Name
			t := v.Sub(time.Now())
			tmp += fmt.Sprintf("- `%v (ends in ~%.0f minutes)`\n", c, t.Minutes())
		}
	}
	i.UserMsg(m.User, tmp)
	return nil
}

func cmd_roll(i *Instance, m *Message, args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("Missing argument.")
	}

	ret, res, e := dice.Roll(args[0])
	if e != nil {
		return fmt.Errorf("Bad dice format.")
	}
	i.ChannelMsg(m.Channel, fmt.Sprintf("@%s has rolled: %v, %v\n", i.Users[m.User].Name, ret, res))
	return nil
}

func cmd_wiki(i *Instance, m *Message, args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("Missing argument.")
	}

	u := WIKI_URL + url.QueryEscape(args[0])
	doc, e := goquery.NewDocument(u)
	if e != nil {
		return fmt.Errorf("Couldn't open the wiki for you (%v).", u)
	}

	node := doc.Find("div #mw-content-text > p").First()
	text := strings.TrimSpace(node.Text())
	if node.Length() < 1 || len(text) < 1 {
		return fmt.Errorf("Couldn't quote that page for you (%v).", u)
	}

	i.ChannelMsg(m.Channel, fmt.Sprintf(">>>%s\n`Source: %s`", text, u))
	return nil
}

func cmd_pun(i *Instance, m *Message, args []string) error {
	resp, e := http.Get("http://www.punoftheday.com/cgi-bin/arandompun.pl")
	if e != nil {
		return fmt.Errorf("Sorry, couldn't make up a pun for you.")
	}
	defer resp.Body.Close()
	body, e := ioutil.ReadAll(resp.Body)
	if e != nil {
		return fmt.Errorf("Sorry, couldn't make up a pun for you.")
	}
	// Yep this is a fucking mess.
	s := strings.TrimSpace(string(body))
	s = strings.TrimPrefix(s, "document.write('"")
	index := strings.Index(s, """)
	s = html.UnescapeString(s[:index])
	i.ChannelMsg(m.Channel, fmt.Sprintf(">>>%v\n`© 1996-2011 PunoftheDay.com`", s))
	return nil
}

type CatFacts struct {
	Facts   []string
	Success string
}

var LAST_CATFACT_TIME time.Time

func cmd_catfact(i *Instance, m *Message, args []string) error {
	now := time.Now()
	if now.After(LAST_CATFACT_TIME) != true {
		return fmt.Errorf("Sorry, I can only show a single catfact per day (or else I'll get banned!).")
	}
	LAST_CATFACT_TIME = now.Add(24 * time.Hour)

	resp, e := http.Get("http://catfacts-api.appspot.com/api/facts?number=1")
	if e != nil {
		return fmt.Errorf("Sorry, couldn't find any cat facts.")
	}

	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	var cf CatFacts
	e = decoder.Decode(&cf)
	if e != nil {
		return fmt.Errorf("Sorry, couldn't find any cat facts.")
	}

	if len(cf.Facts) > 0 {
		msg := fmt.Sprintf(">>>%v\n`© http://catfacts-api.appspot.com`", cf.Facts[0])
		i.ChannelMsg(m.Channel, msg)
		return nil
	}
	return fmt.Errorf("Sorry, couldn't find any cat facts.")
}

func cmd_goon(i *Instance, m *Message, args []string) error {
	q := i.random_goon_quote()
	msg := fmt.Sprintf(">>>%v\n`%v`", q.Quote, q.File)
	i.ChannelMsg(m.Channel, msg)
	return nil
}
func cmd_apollo(i *Instance, m *Message, args []string) error {
	q := i.random_apollo_quote()
	msg := fmt.Sprintf(">>>%v\n`%v`", q.Quote, q.File)
	i.ChannelMsg(m.Channel, msg)
	return nil
}

func cmd_channels(i *Instance, m *Message, args []string) error {
	tmp := fmt.Sprintf("Known channels (%d):\n", len(i.Channels))

	for id, c := range i.Channels {
		tmp += fmt.Sprintf("`#%s %s - %s`\n", id, c.Name, c.Topic.Value)
	}
	i.UserMsg(m.User, tmp)
	return nil
}

func cmd_users(i *Instance, m *Message, args []string) error {
	tmp := fmt.Sprintf("Known users (%d):\n", len(i.Users))

	for id, u := range i.Users {
		tmp += fmt.Sprintf("`#%s %s - %s`\n", id, u.Name, u.Presence)
	}
	i.UserMsg(m.User, tmp)
	return nil
}

func cmd_start_talking(i *Instance, m *Message, args []string) error {
	if m.Channel != TALK_CHANNEL {
		return fmt.Errorf("Nuhuh, to shy to talk there.")
	}

	if i.talking {
		i.ChannelMsg(TALK_CHANNEL, "I'm already talking! Stop bothering me!")
	} else {
		i.ChannelMsg(TALK_CHANNEL, "Helloooo..")
	}
	i.talking = true
	return nil
}

func cmd_stop_talking(i *Instance, m *Message, args []string) error {
	if m.Channel != TALK_CHANNEL {
		return nil
	}

	if i.talking {
		i.ChannelMsg(TALK_CHANNEL, "kthx bye!")
	} else {
		i.ChannelMsg(TALK_CHANNEL, "Haven't said a word, Cap'.")
	}
	i.talking = false
	return nil
}
<commit_msg>Add arrest command (heh..).<commit_after>package slackbot

import (
	"encoding/json"
	"fmt"
	"html"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/justinian/dice"
)

type Command struct {
	Name string
	Desc string
	Func func(*Instance, *Message, []string) error
}

var COMMANDS []*Command

func init() {
	COMMANDS = []*Command{
		&Command{"apollo", "Tell a random quote from Apollo's source code.", cmd_apollo},
		&Command{"arrest", "Arrest someone.", cmd_arrest},
		&Command{"catfact", "Tell a random cat fact.", cmd_catfact},
		&Command{"channels", "Show a list of known channels.", cmd_channels},
		&Command{"duck", "Quack.", cmd_duck},
		&Command{"goon", "Tell a random quote from Goon's source code.", cmd_goon},
		&Command{"help", "Show a list of commands.", cmd_help},
		&Command{"hi", "Start a conversation.", cmd_start_talking},
		&Command{"mute", "Mute my messages to this channel, for a while.", cmd_mute},
		&Command{"pun", "Tell a random pun.", cmd_pun},
		&Command{"roll", "Throw a dice roll.", cmd_roll},
		&Command{"shutup", "Ooor be silent.", cmd_stop_talking},
		&Command{"status", "Show my current status.", cmd_status},
		&Command{"users", "Show a list of known users.", cmd_users},
		&Command{"vote", "start/stop a vote or vote yes/no during a vote.", cmd_vote},
		&Command{"wiki", "Quote a page from our SS13 wiki.", cmd_wiki},
	}
}

func cmd_help(i *Instance, m *Message, args []string) error {
	tmp := "Available commands:\n"
	for _, c := range COMMANDS {
		tmp += fmt.Sprintf("- `%s\t%s`\n", c.Name, c.Desc)
	}
	i.UserMsg(m.User, tmp)
	return nil
}

func cmd_duck(i *Instance, m *Message, args []string) error {
	var quotes = []string{
		"Quack",
		"Wenk",
		"O RLY?",
		"Kraaawk!",
		"Bwaak",
		"Chirp",
	}
	i.ChannelMsg(m.Channel, quotes[rand.Intn(len(quotes))])
	return nil
}

func cmd_vote(i *Instance, m *Message, args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("Missing argument.")
	}

	if args[0] == "start" {
		i.StartVote(m.Channel)
		i.ChannelMsg(m.Channel, fmt.Sprintf("@%s has started a new vote! Tell me if you would like to `vote yes` or `vote no` on it.", i.Users[m.User].Name))
	} else if args[0] == "stop" {
		votes := i.StopVote(m.Channel)
		var result string
		if votes > 0 {
			result = "*YES* has won"
		} else if votes < 0 {
			result = "*NO* has won"
		} else {
			result = "It's a tie! No one won"
		}
		i.ChannelMsg(m.Channel, fmt.Sprintf("@%s has stopped the vote! The result is...\n%s (score: %d)!", i.Users[m.User].Name, result, votes))
	} else if args[0] == "yes" {
		i.Vote(+1, m.Channel)
	} else if args[0] == "no" {
		i.Vote(-1, m.Channel)
	}
	return nil
}

func cmd_mute(i *Instance, m *Message, args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("Missing argument.")
	}

	dur, e := time.ParseDuration(args[0])
	if e != nil {
		return fmt.Errorf("Couldn't parse a duration.")
	}
	if dur.Minutes() < 1 {
		dur = time.Duration(5) * time.Minute
	}
	if dur.Minutes() > 60 {
		dur = time.Duration(60) * time.Minute
	}

	i.ChannelMsg(m.Channel, fmt.Sprintf("I will now shut up for ~%.0f minutes.", dur.Minutes()))
	go func() {
		time.Sleep(5)
		i.Mute(m.Channel, dur)
	}()
	return nil
}

func cmd_status(i *Instance, m *Message, args []string) error {
	tmp := "My current status:\n"
	tmp += fmt.Sprintf("`Running: %v`\n", i.running)
	tmp += fmt.Sprintf("`Version: v%v`\n", VERSION)
	tmp += fmt.Sprintf("`Debug mode: %v`\n", i.Debug)
	tmp += fmt.Sprintf("`Latency: %v`\n", i.Latency)
	tmp += fmt.Sprintf("`My User ID: @%v`\n", i.BotId)
	tmp += fmt.Sprintf("`Tracked channels: %v`\n", len(i.Channels))
	tmp += fmt.Sprintf("`Tracked users: %v`\n", len(i.Users))
	tmp += fmt.Sprintf("`No. of Goon quotes: %v`\n", len(i.goon))
	tmp += fmt.Sprintf("`No. of Apollo quotes: %v`\n", len(i.apollo))
	tmp += fmt.Sprintf("\nI am currently muted on channels:\n")
	if len(i.mutes) < 1 {
		tmp += "- `None`\n"
	} else {
		for k, v := range i.mutes {
			c := i.Channels[k].Name
			t := v.Sub(time.Now())
			tmp += fmt.Sprintf("- `%v (ends in ~%.0f minutes)`\n", c, t.Minutes())
		}
	}
	i.UserMsg(m.User, tmp)
	return nil
}

func cmd_roll(i *Instance, m *Message, args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("Missing argument.")
	}

	ret, res, e := dice.Roll(args[0])
	if e != nil {
		return fmt.Errorf("Bad dice format.")
	}
	i.ChannelMsg(m.Channel, fmt.Sprintf("@%s has rolled: %v, %v\n", i.Users[m.User].Name, ret, res))
	return nil
}

func cmd_wiki(i *Instance, m *Message, args []string) error {
	if len(args) < 1 {
		return fmt.Errorf("Missing argument.")
	}

	u := WIKI_URL + url.QueryEscape(args[0])
	doc, e := goquery.NewDocument(u)
	if e != nil {
		return fmt.Errorf("Couldn't open the wiki for you (%v).", u)
	}

	node := doc.Find("div #mw-content-text > p").First()
	text := strings.TrimSpace(node.Text())
	if node.Length() < 1 || len(text) < 1 {
		return fmt.Errorf("Couldn't quote that page for you (%v).", u)
	}

	i.ChannelMsg(m.Channel, fmt.Sprintf(">>>%s\n`Source: %s`", text, u))
	return nil
}

func cmd_pun(i *Instance, m *Message, args []string) error {
	resp, e := http.Get("http://www.punoftheday.com/cgi-bin/arandompun.pl")
	if e != nil {
		return fmt.Errorf("Sorry, couldn't make up a pun for you.")
	}
	defer resp.Body.Close()
	body, e := ioutil.ReadAll(resp.Body)
	if e != nil {
		return fmt.Errorf("Sorry, couldn't make up a pun for you.")
	}
	// Yep this is a fucking mess.
	s := strings.TrimSpace(string(body))
	s = strings.TrimPrefix(s, "document.write('"")
	index := strings.Index(s, """)
	s = html.UnescapeString(s[:index])
	i.ChannelMsg(m.Channel, fmt.Sprintf(">>>%v\n`© 1996-2011 PunoftheDay.com`", s))
	return nil
}

type CatFacts struct {
	Facts   []string
	Success string
}

var LAST_CATFACT_TIME time.Time

func cmd_catfact(i *Instance, m *Message, args []string) error {
	now := time.Now()
	if now.After(LAST_CATFACT_TIME) != true {
		return fmt.Errorf("Sorry, I can only show a single catfact per day (or else I'll get banned!).")
	}
	LAST_CATFACT_TIME = now.Add(24 * time.Hour)

	resp, e := http.Get("http://catfacts-api.appspot.com/api/facts?number=1")
	if e != nil {
		return fmt.Errorf("Sorry, couldn't find any cat facts.")
	}

	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	var cf CatFacts
	e = decoder.Decode(&cf)
	if e != nil {
		return fmt.Errorf("Sorry, couldn't find any cat facts.")
	}

	if len(cf.Facts) > 0 {
		msg := fmt.Sprintf(">>>%v\n`© http://catfacts-api.appspot.com`", cf.Facts[0])
		i.ChannelMsg(m.Channel, msg)
		return nil
	}
	return fmt.Errorf("Sorry, couldn't find any cat facts.")
}

func cmd_goon(i *Instance, m *Message, args []string) error {
	q := i.random_goon_quote()
	msg := fmt.Sprintf(">>>%v\n`%v`", q.Quote, q.File)
	i.ChannelMsg(m.Channel, msg)
	return nil
}
func cmd_apollo(i *Instance, m *Message, args []string) error {
	q := i.random_apollo_quote()
	msg := fmt.Sprintf(">>>%v\n`%v`", q.Quote, q.File)
	i.ChannelMsg(m.Channel, msg)
	return nil
}

func cmd_channels(i *Instance, m *Message, args []string) error {
	tmp := fmt.Sprintf("Known channels (%d):\n", len(i.Channels))

	for id, c := range i.Channels {
		tmp += fmt.Sprintf("`#%s %s - %s`\n", id, c.Name, c.Topic.Value)
	}
	i.UserMsg(m.User, tmp)
	return nil
}

func cmd_users(i *Instance, m *Message, args []string) error {
	tmp := fmt.Sprintf("Known users (%d):\n", len(i.Users))

	for id, u := range i.Users {
		tmp += fmt.Sprintf("`#%s %s - %s`\n", id, u.Name, u.Presence)
	}
	i.UserMsg(m.User, tmp)
	return nil
}

func cmd_start_talking(i *Instance, m *Message, args []string) error {
	if m.Channel != TALK_CHANNEL {
		return fmt.Errorf("Nuhuh, to shy to talk there.")
	}

	if i.talking {
		i.ChannelMsg(TALK_CHANNEL, "I'm already talking! Stop bothering me!")
	} else {
		i.ChannelMsg(TALK_CHANNEL, "Helloooo..")
	}
	i.talking = true
	return nil
}

func cmd_stop_talking(i *Instance, m *Message, args []string) error {
	if m.Channel != TALK_CHANNEL {
		return nil
	}

	if i.talking {
		i.ChannelMsg(TALK_CHANNEL, "kthx bye!")
	} else {
		i.ChannelMsg(TALK_CHANNEL, "Haven't said a word, Cap'.")
	}
	i.talking = false
	return nil
}

func cmd_arrest(i *Instance, m *Message, args []string) error {
	tmp := "arrest " + strings.Join(args, " ")
	resp, e := i.cleverbot.Ask(tmp)
	if e != nil {
		fmt.Println("cleverbot error:", e)
		resp = "But I don't want to!"
	}
	i.ChannelMsg(m.Channel, resp)
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package hash

import (
	"crypto/sha1"
	"encoding/hex"
	"fmt"
)

type Sig [sha1.Size]byte

func Hash(data []byte) Sig {
	return Sig(sha1.Sum(data))
}

func String(data []byte) string {
	sig := Hash(data)
	return sig.String()
}

func (sig *Sig) String() string {
	return hex.EncodeToString((*sig)[:])
}

func FromString(str string) (Sig, error) {
	bin, err := hex.DecodeString(str)
	if err != nil {
		return Sig{}, fmt.Errorf("failed to decode sig '%v': %v", str, err)
	}
	if len(bin) != len(Sig{}) {
		return Sig{}, fmt.Errorf("failed to decode sig '%v': bad len", str)
	}
	var sig Sig
	for i, v := range bin {
		sig[i] = v
	}
	return sig, err
}
<commit_msg>pkg/hash: allow to hash multiple byte slices<commit_after>// Copyright 2016 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package hash

import (
	"bytes"
	"crypto/sha1"
	"encoding/binary"
	"encoding/hex"
	"fmt"
)

type Sig [sha1.Size]byte

func Hash(pieces ...[]byte) Sig {
	h := sha1.New()
	for _, data := range pieces {
		h.Write(data)
	}
	var sig Sig
	copy(sig[:], h.Sum(nil))
	return sig
}

func String(pieces ...[]byte) string {
	sig := Hash(pieces...)
	return sig.String()
}

func (sig *Sig) String() string {
	return hex.EncodeToString((*sig)[:])
}

// Truncate64 returns first 64 bits of the hash as int64.
func (sig *Sig) Truncate64() int64 {
	var v int64
	if err := binary.Read(bytes.NewReader((*sig)[:]), binary.LittleEndian, &v); err != nil {
		panic(fmt.Sprintf("failed convert hash to id: %v", err))
	}
	return v
}

func FromString(str string) (Sig, error) {
	bin, err := hex.DecodeString(str)
	if err != nil {
		return Sig{}, fmt.Errorf("failed to decode sig '%v': %v", str, err)
	}
	if len(bin) != len(Sig{}) {
		return Sig{}, fmt.Errorf("failed to decode sig '%v': bad len", str)
	}
	var sig Sig
	for i, v := range bin {
		sig[i] = v
	}
	return sig, err
}
<|endoftext|>


<commit_before>// Copyright 2016-2017 Authors of Cilium
//
// 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 policy

import (
	"bytes"
	"encoding/json"
	"fmt"
	"sort"
	"strconv"

	"github.com/cilium/cilium/api/v1/models"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/policy/api"
	"github.com/cilium/cilium/pkg/u8proto"
)

// L7DataMap contains a map of L7 rules per endpoint where key is a hash of EndpointSelector
type L7DataMap map[api.EndpointSelector]api.L7Rules

func (l7 L7DataMap) MarshalJSON() ([]byte, error) {
	if len(l7) == 0 {
		return []byte("[]"), nil
	}

	/* First, create a sorted slice of the selectors so we can get
	 * consistent JSON output */
	selectors := make(api.EndpointSelectorSlice, 0, len(l7))
	for es := range l7 {
		selectors = append(selectors, es)
	}
	sort.Sort(selectors)

	/* Now we can iterate the slice and generate JSON entries. */
	var err error
	buffer := bytes.NewBufferString("[")
	for _, es := range selectors {
		buffer.WriteString("{\"")
		buffer.WriteString(es.LabelSelectorString())
		buffer.WriteString("\":")
		b, err := json.Marshal(l7[es])
		if err == nil {
			buffer.Write(b)
		} else {
			buffer.WriteString("\"L7DataMap error: ")
			buffer.WriteString(err.Error())
			buffer.WriteString("\"")
		}
		buffer.WriteString("},")
	}
	buffer.Truncate(buffer.Len() - 1) // Drop the final ","
	buffer.WriteString("]")

	return buffer.Bytes(), err
}

// L7ParserType is the type used to indicate what L7 parser to use and
// defines all supported types of L7 parsers
type L7ParserType string

const (
	// ParserTypeNone represents the case where no parser type is provided.
	ParserTypeNone L7ParserType = ""
	// ParserTypeHTTP specifies a HTTP parser type
	ParserTypeHTTP L7ParserType = "http"
	// ParserTypeKafka specifies a Kafka parser type
	ParserTypeKafka L7ParserType = "kafka"
)

type L4Filter struct {
	// Port is the destination port to allow
	Port int `json:"port"`
	// Protocol is the L4 protocol to allow or NONE
	Protocol api.L4Proto `json:"protocol"`
	// U8Proto is the Protocol in numeric format, or 0 for NONE
	U8Proto u8proto.U8proto `json:"-"`
	// Endpoints limits the labels for allowing traffic (to / from).
	Endpoints api.EndpointSelectorSlice `json:"-"`
	// L7Parser specifies the L7 protocol parser (optional). If specified as
	// an empty string, then means that no L7 proxy redirect is performed.
	L7Parser L7ParserType `json:"-"`
	// L7RulesPerEp is a list of L7 rules per endpoint passed to the L7 proxy (optional)
	L7RulesPerEp L7DataMap `json:"l7-rules,omitempty"`
	// Ingress is true if filter applies at ingress; false if it applies at egress.
	Ingress bool `json:"-"`
	// The rule labels of this Filter
	DerivedFromRules labels.LabelArrayList `json:"-"`
}

// AllowsAllAtL3 returns whether this L4Filter applies to all endpoints at L3.
func (l4 *L4Filter) AllowsAllAtL3() bool {
	return l4.Endpoints.SelectsAllEndpoints()
}

// GetRelevantRules returns the relevant rules based on the source and
// destination addressing/identity information.
func (l7 L7DataMap) GetRelevantRules(identity *identity.Identity) api.L7Rules {
	rules := api.L7Rules{}

	if identity != nil {
		for selector, endpointRules := range l7 {
			if selector.Matches(identity.Labels.LabelArray()) {
				rules.HTTP = append(rules.HTTP, endpointRules.HTTP...)
				rules.Kafka = append(rules.Kafka, endpointRules.Kafka...)
			}
		}
	}

	// Rules applying to all sources are always appended
	if r, ok := l7[api.WildcardEndpointSelector]; ok {
		rules.HTTP = append(rules.HTTP, r.HTTP...)
		rules.Kafka = append(rules.Kafka, r.Kafka...)
	}

	return rules
}

func (l7 L7DataMap) addRulesForEndpoints(rules api.L7Rules, endpoints []api.EndpointSelector) {
	if rules.Len() == 0 {
		return
	}

	if len(endpoints) > 0 {
		for _, epsel := range endpoints {
			l7[epsel] = rules
		}
	} else {
		// If there are no explicit fromEps, have a 'special' wildcard endpoint.
		l7[api.WildcardEndpointSelector] = rules
	}
}

// CreateL4IngressFilter creates an L4Filter for the specified api.PortProtocol in
// the ingress direction for a particular protocol.
// This L4Filter will only apply to endpoints covered by `fromEndpoints`.
// `rule` allows a series of L7 rules to be associated with this L4Filter.
func CreateL4IngressFilter(fromEndpoints api.EndpointSelectorSlice, rule api.PortRule, port api.PortProtocol,
	protocol api.L4Proto, ruleLabels labels.LabelArray) L4Filter {

	// already validated via PortRule.Validate()
	p, _ := strconv.ParseUint(port.Port, 0, 16)
	// already validated via L4Proto.Validate()
	u8p, _ := u8proto.ParseProtocol(string(protocol))

	filterEndpoints := fromEndpoints
	if fromEndpoints.SelectsAllEndpoints() {
		filterEndpoints = api.EndpointSelectorSlice{api.WildcardEndpointSelector}

	}

	l4 := L4Filter{
		Port:             int(p),
		Protocol:         protocol,
		U8Proto:          u8p,
		L7RulesPerEp:     make(L7DataMap),
		Endpoints:        filterEndpoints,
		DerivedFromRules: labels.LabelArrayList{ruleLabels},
	}

	l4.Ingress = true

	if protocol == api.ProtoTCP && rule.Rules != nil {
		switch {
		case len(rule.Rules.HTTP) > 0:
			l4.L7Parser = ParserTypeHTTP
		case len(rule.Rules.Kafka) > 0:
			l4.L7Parser = ParserTypeKafka
		}

		l4.L7RulesPerEp.addRulesForEndpoints(*rule.Rules, fromEndpoints)
	}

	return l4
}

// CreateL4EgressFilter creates an L4Filter for the specified api.PortProtocol in
// the egress direction for a particular protocol.
// This L4Filter will only apply to endpoints covered by `fromEndpoints`.
// `rule` allows a series of L7 rules to be associated with this L4Filter.
func CreateL4EgressFilter(toEndpoints api.EndpointSelectorSlice, rule api.PortRule, port api.PortProtocol,
	protocol api.L4Proto, ruleLabels labels.LabelArray) L4Filter {

	// already validated via PortRule.Validate()
	p, _ := strconv.ParseUint(port.Port, 0, 16)
	// already validated via L4Proto.Validate()
	u8p, _ := u8proto.ParseProtocol(string(protocol))

	filterEndpoints := toEndpoints
	if toEndpoints.SelectsAllEndpoints() {
		filterEndpoints = api.EndpointSelectorSlice{api.WildcardEndpointSelector}

	}

	l4 := L4Filter{
		Port:             int(p),
		Protocol:         protocol,
		U8Proto:          u8p,
		L7RulesPerEp:     make(L7DataMap),
		Endpoints:        filterEndpoints,
		DerivedFromRules: labels.LabelArrayList{ruleLabels},
	}

	l4.Ingress = false

	if protocol == api.ProtoTCP && rule.Rules != nil {
		switch {
		case len(rule.Rules.HTTP) > 0:
			l4.L7Parser = ParserTypeHTTP
		case len(rule.Rules.Kafka) > 0:
			l4.L7Parser = ParserTypeKafka
		}

		l4.L7RulesPerEp.addRulesForEndpoints(*rule.Rules, toEndpoints)
	}

	return l4
}

// IsRedirect returns true if the L4 filter contains a port redirection
func (l4 *L4Filter) IsRedirect() bool {
	return l4.L7Parser != ParserTypeNone
}

// MarshalIndent returns the `L4Filter` in indented JSON string.
func (l4 *L4Filter) MarshalIndent() string {
	b, err := json.MarshalIndent(l4, "", "  ")
	if err != nil {
		b = []byte("\"L4Filter error: " + err.Error() + "\"")
	}
	return string(b)
}

// String returns the `L4Filter` in a human-readable string.
func (l4 L4Filter) String() string {
	b, err := json.Marshal(l4)
	if err != nil {
		return err.Error()
	}
	return string(b)
}

func (l4 L4Filter) matchesLabels(labels labels.LabelArray) bool {
	if l4.AllowsAllAtL3() {
		return true
	} else if len(labels) == 0 {
		return false
	}

	for _, sel := range l4.Endpoints {
		if sel.Matches(labels) {
			return true
		}
	}

	return false
}

// L4PolicyMap is a list of L4 filters indexable by protocol/port
// key format: "port/proto"
type L4PolicyMap map[string]L4Filter

// HasRedirect returns true if at least one L4 filter contains a port
// redirection
func (l4 L4PolicyMap) HasRedirect() bool {
	for _, f := range l4 {
		if f.IsRedirect() {
			return true
		}
	}

	return false
}

// containsAllL3L4 checks if the L4PolicyMap contains all L4 ports in `ports`.
// For L4Filters that specify ToEndpoints or FromEndpoints, uses `labels` to
// determine whether the policy allows L4 communication between the corresponding
// endpoints.
// Returns api.Denied in the following conditions:
// * If the `L4PolicyMap` has at least one rule and `ports` is empty.
// * If a single port is not present in the `L4PolicyMap`.
// * If a port is present in the `L4PolicyMap`, but it applies ToEndpoints or
// FromEndpoints constraints that require labels not present in `labels`.
// Otherwise, returns api.Allowed.
func (l4 L4PolicyMap) containsAllL3L4(labels labels.LabelArray, ports []*models.Port) api.Decision {
	if len(l4) == 0 {
		return api.Allowed
	}

	if len(ports) == 0 {
		return api.Denied
	}

	for _, l4Ctx := range ports {
		lwrProtocol := l4Ctx.Protocol
		switch lwrProtocol {
		case "", models.PortProtocolANY:
			tcpPort := fmt.Sprintf("%d/TCP", l4Ctx.Port)
			tcpFilter, tcpmatch := l4[tcpPort]
			if tcpmatch {
				tcpmatch = tcpFilter.matchesLabels(labels)
			}
			udpPort := fmt.Sprintf("%d/UDP", l4Ctx.Port)
			udpFilter, udpmatch := l4[udpPort]
			if udpmatch {
				udpmatch = udpFilter.matchesLabels(labels)
			}
			if !tcpmatch && !udpmatch {
				return api.Denied
			}
		default:
			port := fmt.Sprintf("%d/%s", l4Ctx.Port, lwrProtocol)
			filter, match := l4[port]
			if !match || !filter.matchesLabels(labels) {
				return api.Denied
			}
		}
	}
	return api.Allowed
}

type L4Policy struct {
	Ingress L4PolicyMap
	Egress  L4PolicyMap

	// Revision is the repository revision used to generate this policy.
	Revision uint64
}

func NewL4Policy() *L4Policy {
	return &L4Policy{
		Ingress:  L4PolicyMap{},
		Egress:   L4PolicyMap{},
		Revision: 0,
	}
}

// IngressCoversContext checks if the receiver's ingress `L4Policy` contains
// all `dPorts` and `labels`.
func (l4 *L4PolicyMap) IngressCoversContext(ctx *SearchContext) api.Decision {
	return l4.containsAllL3L4(ctx.From, ctx.DPorts)
}

// EgressCoversContext checks if the receiver's egress `L4Policy` contains
// all `dPorts` and `labels`.
func (l4 *L4PolicyMap) EgressCoversContext(ctx *SearchContext) api.Decision {
	return l4.containsAllL3L4(ctx.To, ctx.DPorts)
}

// HasRedirect returns true if the L4 policy contains at least one port redirection
func (l4 *L4Policy) HasRedirect() bool {
	return l4 != nil && (l4.Ingress.HasRedirect() || l4.Egress.HasRedirect())
}

// RequiresConntrack returns true if if the L4 configuration requires
// connection tracking to be enabled.
func (l4 *L4Policy) RequiresConntrack() bool {
	return l4 != nil && (len(l4.Ingress) > 0 || len(l4.Egress) > 0)
}

func (l4 *L4Policy) GetModel() *models.L4Policy {
	if l4 == nil {
		return nil
	}

	ingress := []*models.PolicyRule{}
	for _, v := range l4.Ingress {
		ingress = append(ingress, &models.PolicyRule{
			Rule:             v.MarshalIndent(),
			DerivedFromRules: v.DerivedFromRules.GetModel(),
		})
	}

	egress := []*models.PolicyRule{}
	for _, v := range l4.Egress {
		egress = append(egress, &models.PolicyRule{
			Rule:             v.MarshalIndent(),
			DerivedFromRules: v.DerivedFromRules.GetModel(),
		})
	}

	return &models.L4Policy{
		Ingress: ingress,
		Egress:  egress,
	}
}
<commit_msg>policy: Refactor L4Filter creation<commit_after>// Copyright 2016-2018 Authors of Cilium
//
// 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 policy

import (
	"bytes"
	"encoding/json"
	"fmt"
	"sort"
	"strconv"

	"github.com/cilium/cilium/api/v1/models"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/policy/api"
	"github.com/cilium/cilium/pkg/u8proto"
)

// L7DataMap contains a map of L7 rules per endpoint where key is a hash of EndpointSelector
type L7DataMap map[api.EndpointSelector]api.L7Rules

func (l7 L7DataMap) MarshalJSON() ([]byte, error) {
	if len(l7) == 0 {
		return []byte("[]"), nil
	}

	/* First, create a sorted slice of the selectors so we can get
	 * consistent JSON output */
	selectors := make(api.EndpointSelectorSlice, 0, len(l7))
	for es := range l7 {
		selectors = append(selectors, es)
	}
	sort.Sort(selectors)

	/* Now we can iterate the slice and generate JSON entries. */
	var err error
	buffer := bytes.NewBufferString("[")
	for _, es := range selectors {
		buffer.WriteString("{\"")
		buffer.WriteString(es.LabelSelectorString())
		buffer.WriteString("\":")
		b, err := json.Marshal(l7[es])
		if err == nil {
			buffer.Write(b)
		} else {
			buffer.WriteString("\"L7DataMap error: ")
			buffer.WriteString(err.Error())
			buffer.WriteString("\"")
		}
		buffer.WriteString("},")
	}
	buffer.Truncate(buffer.Len() - 1) // Drop the final ","
	buffer.WriteString("]")

	return buffer.Bytes(), err
}

// L7ParserType is the type used to indicate what L7 parser to use and
// defines all supported types of L7 parsers
type L7ParserType string

const (
	// ParserTypeNone represents the case where no parser type is provided.
	ParserTypeNone L7ParserType = ""
	// ParserTypeHTTP specifies a HTTP parser type
	ParserTypeHTTP L7ParserType = "http"
	// ParserTypeKafka specifies a Kafka parser type
	ParserTypeKafka L7ParserType = "kafka"
)

type L4Filter struct {
	// Port is the destination port to allow
	Port int `json:"port"`
	// Protocol is the L4 protocol to allow or NONE
	Protocol api.L4Proto `json:"protocol"`
	// U8Proto is the Protocol in numeric format, or 0 for NONE
	U8Proto u8proto.U8proto `json:"-"`
	// Endpoints limits the labels for allowing traffic (to / from).
	Endpoints api.EndpointSelectorSlice `json:"-"`
	// L7Parser specifies the L7 protocol parser (optional). If specified as
	// an empty string, then means that no L7 proxy redirect is performed.
	L7Parser L7ParserType `json:"-"`
	// L7RulesPerEp is a list of L7 rules per endpoint passed to the L7 proxy (optional)
	L7RulesPerEp L7DataMap `json:"l7-rules,omitempty"`
	// Ingress is true if filter applies at ingress; false if it applies at egress.
	Ingress bool `json:"-"`
	// The rule labels of this Filter
	DerivedFromRules labels.LabelArrayList `json:"-"`
}

// AllowsAllAtL3 returns whether this L4Filter applies to all endpoints at L3.
func (l4 *L4Filter) AllowsAllAtL3() bool {
	return l4.Endpoints.SelectsAllEndpoints()
}

// GetRelevantRules returns the relevant rules based on the source and
// destination addressing/identity information.
func (l7 L7DataMap) GetRelevantRules(identity *identity.Identity) api.L7Rules {
	rules := api.L7Rules{}

	if identity != nil {
		for selector, endpointRules := range l7 {
			if selector.Matches(identity.Labels.LabelArray()) {
				rules.HTTP = append(rules.HTTP, endpointRules.HTTP...)
				rules.Kafka = append(rules.Kafka, endpointRules.Kafka...)
			}
		}
	}

	// Rules applying to all sources are always appended
	if r, ok := l7[api.WildcardEndpointSelector]; ok {
		rules.HTTP = append(rules.HTTP, r.HTTP...)
		rules.Kafka = append(rules.Kafka, r.Kafka...)
	}

	return rules
}

func (l7 L7DataMap) addRulesForEndpoints(rules api.L7Rules, endpoints []api.EndpointSelector) {
	if rules.Len() == 0 {
		return
	}

	if len(endpoints) > 0 {
		for _, epsel := range endpoints {
			l7[epsel] = rules
		}
	} else {
		// If there are no explicit fromEps, have a 'special' wildcard endpoint.
		l7[api.WildcardEndpointSelector] = rules
	}
}

// CreateL4Filter creates a filter for L4 policy that applies to the specified
// endpoints and port/protocol, with reference to the original rules that the
// filter is derived from. This filter may be associated with a series of L7
// rules via the `rule` parameter.
func CreateL4Filter(peerEndpoints api.EndpointSelectorSlice, rule api.PortRule, port api.PortProtocol,
	protocol api.L4Proto, ruleLabels labels.LabelArray, ingress bool) L4Filter {

	// already validated via PortRule.Validate()
	p, _ := strconv.ParseUint(port.Port, 0, 16)
	// already validated via L4Proto.Validate()
	u8p, _ := u8proto.ParseProtocol(string(protocol))

	filterEndpoints := peerEndpoints
	if peerEndpoints.SelectsAllEndpoints() {
		filterEndpoints = api.EndpointSelectorSlice{api.WildcardEndpointSelector}
	}

	l4 := L4Filter{
		Port:             int(p),
		Protocol:         protocol,
		U8Proto:          u8p,
		L7RulesPerEp:     make(L7DataMap),
		Endpoints:        filterEndpoints,
		DerivedFromRules: labels.LabelArrayList{ruleLabels},
		Ingress:          ingress,
	}

	if protocol == api.ProtoTCP && rule.Rules != nil {
		switch {
		case len(rule.Rules.HTTP) > 0:
			l4.L7Parser = ParserTypeHTTP
		case len(rule.Rules.Kafka) > 0:
			l4.L7Parser = ParserTypeKafka
		}
		l4.L7RulesPerEp.addRulesForEndpoints(*rule.Rules, peerEndpoints)
	}

	return l4
}

// CreateL4IngressFilter creates a filter for L4 policy that applies to the
// specified endpoints and port/protocol for ingress traffic, with reference
// to the original rules that the filter is derived from. This filter may be
// associated with a series of L7 rules via the `rule` parameter.
func CreateL4IngressFilter(fromEndpoints api.EndpointSelectorSlice, rule api.PortRule, port api.PortProtocol,
	protocol api.L4Proto, ruleLabels labels.LabelArray) L4Filter {

	return CreateL4Filter(fromEndpoints, rule, port, protocol, ruleLabels, true)
}

// CreateL4EgressFilter creates a filter for L4 policy that applies to the
// specified endpoints and port/protocol for egress traffic, with reference
// to the original rules that the filter is derived from. This filter may be
// associated with a series of L7 rules via the `rule` parameter.
func CreateL4EgressFilter(toEndpoints api.EndpointSelectorSlice, rule api.PortRule, port api.PortProtocol,
	protocol api.L4Proto, ruleLabels labels.LabelArray) L4Filter {

	return CreateL4Filter(toEndpoints, rule, port, protocol, ruleLabels, false)
}

// IsRedirect returns true if the L4 filter contains a port redirection
func (l4 *L4Filter) IsRedirect() bool {
	return l4.L7Parser != ParserTypeNone
}

// MarshalIndent returns the `L4Filter` in indented JSON string.
func (l4 *L4Filter) MarshalIndent() string {
	b, err := json.MarshalIndent(l4, "", "  ")
	if err != nil {
		b = []byte("\"L4Filter error: " + err.Error() + "\"")
	}
	return string(b)
}

// String returns the `L4Filter` in a human-readable string.
func (l4 L4Filter) String() string {
	b, err := json.Marshal(l4)
	if err != nil {
		return err.Error()
	}
	return string(b)
}

func (l4 L4Filter) matchesLabels(labels labels.LabelArray) bool {
	if l4.AllowsAllAtL3() {
		return true
	} else if len(labels) == 0 {
		return false
	}

	for _, sel := range l4.Endpoints {
		if sel.Matches(labels) {
			return true
		}
	}

	return false
}

// L4PolicyMap is a list of L4 filters indexable by protocol/port
// key format: "port/proto"
type L4PolicyMap map[string]L4Filter

// HasRedirect returns true if at least one L4 filter contains a port
// redirection
func (l4 L4PolicyMap) HasRedirect() bool {
	for _, f := range l4 {
		if f.IsRedirect() {
			return true
		}
	}

	return false
}

// containsAllL3L4 checks if the L4PolicyMap contains all L4 ports in `ports`.
// For L4Filters that specify ToEndpoints or FromEndpoints, uses `labels` to
// determine whether the policy allows L4 communication between the corresponding
// endpoints.
// Returns api.Denied in the following conditions:
// * If the `L4PolicyMap` has at least one rule and `ports` is empty.
// * If a single port is not present in the `L4PolicyMap`.
// * If a port is present in the `L4PolicyMap`, but it applies ToEndpoints or
// FromEndpoints constraints that require labels not present in `labels`.
// Otherwise, returns api.Allowed.
func (l4 L4PolicyMap) containsAllL3L4(labels labels.LabelArray, ports []*models.Port) api.Decision {
	if len(l4) == 0 {
		return api.Allowed
	}

	if len(ports) == 0 {
		return api.Denied
	}

	for _, l4Ctx := range ports {
		lwrProtocol := l4Ctx.Protocol
		switch lwrProtocol {
		case "", models.PortProtocolANY:
			tcpPort := fmt.Sprintf("%d/TCP", l4Ctx.Port)
			tcpFilter, tcpmatch := l4[tcpPort]
			if tcpmatch {
				tcpmatch = tcpFilter.matchesLabels(labels)
			}
			udpPort := fmt.Sprintf("%d/UDP", l4Ctx.Port)
			udpFilter, udpmatch := l4[udpPort]
			if udpmatch {
				udpmatch = udpFilter.matchesLabels(labels)
			}
			if !tcpmatch && !udpmatch {
				return api.Denied
			}
		default:
			port := fmt.Sprintf("%d/%s", l4Ctx.Port, lwrProtocol)
			filter, match := l4[port]
			if !match || !filter.matchesLabels(labels) {
				return api.Denied
			}
		}
	}
	return api.Allowed
}

type L4Policy struct {
	Ingress L4PolicyMap
	Egress  L4PolicyMap

	// Revision is the repository revision used to generate this policy.
	Revision uint64
}

func NewL4Policy() *L4Policy {
	return &L4Policy{
		Ingress:  L4PolicyMap{},
		Egress:   L4PolicyMap{},
		Revision: 0,
	}
}

// IngressCoversContext checks if the receiver's ingress `L4Policy` contains
// all `dPorts` and `labels`.
func (l4 *L4PolicyMap) IngressCoversContext(ctx *SearchContext) api.Decision {
	return l4.containsAllL3L4(ctx.From, ctx.DPorts)
}

// EgressCoversContext checks if the receiver's egress `L4Policy` contains
// all `dPorts` and `labels`.
func (l4 *L4PolicyMap) EgressCoversContext(ctx *SearchContext) api.Decision {
	return l4.containsAllL3L4(ctx.To, ctx.DPorts)
}

// HasRedirect returns true if the L4 policy contains at least one port redirection
func (l4 *L4Policy) HasRedirect() bool {
	return l4 != nil && (l4.Ingress.HasRedirect() || l4.Egress.HasRedirect())
}

// RequiresConntrack returns true if if the L4 configuration requires
// connection tracking to be enabled.
func (l4 *L4Policy) RequiresConntrack() bool {
	return l4 != nil && (len(l4.Ingress) > 0 || len(l4.Egress) > 0)
}

func (l4 *L4Policy) GetModel() *models.L4Policy {
	if l4 == nil {
		return nil
	}

	ingress := []*models.PolicyRule{}
	for _, v := range l4.Ingress {
		ingress = append(ingress, &models.PolicyRule{
			Rule:             v.MarshalIndent(),
			DerivedFromRules: v.DerivedFromRules.GetModel(),
		})
	}

	egress := []*models.PolicyRule{}
	for _, v := range l4.Egress {
		egress = append(egress, &models.PolicyRule{
			Rule:             v.MarshalIndent(),
			DerivedFromRules: v.DerivedFromRules.GetModel(),
		})
	}

	return &models.L4Policy{
		Ingress: ingress,
		Egress:  egress,
	}
}
<|endoftext|>