<commit_before>package server

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rsa"
	"errors"
	"fmt"

	// Rather than crypto/x509 as ct allows disabling critical extension checks.
	"github.com/google/certificate-transparency-go/x509"
	"github.com/google/go-tpm-tools/internal"
	pb "github.com/google/go-tpm-tools/proto/attest"
	tpmpb "github.com/google/go-tpm-tools/proto/tpm"
	"github.com/google/go-tpm/tpm2"
	"google.golang.org/protobuf/proto"
)

// The hash algorithms we support, in their preferred order of use.
var supportedHashAlgs = []tpm2.Algorithm{
	tpm2.AlgSHA512, tpm2.AlgSHA384, tpm2.AlgSHA256, tpm2.AlgSHA1,
}

// VerifyOpts allows for customizing the functionality of VerifyAttestation.
type VerifyOpts struct {
	// The nonce used when calling client.Attest
	Nonce []byte
	// Trusted public keys that can be used to directly verify the key used for
	// attestation. This option should be used if you already know the AK, as
	// it provides the highest level of assurance.
	TrustedAKs []crypto.PublicKey
	// Allow attestations to be verified using SHA-1. This defaults to false
	// because SHA-1 is a weak hash algorithm with known collision attacks.
	// However, setting this to true may be necessary if the client only
	// supports the legacy event log format. This is the case on older Linux
	// distributions (such as Debian 10).
	AllowSHA1 bool
	// A collection of trusted root CAs that are used to sign AK certificates.
	// The TrustedAKs are used first, followed by TrustRootCerts and
	// IntermediateCerts.
	// Adding a specific TPM manufacturer's root and intermediate CAs means all
	// TPMs signed by that CA will be trusted.
	TrustedRootCerts  *x509.CertPool
	IntermediateCerts *x509.CertPool
}

// VerifyAttestation performs the following checks on an Attestation:
//    - the AK used to generate the attestation is trusted (based on VerifyOpts)
//    - the provided signature is generated by the trusted AK public key
//    - the signature signs the provided quote data
//    - the quote data starts with TPM_GENERATED_VALUE
//    - the quote data is a valid TPMS_QUOTE_INFO
//    - the quote data was taken over the provided PCRs
//    - the provided PCR values match the quote data internal digest
//    - the provided opts.Nonce matches that in the quote data
//    - the provided eventlog matches the provided PCR values
//
// After this, the eventlog is parsed and the corresponding MachineState is
// returned. This design prevents unverified MachineStates from being used.
func VerifyAttestation(attestation *pb.Attestation, opts VerifyOpts) (*pb.MachineState, error) {
	// Verify the AK
	akPubArea, err := tpm2.DecodePublic(attestation.GetAkPub())
	if err != nil {
		return nil, fmt.Errorf("failed to decode AK public area: %w", err)
	}
	akPubKey, err := akPubArea.Key()
	if err != nil {
		return nil, fmt.Errorf("failed to get AK public key: %w", err)
	}
	if err := checkAKTrusted(akPubKey, attestation.GetAkCert(), opts); err != nil {
		return nil, fmt.Errorf("failed to validate AK: %w", err)
	}

	// Verify the signing hash algorithm
	signHashAlg, err := internal.GetSigningHashAlg(akPubArea)
	if err != nil {
		return nil, fmt.Errorf("bad AK public area: %w", err)
	}
	if err = checkHashAlgSupported(signHashAlg, opts); err != nil {
		return nil, fmt.Errorf("in AK public area: %w", err)
	}

	// Attempt to replay the log against our PCRs in order of hash preference
	var lastErr error
	for _, quote := range supportedQuotes(attestation.GetQuotes()) {
		// Verify the Quote
		if err = internal.VerifyQuote(quote, akPubKey, opts.Nonce); err != nil {
			lastErr = fmt.Errorf("failed to verify quote: %w", err)
			continue
		}

		// Parse event logs and replay the events against the provided PCRs
		pcrs := quote.GetPcrs()
		state, err := parsePCClientEventLog(attestation.GetEventLog(), pcrs)
		if err != nil {
			lastErr = fmt.Errorf("failed to validate the PCClient event log: %w", err)
			continue
		}

		celState, err := parseCanonicalEventLog(attestation.GetCanonicalEventLog(), pcrs)
		if err != nil {
			lastErr = fmt.Errorf("failed to validate the Canonical event log: %w", err)
			continue
		}

		proto.Merge(celState, state)

		// Verify the PCR hash algorithm. We have this check here (instead of at
		// the start of the loop) so that the user gets a "SHA-1 not supported"
		// error only if allowing SHA-1 support would actually allow the log
		// to be verified. This makes debugging failed verifications easier.
		pcrHashAlg := tpm2.Algorithm(pcrs.GetHash())
		if err = checkHashAlgSupported(pcrHashAlg, opts); err != nil {
			lastErr = fmt.Errorf("when verifying PCRs: %w", err)
			continue
		}

		return celState, nil
	}

	if lastErr != nil {
		return nil, lastErr
	}
	return nil, fmt.Errorf("attestation does not contain a supported quote")
}

func pubKeysEqual(k1 crypto.PublicKey, k2 crypto.PublicKey) bool {
	switch key := k1.(type) {
	case *rsa.PublicKey:
		return key.Equal(k2)
	case *ecdsa.PublicKey:
		return key.Equal(k2)
	default:
		return false
	}
}

// Checks if the provided AK public key can be trusted
func checkAKTrusted(ak crypto.PublicKey, akCertBytes []byte, opts VerifyOpts) error {
	checkPub := len(opts.TrustedAKs) > 0
	checkCert := opts.TrustedRootCerts != nil && len(opts.TrustedRootCerts.Subjects()) > 0
	if !checkPub && !checkCert {
		return fmt.Errorf("no trust mechanism provided, either use TrustedAKs or TrustedRootCerts")
	}
	if checkPub && checkCert {
		return fmt.Errorf("multiple trust mechanisms provided, only use one of TrustedAKs or TrustedRootCerts")
	}

	// Check against known AKs
	if checkPub {
		for _, trusted := range opts.TrustedAKs {
			if pubKeysEqual(ak, trusted) {
				return nil
			}
		}
		return fmt.Errorf("public key is not trusted")
	}

	// Check if the AK Cert chains to a trusted root
	if len(akCertBytes) == 0 {
		return errors.New("no certificate provided in attestation")
	}
	akCert, err := x509.ParseCertificate(akCertBytes)
	if err != nil {
		return fmt.Errorf("failed to parse certificate: %w", err)
	}
	if !pubKeysEqual(ak, akCert.PublicKey) {
		return fmt.Errorf("mismatch between public key and certificate")
	}

	x509Opts := x509.VerifyOptions{
		Roots:         opts.TrustedRootCerts,
		Intermediates: opts.IntermediateCerts,
		// x509 (both ct and crypto) marks the SAN extension unhandled if SAN
		// does not parse any of DNSNames, EmailAddresses, IPAddresses, or URIs.
		// https://cs.opensource.google/go/go/+/master:src/crypto/x509/parser.go;l=668-678
		DisableCriticalExtensionChecks: true,
		// The default key usage (ExtKeyUsageServerAuth) is not appropriate for
		// an Attestation Key: ExtKeyUsage of
		// - https://oidref.com/2.23.133.8.1
		// - https://oidref.com/2.23.133.8.3
		// https://pkg.go.dev/crypto/x509#VerifyOptions
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsage(x509.ExtKeyUsageAny)},
	}
	if _, err := akCert.Verify(x509Opts); err != nil {
		return fmt.Errorf("failed to verify certificate against trusted roots: %v", err)
	}
	return nil
}

func checkHashAlgSupported(hash tpm2.Algorithm, opts VerifyOpts) error {
	if hash == tpm2.AlgSHA1 && !opts.AllowSHA1 {
		return fmt.Errorf("SHA-1 is not allowed for verification (set VerifyOpts.AllowSHA1 to true to allow)")
	}
	for _, alg := range supportedHashAlgs {
		if hash == alg {
			return nil
		}
	}
	return fmt.Errorf("unsupported hash algorithm: %v", hash)
}

// Retrieve the supported quotes in order of hash preference
func supportedQuotes(quotes []*tpmpb.Quote) []*tpmpb.Quote {
	out := make([]*tpmpb.Quote, 0, len(quotes))
	for _, alg := range supportedHashAlgs {
		for _, quote := range quotes {
			if tpm2.Algorithm(quote.GetPcrs().GetHash()) == alg {
				out = append(out, quote)
				break
			}
		}
	}
	return out
}
<commit_msg>server: Support all public key types in pubKeysEqual<commit_after>package server

import (
	"crypto"
	"errors"
	"fmt"

	// Rather than crypto/x509 as ct allows disabling critical extension checks.
	"github.com/google/certificate-transparency-go/x509"
	"github.com/google/go-tpm-tools/internal"
	pb "github.com/google/go-tpm-tools/proto/attest"
	tpmpb "github.com/google/go-tpm-tools/proto/tpm"
	"github.com/google/go-tpm/tpm2"
	"google.golang.org/protobuf/proto"
)

// The hash algorithms we support, in their preferred order of use.
var supportedHashAlgs = []tpm2.Algorithm{
	tpm2.AlgSHA512, tpm2.AlgSHA384, tpm2.AlgSHA256, tpm2.AlgSHA1,
}

// VerifyOpts allows for customizing the functionality of VerifyAttestation.
type VerifyOpts struct {
	// The nonce used when calling client.Attest
	Nonce []byte
	// Trusted public keys that can be used to directly verify the key used for
	// attestation. This option should be used if you already know the AK, as
	// it provides the highest level of assurance.
	TrustedAKs []crypto.PublicKey
	// Allow attestations to be verified using SHA-1. This defaults to false
	// because SHA-1 is a weak hash algorithm with known collision attacks.
	// However, setting this to true may be necessary if the client only
	// supports the legacy event log format. This is the case on older Linux
	// distributions (such as Debian 10).
	AllowSHA1 bool
	// A collection of trusted root CAs that are used to sign AK certificates.
	// The TrustedAKs are used first, followed by TrustRootCerts and
	// IntermediateCerts.
	// Adding a specific TPM manufacturer's root and intermediate CAs means all
	// TPMs signed by that CA will be trusted.
	TrustedRootCerts  *x509.CertPool
	IntermediateCerts *x509.CertPool
}

// VerifyAttestation performs the following checks on an Attestation:
//    - the AK used to generate the attestation is trusted (based on VerifyOpts)
//    - the provided signature is generated by the trusted AK public key
//    - the signature signs the provided quote data
//    - the quote data starts with TPM_GENERATED_VALUE
//    - the quote data is a valid TPMS_QUOTE_INFO
//    - the quote data was taken over the provided PCRs
//    - the provided PCR values match the quote data internal digest
//    - the provided opts.Nonce matches that in the quote data
//    - the provided eventlog matches the provided PCR values
//
// After this, the eventlog is parsed and the corresponding MachineState is
// returned. This design prevents unverified MachineStates from being used.
func VerifyAttestation(attestation *pb.Attestation, opts VerifyOpts) (*pb.MachineState, error) {
	// Verify the AK
	akPubArea, err := tpm2.DecodePublic(attestation.GetAkPub())
	if err != nil {
		return nil, fmt.Errorf("failed to decode AK public area: %w", err)
	}
	akPubKey, err := akPubArea.Key()
	if err != nil {
		return nil, fmt.Errorf("failed to get AK public key: %w", err)
	}
	if err := checkAKTrusted(akPubKey, attestation.GetAkCert(), opts); err != nil {
		return nil, fmt.Errorf("failed to validate AK: %w", err)
	}

	// Verify the signing hash algorithm
	signHashAlg, err := internal.GetSigningHashAlg(akPubArea)
	if err != nil {
		return nil, fmt.Errorf("bad AK public area: %w", err)
	}
	if err = checkHashAlgSupported(signHashAlg, opts); err != nil {
		return nil, fmt.Errorf("in AK public area: %w", err)
	}

	// Attempt to replay the log against our PCRs in order of hash preference
	var lastErr error
	for _, quote := range supportedQuotes(attestation.GetQuotes()) {
		// Verify the Quote
		if err = internal.VerifyQuote(quote, akPubKey, opts.Nonce); err != nil {
			lastErr = fmt.Errorf("failed to verify quote: %w", err)
			continue
		}

		// Parse event logs and replay the events against the provided PCRs
		pcrs := quote.GetPcrs()
		state, err := parsePCClientEventLog(attestation.GetEventLog(), pcrs)
		if err != nil {
			lastErr = fmt.Errorf("failed to validate the PCClient event log: %w", err)
			continue
		}

		celState, err := parseCanonicalEventLog(attestation.GetCanonicalEventLog(), pcrs)
		if err != nil {
			lastErr = fmt.Errorf("failed to validate the Canonical event log: %w", err)
			continue
		}

		proto.Merge(celState, state)

		// Verify the PCR hash algorithm. We have this check here (instead of at
		// the start of the loop) so that the user gets a "SHA-1 not supported"
		// error only if allowing SHA-1 support would actually allow the log
		// to be verified. This makes debugging failed verifications easier.
		pcrHashAlg := tpm2.Algorithm(pcrs.GetHash())
		if err = checkHashAlgSupported(pcrHashAlg, opts); err != nil {
			lastErr = fmt.Errorf("when verifying PCRs: %w", err)
			continue
		}

		return celState, nil
	}

	if lastErr != nil {
		return nil, lastErr
	}
	return nil, fmt.Errorf("attestation does not contain a supported quote")
}

func pubKeysEqual(k1 crypto.PublicKey, k2 crypto.PublicKey) bool {
	// Common interface for all public keys (see crypto.PublicKey documentation)
	type publicKey interface {
		Equal(crypto.PublicKey) bool
	}
	if key, ok := k1.(publicKey); ok {
		return key.Equal(k2)
	}
	return false
}

// Checks if the provided AK public key can be trusted
func checkAKTrusted(ak crypto.PublicKey, akCertBytes []byte, opts VerifyOpts) error {
	checkPub := len(opts.TrustedAKs) > 0
	checkCert := opts.TrustedRootCerts != nil && len(opts.TrustedRootCerts.Subjects()) > 0
	if !checkPub && !checkCert {
		return fmt.Errorf("no trust mechanism provided, either use TrustedAKs or TrustedRootCerts")
	}
	if checkPub && checkCert {
		return fmt.Errorf("multiple trust mechanisms provided, only use one of TrustedAKs or TrustedRootCerts")
	}

	// Check against known AKs
	if checkPub {
		for _, trusted := range opts.TrustedAKs {
			if pubKeysEqual(ak, trusted) {
				return nil
			}
		}
		return fmt.Errorf("public key is not trusted")
	}

	// Check if the AK Cert chains to a trusted root
	if len(akCertBytes) == 0 {
		return errors.New("no certificate provided in attestation")
	}
	akCert, err := x509.ParseCertificate(akCertBytes)
	if err != nil {
		return fmt.Errorf("failed to parse certificate: %w", err)
	}
	if !pubKeysEqual(ak, akCert.PublicKey) {
		return fmt.Errorf("mismatch between public key and certificate")
	}

	x509Opts := x509.VerifyOptions{
		Roots:         opts.TrustedRootCerts,
		Intermediates: opts.IntermediateCerts,
		// x509 (both ct and crypto) marks the SAN extension unhandled if SAN
		// does not parse any of DNSNames, EmailAddresses, IPAddresses, or URIs.
		// https://cs.opensource.google/go/go/+/master:src/crypto/x509/parser.go;l=668-678
		DisableCriticalExtensionChecks: true,
		// The default key usage (ExtKeyUsageServerAuth) is not appropriate for
		// an Attestation Key: ExtKeyUsage of
		// - https://oidref.com/2.23.133.8.1
		// - https://oidref.com/2.23.133.8.3
		// https://pkg.go.dev/crypto/x509#VerifyOptions
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsage(x509.ExtKeyUsageAny)},
	}
	if _, err := akCert.Verify(x509Opts); err != nil {
		return fmt.Errorf("failed to verify certificate against trusted roots: %v", err)
	}
	return nil
}

func checkHashAlgSupported(hash tpm2.Algorithm, opts VerifyOpts) error {
	if hash == tpm2.AlgSHA1 && !opts.AllowSHA1 {
		return fmt.Errorf("SHA-1 is not allowed for verification (set VerifyOpts.AllowSHA1 to true to allow)")
	}
	for _, alg := range supportedHashAlgs {
		if hash == alg {
			return nil
		}
	}
	return fmt.Errorf("unsupported hash algorithm: %v", hash)
}

// Retrieve the supported quotes in order of hash preference
func supportedQuotes(quotes []*tpmpb.Quote) []*tpmpb.Quote {
	out := make([]*tpmpb.Quote, 0, len(quotes))
	for _, alg := range supportedHashAlgs {
		for _, quote := range quotes {
			if tpm2.Algorithm(quote.GetPcrs().GetHash()) == alg {
				out = append(out, quote)
				break
			}
		}
	}
	return out
}
<|endoftext|>


<commit_before>// Package uuiddirs provides functions to split up a UUID
// into a series of sub-directories so that an unlimited number
// of UUIDs can be used as directories.
//
// Example:
//   The UUID f0498fad-437c-4954-ad82-8ec2cc202628 maps to the path
//   f0/498/fad/437c4954/ad828ec2cc202628
package uuiddir

import (
	"encoding/hex"
	"path"
	"strings"

	"github.com/domonda/errors"

	fs "github.com/ungerik/go-fs"
)

// Split a UUID into 5 hex strings.
// Example:
//   Splitting the UUID f0498fad-437c-4954-ad82-8ec2cc202628 returns
//   []string{"f0", "498", "fad", "437c4954", "ad828ec2cc202628"}
func Split(uuid [16]byte) []string {
	hexStr := hex.EncodeToString(uuid[:])
	return []string{
		hexStr[0:2],
		hexStr[2:5],
		hexStr[5:8],
		hexStr[8:16],
		hexStr[16:32],
	}
}

// Join returns a directory with the splitted UUID and pathParts joined to baseDir.
func Join(baseDir fs.File, uuid [16]byte, pathParts ...string) fs.File {
	return baseDir.Join(append(Split(uuid), pathParts...)...)
}

// Parse the path of uuidDir for a UUID
func Parse(uuidDir fs.File) (uuid [16]byte, err error) {
	uuidPath := strings.TrimSuffix(uuidDir.PathWithSlashes(), "/")
	if len(uuidPath) < 36 {
		return nilUUID, errors.Errorf("path can't be parsed as UUID: %q", string(uuidDir))
	}
	return ParseString(uuidPath[len(uuidPath)-36:])
}

// FormatString returns the splitted UUID joined with slashes.
// It's the inverse to ParseString.
func FormatString(uuid [16]byte) string {
	return path.Join(Split(uuid)...)
}

// ParseString parses a 36 character string (like returned from FormatString) as UUID.
func ParseString(uuidPath string) (uuid [16]byte, err error) {
	if len(uuidPath) != 36 {
		return nilUUID, errors.Errorf("path can't be parsed as UUID: %q", uuidPath)
	}
	uuidPath = strings.Replace(uuidPath, "/", "", 4)
	if len(uuidPath) != 32 {
		return nilUUID, errors.Errorf("path can't be parsed as UUID: %q", uuidPath)
	}
	b, err := hex.DecodeString(uuidPath)
	if err != nil {
		return nilUUID, errors.Errorf("path can't be parsed as UUID: %q", uuidPath)
	}
	copy(uuid[:], b)
	return uuid, validateUUID(uuid)
}

// Enum calls callback for every directory that represents an UUID under baseDir.
func Enum(baseDir fs.File, callback func(uuidDir fs.File, uuid [16]byte) error) error {
	return baseDir.ListDir(func(level0Dir fs.File) error {
		if !level0Dir.Exists() || level0Dir.IsHidden() {
			return nil
		}
		if !level0Dir.IsDir() {
			// fmt.Println("Directory expected but found file:", level0Dir)
			return nil
		}
		return level0Dir.ListDir(func(level1Dir fs.File) error {
			if !level1Dir.Exists() || level1Dir.IsHidden() {
				return nil
			}
			if !level1Dir.IsDir() {
				// fmt.Println("Directory expected but found file:", level1Dir)
				return nil
			}
			return level1Dir.ListDir(func(level2Dir fs.File) error {
				if !level2Dir.Exists() || level2Dir.IsHidden() {
					return nil
				}
				if !level2Dir.IsDir() {
					// fmt.Println("Directory expected but found file:", level2Dir)
					return nil
				}
				return level2Dir.ListDir(func(level3Dir fs.File) error {
					if !level3Dir.Exists() || level3Dir.IsHidden() {
						return nil
					}
					if !level3Dir.IsDir() {
						// fmt.Println("Directory expected but found file:", level3Dir)
						return nil
					}
					return level3Dir.ListDir(func(uuidDir fs.File) error {
						if !uuidDir.Exists() || uuidDir.IsHidden() {
							return nil
						}
						if !uuidDir.IsDir() {
							// fmt.Println("Directory expected but found file:", uuidDir)
							return nil
						}
						uuid, err := Parse(uuidDir)
						if err != nil {
							return err
						}
						return callback(uuidDir, uuid)
					})
				})
			})
		})
	})
}

// RemoveDir deletes uuidDir recursevely and all directories
// upward of uuidDir until but not including baseDir.
func RemoveDir(baseDir, uuidDir fs.File) error {
	basePath := baseDir.Path()
	uuidPath := uuidDir.Path()
	if !strings.HasPrefix(uuidPath, basePath) || baseDir.FileSystem() != uuidDir.FileSystem() {
		return errors.Errorf("uuidDir(%q) is not a sub directory of baseDir(%q)", uuidPath, basePath)
	}
	if uuidPath == basePath {
		return nil
	}

	// fmt.Println("deleting", uuidDir.Path())
	err := uuidDir.RemoveRecursive()
	if err != nil {
		return err
	}
	for {
		uuidDir = uuidDir.Dir()
		if uuidDir.Path() == basePath || !uuidDir.IsEmptyDir() {
			return nil
		}
		// fmt.Println("deleting", uuidDir.Path())
		err = uuidDir.Remove()
		if err != nil {
			return err
		}
	}
}

// Make sub-directories under baseDir for the passed UUID
func Make(baseDir fs.File, uuid [16]byte) (uuidDir fs.File, err error) {
	uuidDir = Join(baseDir, uuid)
	return uuidDir, baseDir.MakeAllDirs()
}

// Remove the sub-directories under baseDir for the passed UUID
func Remove(baseDir fs.File, uuid [16]byte) error {
	return RemoveDir(baseDir, Join(baseDir, uuid))
}
<commit_msg>use fmt.Errorf<commit_after>// Package uuiddirs provides functions to split up a UUID
// into a series of sub-directories so that an unlimited number
// of UUIDs can be used as directories.
//
// Example:
//   The UUID f0498fad-437c-4954-ad82-8ec2cc202628 maps to the path
//   f0/498/fad/437c4954/ad828ec2cc202628
package uuiddir

import (
	"encoding/hex"
	"fmt"
	"path"
	"strings"

	fs "github.com/ungerik/go-fs"
)

// Split a UUID into 5 hex strings.
// Example:
//   Splitting the UUID f0498fad-437c-4954-ad82-8ec2cc202628 returns
//   []string{"f0", "498", "fad", "437c4954", "ad828ec2cc202628"}
func Split(uuid [16]byte) []string {
	hexStr := hex.EncodeToString(uuid[:])
	return []string{
		hexStr[0:2],
		hexStr[2:5],
		hexStr[5:8],
		hexStr[8:16],
		hexStr[16:32],
	}
}

// Join returns a directory with the splitted UUID and pathParts joined to baseDir.
func Join(baseDir fs.File, uuid [16]byte, pathParts ...string) fs.File {
	return baseDir.Join(append(Split(uuid), pathParts...)...)
}

// Parse the path of uuidDir for a UUID
func Parse(uuidDir fs.File) (uuid [16]byte, err error) {
	uuidPath := strings.TrimSuffix(uuidDir.PathWithSlashes(), "/")
	if len(uuidPath) < 36 {
		return nilUUID, fmt.Errorf("path can't be parsed as UUID: %q", string(uuidDir))
	}
	return ParseString(uuidPath[len(uuidPath)-36:])
}

// FormatString returns the splitted UUID joined with slashes.
// It's the inverse to ParseString.
func FormatString(uuid [16]byte) string {
	return path.Join(Split(uuid)...)
}

// ParseString parses a 36 character string (like returned from FormatString) as UUID.
func ParseString(uuidPath string) (uuid [16]byte, err error) {
	if len(uuidPath) != 36 {
		return nilUUID, fmt.Errorf("path can't be parsed as UUID: %q", uuidPath)
	}
	uuidPath = strings.Replace(uuidPath, "/", "", 4)
	if len(uuidPath) != 32 {
		return nilUUID, fmt.Errorf("path can't be parsed as UUID: %q", uuidPath)
	}
	b, err := hex.DecodeString(uuidPath)
	if err != nil {
		return nilUUID, fmt.Errorf("path can't be parsed as UUID: %q", uuidPath)
	}
	copy(uuid[:], b)
	return uuid, validateUUID(uuid)
}

// Enum calls callback for every directory that represents an UUID under baseDir.
func Enum(baseDir fs.File, callback func(uuidDir fs.File, uuid [16]byte) error) error {
	return baseDir.ListDir(func(level0Dir fs.File) error {
		if !level0Dir.Exists() || level0Dir.IsHidden() {
			return nil
		}
		if !level0Dir.IsDir() {
			// fmt.Println("Directory expected but found file:", level0Dir)
			return nil
		}
		return level0Dir.ListDir(func(level1Dir fs.File) error {
			if !level1Dir.Exists() || level1Dir.IsHidden() {
				return nil
			}
			if !level1Dir.IsDir() {
				// fmt.Println("Directory expected but found file:", level1Dir)
				return nil
			}
			return level1Dir.ListDir(func(level2Dir fs.File) error {
				if !level2Dir.Exists() || level2Dir.IsHidden() {
					return nil
				}
				if !level2Dir.IsDir() {
					// fmt.Println("Directory expected but found file:", level2Dir)
					return nil
				}
				return level2Dir.ListDir(func(level3Dir fs.File) error {
					if !level3Dir.Exists() || level3Dir.IsHidden() {
						return nil
					}
					if !level3Dir.IsDir() {
						// fmt.Println("Directory expected but found file:", level3Dir)
						return nil
					}
					return level3Dir.ListDir(func(uuidDir fs.File) error {
						if !uuidDir.Exists() || uuidDir.IsHidden() {
							return nil
						}
						if !uuidDir.IsDir() {
							// fmt.Println("Directory expected but found file:", uuidDir)
							return nil
						}
						uuid, err := Parse(uuidDir)
						if err != nil {
							return err
						}
						return callback(uuidDir, uuid)
					})
				})
			})
		})
	})
}

// RemoveDir deletes uuidDir recursevely and all directories
// upward of uuidDir until but not including baseDir.
func RemoveDir(baseDir, uuidDir fs.File) error {
	basePath := baseDir.Path()
	uuidPath := uuidDir.Path()
	if !strings.HasPrefix(uuidPath, basePath) || baseDir.FileSystem() != uuidDir.FileSystem() {
		return fmt.Errorf("uuidDir(%q) is not a sub directory of baseDir(%q)", uuidPath, basePath)
	}
	if uuidPath == basePath {
		return nil
	}

	// fmt.Println("deleting", uuidDir.Path())
	err := uuidDir.RemoveRecursive()
	if err != nil {
		return err
	}
	for {
		uuidDir = uuidDir.Dir()
		if uuidDir.Path() == basePath || !uuidDir.IsEmptyDir() {
			return nil
		}
		// fmt.Println("deleting", uuidDir.Path())
		err = uuidDir.Remove()
		if err != nil {
			return err
		}
	}
}

// Make sub-directories under baseDir for the passed UUID
func Make(baseDir fs.File, uuid [16]byte) (uuidDir fs.File, err error) {
	uuidDir = Join(baseDir, uuid)
	return uuidDir, baseDir.MakeAllDirs()
}

// Remove the sub-directories under baseDir for the passed UUID
func Remove(baseDir fs.File, uuid [16]byte) error {
	return RemoveDir(baseDir, Join(baseDir, uuid))
}
<|endoftext|>


<commit_before>package messages

import (
	"fmt"
	"os"
	"testing"
)

func TestVerifyHistorianMessage(t *testing.T) {
	s := "alexandria-historian-v001:pool.alexandria.io:0.000136008500:316306445.6533333:nr:0.00000500:0.00217:IN9OrF1Kpd5S0x36nXWI0lFjhnS1Z9I9k7cxWJrFUlsfcgwJytZ+GlKP1/tHCijAdGAX6LnOgOtcvI/vMQgVcwA="

	if os.Getenv("F_USER") == "" {
		t.Skip("skipping test; $F_TOKEN not set")
	}
	if os.Getenv("F_TOKEN") == "" {
		t.Skip("skipping test; $F_TOKEN not set")
	}

	// Don't need heavy testing of true address validity
	// The heavy lifting is done by the FlorinCoin daemon
	cases := []struct {
		in  string
		err error
	}{
		{s, nil},
		{s[:len(s)-1] + "a", ErrHistorianMessageBadSignature},
	}

	for _, c := range cases {
		got, err := VerifyHistorianMessage([]byte(c.in))
		if err != c.err {
			t.Errorf("VerifyHistorianMessage(%q) | err == %q, want %q", c.in, err, c.err)
		}
		// ToDo: check the decoded result
		fmt.Printf("%v\n", got)
		//if got != c.out {
		//	t.Errorf("CheckAddress(%q) == %q, want %q", c.in, got, c.out)
		//}
	}
}
<commit_msg>Additional tests<commit_after>package messages

import (
	"fmt"
	"os"
	"strings"
	"testing"
)

func TestVerifyHistorianMessage(t *testing.T) {
	s := "alexandria-historian-v001:pool.alexandria.io:0.000136008500:316306445.6533333:nr:0.00000500:0.00217:IN9OrF1Kpd5S0x36nXWI0lFjhnS1Z9I9k7cxWJrFUlsfcgwJytZ+GlKP1/tHCijAdGAX6LnOgOtcvI/vMQgVcwA="

	if os.Getenv("F_USER") == "" {
		t.Skip("skipping test; $F_TOKEN not set")
	}
	if os.Getenv("F_TOKEN") == "" {
		t.Skip("skipping test; $F_TOKEN not set")
	}

	// Don't need heavy testing of true address validity
	// The heavy lifting is done by the FlorinCoin daemon
	cases := []struct {
		in  string
		err error
	}{
		{s, nil},
		{s[:len(s)-1] + "a", ErrHistorianMessageBadSignature},
		{s[:25], ErrHistorianMessageInvalid},
		{strings.Replace(s, "v001", "v002", 1), ErrHistorianMessageInvalid},
		{strings.Replace(s, "pool.", "notpool.", 1), ErrHistorianMessageInvalid},
	}

	for _, c := range cases {
		got, err := VerifyHistorianMessage([]byte(c.in))
		if err != c.err {
			t.Errorf("VerifyHistorianMessage(%q) | err == %q, want %q", c.in, err, c.err)
		}
		// ToDo: check the decoded result
		fmt.Printf("%v\n", got)
		//if got != c.out {
		//	t.Errorf("CheckAddress(%q) == %q, want %q", c.in, got, c.out)
		//}
	}
}
<|endoftext|>


<commit_before>package h2spec

import (
	"fmt"
	"github.com/bradfitz/http2"
	"io"
	"net"
	"syscall"
)

func WindowUpdateTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9", "WINDOW_UPDATE")

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with an flow control window increment of 0",
		"The endpoint MUST respond with a connection error of type PROTOCOL_ERROR.",
		func(ctx *Context) (expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteWindowUpdate(0, 0)

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with an flow control window increment of 0 on a stream",
		"The endpoint MUST respond with a stream error of type PROTOCOL_ERROR.",
		func(ctx *Context) (expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)
			http2Conn.fr.WriteWindowUpdate(1, 0)

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestStreamError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with a length other than a multiple of 4 octets",
		"The endpoint MUST respond with a connection error of type FRAME_SIZE_ERROR.",
		func(ctx *Context) (expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			fmt.Fprintf(http2Conn.conn, "\x00\x00\x03\x08\x00\x00\x00\x00\x00")
			fmt.Fprintf(http2Conn.conn, "\x00\x00\x01")

			actualCodes := []http2.ErrCode{http2.ErrCodeFrameSize}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestGroup(TheFlowControlWindowTestGroup(ctx))
	tg.AddTestGroup(InitialFlowControlWindowSizeTestGroup(ctx))

	return tg
}

func TheFlowControlWindowTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9.1", "The Flow Control Window")

	tg.AddTestCase(NewTestCase(
		"Sends multiple WINDOW_UPDATE frames on a connection increasing the flow control window to above 2^31-1",
		"The endpoint MUST sends a GOAWAY frame with a FLOW_CONTROL_ERROR code.",
		func(ctx *Context) (expected []Result, actual Result) {
			expected = []Result{
				&ResultFrame{http2.FrameGoAway, FlagDefault, http2.ErrCodeFlowControl},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteWindowUpdate(0, 2147483647)
			http2Conn.fr.WriteWindowUpdate(0, 2147483647)

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						actual = &ResultConnectionClose{}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.GoAwayFrame:
					actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
					if f.ErrCode == http2.ErrCodeFlowControl {
						break loop
					}
				default:
					actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
				}
			}

			return expected, actual
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends multiple WINDOW_UPDATE frames on a stream increasing the flow control window to above 2^31-1",
		"The endpoint MUST sends a RST_STREAM with the error code of FLOW_CONTROL_ERROR code.",
		func(ctx *Context) (expected []Result, actual Result) {
			expected = []Result{
				&ResultFrame{http2.FrameRSTStream, FlagDefault, http2.ErrCodeFlowControl},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			http2Conn.fr.WriteWindowUpdate(1, 2147483647)
			http2Conn.fr.WriteWindowUpdate(1, 2147483647)

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						actual = &ResultConnectionClose{}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.RSTStreamFrame:
					actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
					if f.ErrCode == http2.ErrCodeFlowControl {
						break loop
					}
				default:
					actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
				}
			}

			return expected, actual
		},
	))

	return tg
}

func InitialFlowControlWindowSizeTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9.2", "Initial Flow Control Window Size")

	tg.AddTestCase(NewTestCase(
		"Sends a SETTINGS_INITIAL_WINDOW_SIZE settings with an exceeded maximum window size value",
		"The endpoint MUST respond with a connection error of type FLOW_CONTROL_ERROR.",
		func(ctx *Context) (expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			fmt.Fprintf(http2Conn.conn, "\x00\x00\x06\x04\x00\x00\x00\x00\x00")
			fmt.Fprintf(http2Conn.conn, "\x00\x04\x80\x00\x00\x00")

			actualCodes := []http2.ErrCode{http2.ErrCodeFlowControl}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	return tg
}
<commit_msg>Add test case for basic window control<commit_after>package h2spec

import (
	"errors"
	"fmt"
	"github.com/bradfitz/http2"
	"io"
	"net"
	"syscall"
)

func WindowUpdateTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9", "WINDOW_UPDATE")

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame",
		"The endpoint is expected to send the DATA frame based on the window size.",
		func(ctx *Context) (expected []Result, actual Result) {
			expected = []Result{
				&ResultFrame{http2.FrameData, FlagDefault, ErrCodeDefault},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			settings := http2.Setting{http2.SettingInitialWindowSize, 1}
			http2Conn.fr.WriteSettings(settings)

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = true
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			winUpdated := false

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						actual = &ResultConnectionClose{}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.DataFrame:

					if winUpdated {
						if f.FrameHeader.Length > 10 {
							err := errors.New("The length of DATA frame is invalid.")
							actual = &ResultError{err}
							break loop
						}

						actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
						break loop
					} else {
						if f.FrameHeader.Length != 1 {
							err := errors.New("The length of DATA frame is invalid.")
							actual = &ResultError{err}
							break loop
						}

						http2Conn.fr.WriteWindowUpdate(1, 10)
						winUpdated = true
					}
				default:
					actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
				}
			}

			return expected, actual
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with an flow control window increment of 0",
		"The endpoint MUST respond with a connection error of type PROTOCOL_ERROR.",
		func(ctx *Context) (expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteWindowUpdate(0, 0)

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with an flow control window increment of 0 on a stream",
		"The endpoint MUST respond with a stream error of type PROTOCOL_ERROR.",
		func(ctx *Context) (expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)
			http2Conn.fr.WriteWindowUpdate(1, 0)

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestStreamError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with a length other than a multiple of 4 octets",
		"The endpoint MUST respond with a connection error of type FRAME_SIZE_ERROR.",
		func(ctx *Context) (expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			fmt.Fprintf(http2Conn.conn, "\x00\x00\x03\x08\x00\x00\x00\x00\x00")
			fmt.Fprintf(http2Conn.conn, "\x00\x00\x01")

			actualCodes := []http2.ErrCode{http2.ErrCodeFrameSize}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestGroup(TheFlowControlWindowTestGroup(ctx))
	tg.AddTestGroup(InitialFlowControlWindowSizeTestGroup(ctx))

	return tg
}

func TheFlowControlWindowTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9.1", "The Flow Control Window")

	tg.AddTestCase(NewTestCase(
		"Sends multiple WINDOW_UPDATE frames on a connection increasing the flow control window to above 2^31-1",
		"The endpoint MUST sends a GOAWAY frame with a FLOW_CONTROL_ERROR code.",
		func(ctx *Context) (expected []Result, actual Result) {
			expected = []Result{
				&ResultFrame{http2.FrameGoAway, FlagDefault, http2.ErrCodeFlowControl},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteWindowUpdate(0, 2147483647)
			http2Conn.fr.WriteWindowUpdate(0, 2147483647)

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						actual = &ResultConnectionClose{}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.GoAwayFrame:
					actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
					if f.ErrCode == http2.ErrCodeFlowControl {
						break loop
					}
				default:
					actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
				}
			}

			return expected, actual
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends multiple WINDOW_UPDATE frames on a stream increasing the flow control window to above 2^31-1",
		"The endpoint MUST sends a RST_STREAM with the error code of FLOW_CONTROL_ERROR code.",
		func(ctx *Context) (expected []Result, actual Result) {
			expected = []Result{
				&ResultFrame{http2.FrameRSTStream, FlagDefault, http2.ErrCodeFlowControl},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			http2Conn.fr.WriteWindowUpdate(1, 2147483647)
			http2Conn.fr.WriteWindowUpdate(1, 2147483647)

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						actual = &ResultConnectionClose{}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.RSTStreamFrame:
					actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
					if f.ErrCode == http2.ErrCodeFlowControl {
						break loop
					}
				default:
					actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
				}
			}

			return expected, actual
		},
	))

	return tg
}

func InitialFlowControlWindowSizeTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9.2", "Initial Flow Control Window Size")

	tg.AddTestCase(NewTestCase(
		"Sends a SETTINGS_INITIAL_WINDOW_SIZE settings with an exceeded maximum window size value",
		"The endpoint MUST respond with a connection error of type FLOW_CONTROL_ERROR.",
		func(ctx *Context) (expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			fmt.Fprintf(http2Conn.conn, "\x00\x00\x06\x04\x00\x00\x00\x00\x00")
			fmt.Fprintf(http2Conn.conn, "\x00\x04\x80\x00\x00\x00")

			actualCodes := []http2.ErrCode{http2.ErrCodeFlowControl}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	return tg
}
<|endoftext|>


<commit_before>package controller

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/appscode/log"
	"github.com/graymeta/stow"
	_ "github.com/graymeta/stow/google"
	_ "github.com/graymeta/stow/s3"
	tapi "github.com/k8sdb/apimachinery/api"
	kapi "k8s.io/kubernetes/pkg/api"
	k8serr "k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/unversioned"
	kapps "k8s.io/kubernetes/pkg/apis/apps"
	"k8s.io/kubernetes/pkg/apis/batch"
	"k8s.io/kubernetes/pkg/apis/extensions"
	"k8s.io/kubernetes/pkg/labels"
)

func (w *Controller) EnsureDatabaseSnapshot() {
	resourceName := tapi.ResourceNameDatabaseSnapshot + "." + tapi.V1beta1SchemeGroupVersion.Group

	if _, err := w.Client.Extensions().ThirdPartyResources().Get(resourceName); err != nil {
		if !k8serr.IsNotFound(err) {
			log.Fatalln(err)
		}
	} else {
		return
	}

	thirdPartyResource := &extensions.ThirdPartyResource{
		TypeMeta: unversioned.TypeMeta{
			APIVersion: "extensions/v1beta1",
			Kind:       "ThirdPartyResource",
		},
		ObjectMeta: kapi.ObjectMeta{
			Name: resourceName,
		},
		Versions: []extensions.APIVersion{
			{
				Name: tapi.V1beta1SchemeGroupVersion.Version,
			},
		},
	}

	if _, err := w.Client.Extensions().ThirdPartyResources().Create(thirdPartyResource); err != nil {
		log.Fatalln(err)
	}
}

func (w *Controller) EnsureDeletedDatabase() {
	resourceName := tapi.ResourceNameDeletedDatabase + "." + tapi.V1beta1SchemeGroupVersion.Group

	if _, err := w.Client.Extensions().ThirdPartyResources().Get(resourceName); err != nil {
		if !k8serr.IsNotFound(err) {
			log.Fatalln(err)
		}
	} else {
		return
	}

	thirdPartyResource := &extensions.ThirdPartyResource{
		TypeMeta: unversioned.TypeMeta{
			APIVersion: "extensions/v1beta1",
			Kind:       "ThirdPartyResource",
		},
		ObjectMeta: kapi.ObjectMeta{
			Name: resourceName,
		},
		Versions: []extensions.APIVersion{
			{
				Name: tapi.V1beta1SchemeGroupVersion.Version,
			},
		},
	}

	if _, err := w.Client.Extensions().ThirdPartyResources().Create(thirdPartyResource); err != nil {
		log.Fatalln(err)
	}
}

const (
	LabelSnapshotActive = "elastic.k8sdb.com/status"
)

func (w *Controller) CheckDatabaseSnapshotJob(snapshot *tapi.DatabaseSnapshot, jobName string, checkTime float64) {

	unversionedNow := unversioned.Now()
	snapshot.Status.StartTime = &unversionedNow
	snapshot.Status.Status = tapi.SnapshotRunning

	snapshot.Labels[LabelSnapshotActive] = string(tapi.SnapshotRunning)
	var err error
	if snapshot, err = w.ExtClient.DatabaseSnapshot(snapshot.Namespace).Update(snapshot); err != nil {
		log.Errorln(err)
	}

	var jobSuccess bool = false
	var job *batch.Job

	then := time.Now()
	now := time.Now()
	for now.Sub(then).Minutes() < checkTime {
		log.Debugln("Checking for Job ", jobName)
		job, err = w.Client.Batch().Jobs(snapshot.Namespace).Get(jobName)
		if err != nil {
			break
		}
		log.Debugf("Pods Statuses:	%d Running / %d Succeeded / %d Failed",
			job.Status.Active, job.Status.Succeeded, job.Status.Failed)
		// If job is success
		if job.Status.Succeeded > 0 {
			jobSuccess = true
			break
		} else if job.Status.Failed > 0 {
			break
		}

		time.Sleep(time.Minute)
		now = time.Now()
	}

	podList, err := w.Client.Core().Pods(job.Namespace).List(
		kapi.ListOptions{
			LabelSelector: labels.SelectorFromSet(job.Spec.Selector.MatchLabels),
		},
	)
	if err != nil {
		log.Errorln(err)
		return
	}

	for _, pod := range podList.Items {
		if err := w.Client.Core().Pods(pod.Namespace).Delete(pod.Name, nil); err != nil {
			log.Errorln(err)
		}
	}

	for _, volume := range job.Spec.Template.Spec.Volumes {
		claim := volume.PersistentVolumeClaim
		if claim != nil {
			err := w.Client.Core().PersistentVolumeClaims(job.Namespace).Delete(claim.ClaimName, nil)
			if err != nil {
				log.Errorln(err)
			}
		}
	}

	if err := w.Client.Batch().Jobs(job.Namespace).Delete(job.Name, nil); err != nil {
		log.Errorln(err)
	}

	if snapshot, err = w.ExtClient.DatabaseSnapshot(snapshot.Namespace).Get(snapshot.Name); err != nil {
		log.Errorln(err)
		return
	}

	unversionedNow = unversioned.Now()
	snapshot.Status.CompletionTime = &unversionedNow
	if jobSuccess {
		snapshot.Status.Status = tapi.SnapshotSuccessed
	} else {
		snapshot.Status.Status = tapi.SnapshotFailed
	}

	delete(snapshot.Labels, LabelSnapshotActive)

	if _, err := w.ExtClient.DatabaseSnapshot(snapshot.Namespace).Update(snapshot); err != nil {
		log.Errorln(err)
	}
}

func (w *Controller) CheckStatefulSets(statefulSet *kapps.StatefulSet, checkTime float64) error {
	podName := fmt.Sprintf("%v-%v", statefulSet.Name, 0)

	podReady := false
	then := time.Now()
	now := time.Now()
	for now.Sub(then).Minutes() < checkTime {
		pod, err := w.Client.Core().Pods(statefulSet.Namespace).Get(podName)
		if err != nil {
			if k8serr.IsNotFound(err) {
				time.Sleep(time.Second * 10)
				now = time.Now()
				continue
			} else {
				return err
			}
		}
		log.Debugf("Pod Phase: %v", pod.Status.Phase)

		// If job is success
		if pod.Status.Phase == kapi.PodRunning {
			podReady = true
			break
		}

		time.Sleep(time.Minute)
		now = time.Now()
	}
	if !podReady {
		return errors.New("Database fails to be Ready")
	}
	return nil
}

func (w *Controller) GetVolumeForSnapshot(storage *tapi.StorageSpec, jobName, namespace string) (*kapi.Volume, error) {
	volume := &kapi.Volume{
		Name: "util-volume",
	}
	if storage != nil {
		claim := &kapi.PersistentVolumeClaim{
			ObjectMeta: kapi.ObjectMeta{
				Name:      jobName,
				Namespace: namespace,
				Annotations: map[string]string{
					"volume.beta.kubernetes.io/storage-class": storage.Class,
				},
			},
			Spec: storage.PersistentVolumeClaimSpec,
		}

		if _, err := w.Client.Core().PersistentVolumeClaims(claim.Namespace).Create(claim); err != nil {
			return nil, err
		}

		volume.PersistentVolumeClaim = &kapi.PersistentVolumeClaimVolumeSource{
			ClaimName: claim.Name,
		}
	} else {
		volume.EmptyDir = &kapi.EmptyDirVolumeSource{}
	}
	return volume, nil
}

const (
	keyProvider = "provider"
	keyConfig   = "config"
)

func (w *Controller) CheckBucketAccess(bucketName, secretName, namespace string) error {
	secret, err := w.Client.Core().Secrets(namespace).Get(secretName)
	if err != nil {
		return err
	}

	provider := secret.Data[keyProvider]
	if provider == nil {
		return errors.New("Missing provider key")
	}
	configData := secret.Data[keyConfig]
	if configData == nil {
		return errors.New("Missing config key")
	}

	var config stow.ConfigMap
	if err := json.Unmarshal(configData, &config); err != nil {
		return errors.New("Fail to Unmarshal config data")
	}

	loc, err := stow.Dial(string(provider), config)
	if err != nil {
		return err
	}

	container, err := loc.Container(bucketName)
	if err != nil {
		return err
	}

	r := bytes.NewReader([]byte("CheckBucketAccess"))
	item, err := container.Put(".k8sdb", r, r.Size(), nil)
	if err != nil {
		return err
	}

	if err := container.RemoveItem(item.ID()); err != nil {
		return err
	}
	return nil
}
<commit_msg>Use time.Duration instead of float64 (#9)<commit_after>package controller

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/appscode/log"
	"github.com/graymeta/stow"
	_ "github.com/graymeta/stow/google"
	_ "github.com/graymeta/stow/s3"
	tapi "github.com/k8sdb/apimachinery/api"
	kapi "k8s.io/kubernetes/pkg/api"
	k8serr "k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/unversioned"
	kapps "k8s.io/kubernetes/pkg/apis/apps"
	"k8s.io/kubernetes/pkg/apis/batch"
	"k8s.io/kubernetes/pkg/apis/extensions"
	"k8s.io/kubernetes/pkg/labels"
)

func (w *Controller) EnsureDatabaseSnapshot() {
	resourceName := tapi.ResourceNameDatabaseSnapshot + "." + tapi.V1beta1SchemeGroupVersion.Group

	if _, err := w.Client.Extensions().ThirdPartyResources().Get(resourceName); err != nil {
		if !k8serr.IsNotFound(err) {
			log.Fatalln(err)
		}
	} else {
		return
	}

	thirdPartyResource := &extensions.ThirdPartyResource{
		TypeMeta: unversioned.TypeMeta{
			APIVersion: "extensions/v1beta1",
			Kind:       "ThirdPartyResource",
		},
		ObjectMeta: kapi.ObjectMeta{
			Name: resourceName,
		},
		Versions: []extensions.APIVersion{
			{
				Name: tapi.V1beta1SchemeGroupVersion.Version,
			},
		},
	}

	if _, err := w.Client.Extensions().ThirdPartyResources().Create(thirdPartyResource); err != nil {
		log.Fatalln(err)
	}
}

func (w *Controller) EnsureDeletedDatabase() {
	resourceName := tapi.ResourceNameDeletedDatabase + "." + tapi.V1beta1SchemeGroupVersion.Group

	if _, err := w.Client.Extensions().ThirdPartyResources().Get(resourceName); err != nil {
		if !k8serr.IsNotFound(err) {
			log.Fatalln(err)
		}
	} else {
		return
	}

	thirdPartyResource := &extensions.ThirdPartyResource{
		TypeMeta: unversioned.TypeMeta{
			APIVersion: "extensions/v1beta1",
			Kind:       "ThirdPartyResource",
		},
		ObjectMeta: kapi.ObjectMeta{
			Name: resourceName,
		},
		Versions: []extensions.APIVersion{
			{
				Name: tapi.V1beta1SchemeGroupVersion.Version,
			},
		},
	}

	if _, err := w.Client.Extensions().ThirdPartyResources().Create(thirdPartyResource); err != nil {
		log.Fatalln(err)
	}
}

const (
	LabelSnapshotActive = "elastic.k8sdb.com/status"
)

func (w *Controller) CheckDatabaseSnapshotJob(snapshot *tapi.DatabaseSnapshot, jobName string, checkDuration time.Duration) {

	unversionedNow := unversioned.Now()
	snapshot.Status.StartTime = &unversionedNow
	snapshot.Status.Status = tapi.SnapshotRunning
	snapshot.Labels[LabelSnapshotActive] = string(tapi.SnapshotRunning)
	var err error
	if snapshot, err = w.ExtClient.DatabaseSnapshot(snapshot.Namespace).Update(snapshot); err != nil {
		log.Errorln(err)
	}

	var jobSuccess bool = false
	var job *batch.Job

	then := time.Now()
	now := time.Now()
	for now.Sub(then) < checkDuration {
		log.Debugln("Checking for Job ", jobName)
		job, err = w.Client.Batch().Jobs(snapshot.Namespace).Get(jobName)
		if err != nil {
			break
		}
		log.Debugf("Pods Statuses:	%d Running / %d Succeeded / %d Failed",
			job.Status.Active, job.Status.Succeeded, job.Status.Failed)
		// If job is success
		if job.Status.Succeeded > 0 {
			jobSuccess = true
			break
		} else if job.Status.Failed > 0 {
			break
		}

		time.Sleep(time.Minute)
		now = time.Now()
	}

	podList, err := w.Client.Core().Pods(job.Namespace).List(
		kapi.ListOptions{
			LabelSelector: labels.SelectorFromSet(job.Spec.Selector.MatchLabels),
		},
	)
	if err != nil {
		log.Errorln(err)
		return
	}

	for _, pod := range podList.Items {
		if err := w.Client.Core().Pods(pod.Namespace).Delete(pod.Name, nil); err != nil {
			log.Errorln(err)
		}
	}

	for _, volume := range job.Spec.Template.Spec.Volumes {
		claim := volume.PersistentVolumeClaim
		if claim != nil {
			err := w.Client.Core().PersistentVolumeClaims(job.Namespace).Delete(claim.ClaimName, nil)
			if err != nil {
				log.Errorln(err)
			}
		}
	}

	if err := w.Client.Batch().Jobs(job.Namespace).Delete(job.Name, nil); err != nil {
		log.Errorln(err)
	}

	if snapshot, err = w.ExtClient.DatabaseSnapshot(snapshot.Namespace).Get(snapshot.Name); err != nil {
		log.Errorln(err)
		return
	}

	unversionedNow = unversioned.Now()
	snapshot.Status.CompletionTime = &unversionedNow
	if jobSuccess {
		snapshot.Status.Status = tapi.SnapshotSuccessed
	} else {
		snapshot.Status.Status = tapi.SnapshotFailed
	}

	delete(snapshot.Labels, LabelSnapshotActive)

	if _, err := w.ExtClient.DatabaseSnapshot(snapshot.Namespace).Update(snapshot); err != nil {
		log.Errorln(err)
	}
}

func (w *Controller) CheckStatefulSets(statefulSet *kapps.StatefulSet, checkDuration time.Duration) error {
	podName := fmt.Sprintf("%v-%v", statefulSet.Name, 0)

	podReady := false
	then := time.Now()
	now := time.Now()
	for now.Sub(then) < checkDuration {
		pod, err := w.Client.Core().Pods(statefulSet.Namespace).Get(podName)
		if err != nil {
			if k8serr.IsNotFound(err) {
				time.Sleep(time.Second * 10)
				now = time.Now()
				continue
			} else {
				return err
			}
		}
		log.Debugf("Pod Phase: %v", pod.Status.Phase)

		// If job is success
		if pod.Status.Phase == kapi.PodRunning {
			podReady = true
			break
		}

		time.Sleep(time.Minute)
		now = time.Now()
	}
	if !podReady {
		return errors.New("Database fails to be Ready")
	}
	return nil
}

func (w *Controller) GetVolumeForSnapshot(storage *tapi.StorageSpec, jobName, namespace string) (*kapi.Volume, error) {
	volume := &kapi.Volume{
		Name: "util-volume",
	}
	if storage != nil {
		claim := &kapi.PersistentVolumeClaim{
			ObjectMeta: kapi.ObjectMeta{
				Name:      jobName,
				Namespace: namespace,
				Annotations: map[string]string{
					"volume.beta.kubernetes.io/storage-class": storage.Class,
				},
			},
			Spec: storage.PersistentVolumeClaimSpec,
		}

		if _, err := w.Client.Core().PersistentVolumeClaims(claim.Namespace).Create(claim); err != nil {
			return nil, err
		}

		volume.PersistentVolumeClaim = &kapi.PersistentVolumeClaimVolumeSource{
			ClaimName: claim.Name,
		}
	} else {
		volume.EmptyDir = &kapi.EmptyDirVolumeSource{}
	}
	return volume, nil
}

const (
	keyProvider = "provider"
	keyConfig   = "config"
)

func (w *Controller) CheckBucketAccess(bucketName, secretName, namespace string) error {
	secret, err := w.Client.Core().Secrets(namespace).Get(secretName)
	if err != nil {
		return err
	}

	provider := secret.Data[keyProvider]
	if provider == nil {
		return errors.New("Missing provider key")
	}
	configData := secret.Data[keyConfig]
	if configData == nil {
		return errors.New("Missing config key")
	}

	var config stow.ConfigMap
	if err := json.Unmarshal(configData, &config); err != nil {
		return errors.New("Fail to Unmarshal config data")
	}

	loc, err := stow.Dial(string(provider), config)
	if err != nil {
		return err
	}

	container, err := loc.Container(bucketName)
	if err != nil {
		return err
	}

	r := bytes.NewReader([]byte("CheckBucketAccess"))
	item, err := container.Put(".k8sdb", r, r.Size(), nil)
	if err != nil {
		return err
	}

	if err := container.RemoveItem(item.ID()); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package pipe

import "os"

type Pipe struct {
}

func NewPipe() *Pipe {
	return NewPipeSize(defaultMemBufferSize)
}

func NewPipeSize(size int) *Pipe {
	return newPipe(newMemBufferSize(size))
}

func NewPipeFile(file *os.File, size int) *Pipe {
	return newPipe(newFileBufferSize(file, size))
}

func newPipe(store Buffer) *Pipe {
	panic("TODO")
}

func (p *Pipe) Close() {
	p.CloseReader(nil)
	p.CloseWriter(nil)
}

func (p *Pipe) Reader() Reader {
	return &PipeReader{p}
}

func (p *Pipe) Read(b []byte) (int, error) {
	panic("TODO")
}

func (p *Pipe) Buffered() (int, error) {
	panic("TODO")
}

func (p *Pipe) CloseReader(err error) error {
	panic("TODO")
}

func (p *Pipe) Writer() Writer {
	return &PipeWriter{p}
}

func (p *Pipe) Write(b []byte) (int, error) {
	panic("TODO")
}

func (p *Pipe) Available() (int, error) {
	panic("TODO")
}

func (p *Pipe) CloseWriter(err error) error {
	panic("TODO")
}
<commit_msg>pipe: update pipe.go<commit_after>package pipe

import (
	"io"
	"os"
	"sync"

	"github.com/CodisLabs/codis/pkg/utils/errors"
)

type Pipe struct {
	r, w struct {
		sync.Mutex
		cond *sync.Cond
	}
	mu sync.Mutex

	rerr, werr error

	store Buffer
}

func NewPipe() *Pipe {
	return NewPipeSize(defaultMemBufferSize)
}

func NewPipeSize(size int) *Pipe {
	return newPipe(newMemBufferSize(size))
}

func NewPipeFile(file *os.File, size int) *Pipe {
	return newPipe(newFileBufferSize(file, size))
}

func newPipe(store Buffer) *Pipe {
	p := &Pipe{store: store}
	p.r.cond = sync.NewCond(&p.mu)
	p.w.cond = sync.NewCond(&p.mu)
	return p
}

func (p *Pipe) Close() {
	p.CloseReader(nil)
	p.CloseWriter(nil)
}

func (p *Pipe) Reader() Reader {
	return &PipeReader{p}
}

func (p *Pipe) Read(b []byte) (int, error) {
	p.r.Lock()
	defer p.r.Unlock()
	for {
		n, err := p.readSome(b)
		if err != nil || n != 0 {
			return n, err
		}
		if len(b) == 0 {
			return 0, nil
		}
	}
}

func (p *Pipe) readSome(b []byte) (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.rerr != nil {
		return 0, errors.Trace(io.ErrClosedPipe)
	}
	if len(b) == 0 {
		if p.store.Buffered() != 0 {
			return 0, nil
		}
		return 0, p.werr
	}
	n, err := p.store.ReadSome(b)
	if err != nil || n != 0 {
		p.w.cond.Signal()
		return n, err
	}
	if p.werr != nil {
		return 0, p.werr
	} else {
		p.r.cond.Wait()
		return 0, nil
	}
}

func (p *Pipe) Buffered() (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.rerr != nil {
		return 0, p.rerr
	}
	if n := p.store.Buffered(); n != 0 {
		return n, nil
	} else {
		return 0, p.werr
	}
}

func (p *Pipe) CloseReader(err error) error {
	if err == nil {
		err = errors.Trace(io.ErrClosedPipe)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.rerr == nil {
		p.rerr = err
	}
	p.r.cond.Broadcast()
	p.w.cond.Broadcast()
	return p.store.CloseReader()
}

func (p *Pipe) Writer() Writer {
	return &PipeWriter{p}
}

func (p *Pipe) Write(b []byte) (int, error) {
	p.w.Lock()
	defer p.w.Unlock()
	var nn int
	for {
		n, err := p.writeSome(b)
		if err != nil || n == len(b) {
			return nn + n, err
		}
		nn, b = nn+n, b[n:]
	}
}

func (p *Pipe) writeSome(b []byte) (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.werr != nil {
		return 0, errors.Trace(io.ErrClosedPipe)
	}
	if p.rerr != nil {
		return 0, p.rerr
	}
	if len(b) == 0 {
		return 0, nil
	}
	n, err := p.store.WriteSome(b)
	if err != nil || n != 0 {
		p.r.cond.Signal()
		return n, err
	} else {
		p.w.cond.Wait()
		return 0, nil
	}
}

func (p *Pipe) Available() (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.werr != nil {
		return 0, p.werr
	}
	if p.rerr != nil {
		return 0, p.rerr
	}
	return p.store.Available(), nil
}

func (p *Pipe) CloseWriter(err error) error {
	if err == nil {
		err = errors.Trace(io.EOF)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.werr == nil {
		p.werr = err
	}
	p.r.cond.Broadcast()
	p.w.cond.Broadcast()
	return p.store.CloseWriter()
}
<|endoftext|>


<commit_before>package plugin

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/sgt-kabukiman/kabukibot/bot"
	"github.com/sgt-kabukiman/srapi"
)

type speedruncomConfig struct {
	Interval int
	Mapping  map[string]map[string]string
}

type SpeedrunComPlugin struct {
	config speedruncomConfig
	dict   *bot.Dictionary
}

func NewSpeedrunComPlugin() *SpeedrunComPlugin {
	return &SpeedrunComPlugin{}
}

func (self *SpeedrunComPlugin) Name() string {
	return "speedruncom"
}

func (self *SpeedrunComPlugin) Permissions() []string {
	return []string{"use_speedruncom_commands"}
}

func (self *SpeedrunComPlugin) Setup(bot *bot.Kabukibot) {
	self.config = speedruncomConfig{}
	self.dict = bot.Dictionary()

	err := bot.Configuration().PluginConfig(self.Name(), &self.config)
	if err != nil {
		bot.Logger().Warn("Could not load 'speedruncom' plugin configuration: %s", err)
	}

	go self.updater()
}

func (self *SpeedrunComPlugin) updater() {
	interval := time.Duration(self.config.Interval * int(time.Minute))

	for {
		for gameID, catList := range self.config.Mapping {
			game, err := srapi.GameByID(gameID, srapi.NoEmbeds)
			if err != nil {
				continue
			}

			leaderboards, err := game.Records(nil, "players,regions,platforms,category")
			if err != nil {
				continue
			}

			for _, leaderboard := range leaderboards.Data {
				if len(leaderboard.Runs) == 0 {
					continue
				}

				cat, err := leaderboard.Category(srapi.NoEmbeds)
				if err != nil {
					continue
				}

				dictKey, okay := catList[cat.ID]
				if !okay {
					continue
				}

				wr := leaderboard.Runs[0]
				formatted := formatWorldRecord(&wr.Run, game, cat, nil, nil, nil)

				self.dict.Set(dictKey, formatted)
			}
		}

		time.Sleep(interval)
	}
}

func (self *SpeedrunComPlugin) CreateWorker(channel bot.Channel) bot.PluginWorker {
	return &speedruncomWorker{
		channel: channel.Name(),
		acl:     channel.ACL(),
	}
}

type speedruncomWorker struct {
	channel string
	acl     *bot.ACL
}

func (self *speedruncomWorker) Enable() {
	// nothing to do for us
}

func (self *speedruncomWorker) Disable() {
	// nothing to do for us
}

func (self *speedruncomWorker) Part() {
	// nothing to do for us
}

func (self *speedruncomWorker) Shutdown() {
	// nothing to do for us
}

func (self *speedruncomWorker) HandleTextMessage(msg *bot.TextMessage, sender bot.Sender) {
	if msg.IsProcessed() || msg.IsFromBot() {
		return
	}

	if msg.IsCommand("wr") {
		self.handleWorldRecordCommand(msg, sender)
		msg.SetProcessed()
	}
}

var cleanerRegexp = regexp.MustCompile(`[^a-zA-Z0-9]`)

func (self *speedruncomWorker) handleWorldRecordCommand(msg *bot.TextMessage, sender bot.Sender) {
	if !self.acl.IsAllowed(msg.User, "use_speedrun_commands") {
		return
	}

	args := msg.Arguments()

	if len(args) == 0 {
		sender.Respond("you have to give me a game abbreviation.")
		return
	}

	gameIdentifier := args[0]

	var category *srapi.Category

	// try to find the game
	game, err := srapi.GameByAbbreviation(gameIdentifier, "categories")
	if err != nil {
		sender.Respond("I could not find a game with the abbreviation \"" + gameIdentifier + "\".")
		return
	}

	// assume all further args form the category, like "All Missions" or "Any%";
	// we normalise the value to make it -- hopefully -- easier to find the correct category

	if len(args) > 1 {
		catIdentifier := cleanerRegexp.ReplaceAllString(strings.ToLower(strings.Join(args[1:], "")), "")

		categories, err := game.Categories(nil, nil, srapi.NoEmbeds)
		catNames := []string{}

		if err == nil {
			for _, cat := range categories {
				id := cleanerRegexp.ReplaceAllString(strings.ToLower(cat.Name), "")

				if cat.Type == "per-game" {
					catNames = append(catNames, cat.Name)
				}

				if id == catIdentifier {
					category = cat
				}
			}
		}

		if category == nil {
			sender.Respond("I could not find a category named \"" + strings.Join(args[1:], " ") + "\". Available categories are: " + bot.HumanJoin(catNames, ", "))
			return
		} else if category.Type != "per-game" {
			sender.Respond(category.Name + " is a IL category; cannot report records for now. Sorry.")
			return
		}
	}

	var lb *srapi.Leaderboard

	// fetch the leaderboard, if possible (only available for games with full-game categories by default)
	if category == nil {
		lb, err = game.PrimaryLeaderboard(&srapi.LeaderboardOptions{Top: 1}, "players,platforms,regions,category")
		if err != nil || lb == nil {
			sender.Respond(game.Names.International + " does not use full-game categories by default, so I don't know what category or level you are referring to.")
			return
		}

		category, err = lb.Category(srapi.NoEmbeds)
		if err != nil {
			sender.Respond("the data from speedrun.com is invalid, cannot procede. Sorry. Try again later or a with different game.")
			return
		}
	} else {
		lb, err = category.PrimaryLeaderboard(&srapi.LeaderboardOptions{Top: 1}, "players,platforms,regions")
		if err != nil || lb == nil {
			sender.Respond(game.Names.International + " does not have runs for its \"" + category.Name + "\" category.")
			return
		}
	}

	// the leaderboard could be empty
	if len(lb.Runs) == 0 {
		sender.Respond(game.Names.International + ": " + category.Name + " does not have any matching runs yet.")
		return
	}

	// show only the first WR
	firstRun := lb.Runs[0].Run
	formatted := formatWorldRecord(&firstRun, game, category, nil, nil, nil)

	sender.SendText(formatted)
}

func formatWorldRecord(run *srapi.Run, game *srapi.Game, cat *srapi.Category, players []*srapi.Player, region *srapi.Region, platform *srapi.Platform) string {
	var err *srapi.Error

	if game == nil {
		game, err = run.Game(srapi.NoEmbeds)
		if err != nil {
			return "Could not fetch game."
		}
	}

	if cat == nil {
		cat, err = run.Category(srapi.NoEmbeds)
		if err != nil {
			return "Could not fetch category."
		}
	}

	formatted := fmt.Sprintf("WR for %s [%s] is %s", game.Names.International, cat.Name, run.Times.Primary.Format())

	if run.Times.IngameTime.Duration > 0 {
		formatted += " (" + run.Times.IngameTime.Format() + " IGT)"
	}

	// collect player names
	names := []string{}

	if len(players) == 0 {
		players, err = run.Players()
		if err != nil {
			return "Could not fetch players: " + err.Error()
		}
	}

	for _, player := range players {
		names = append(names, player.Name())
	}

	formatted += " by " + bot.HumanJoin(names, ", ")

	if run.Date != nil {
		now := time.Now()
		duration := int(now.Sub(run.Date.Time).Hours() / 24)
		date := ""

		switch duration {
		case 0:
			date = "today"
		case 1:
			date = "yesterday"
		case -1:
			date = "tomorrow"
		default:
			if duration > 0 {
				date = fmt.Sprintf("%d days ago", duration)
			} else {
				date = fmt.Sprintf("in %d days", -duration)
			}
		}

		formatted += ", " + date
	}

	// append platform info
	if platform == nil {
		platform, err = run.Platform()
		if err != nil {
			return "Could not fetch platform."
		}
	}

	showRegion := true

	if platform != nil {
		formatted = formatted + " (played on " + platform.Name
		showRegion = platform.ID != "8zjwp7vo" // do not show on PC
	}

	// append region info

	if showRegion {
		if region == nil {
			region, err = run.Region()
			if err != nil {
				return "Could not fetch region."
			}
		}

		if region != nil {
			formatted += ", " + region.Name
		}
	}

	if platform != nil {
		formatted += ")"
	}

	return formatted + "."
}
<commit_msg>update to recent changes in srapi<commit_after>package plugin

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/sgt-kabukiman/kabukibot/bot"
	"github.com/sgt-kabukiman/srapi"
)

type speedruncomConfig struct {
	Interval int
	Mapping  map[string]map[string]string
}

type SpeedrunComPlugin struct {
	config speedruncomConfig
	dict   *bot.Dictionary
}

func NewSpeedrunComPlugin() *SpeedrunComPlugin {
	return &SpeedrunComPlugin{}
}

func (self *SpeedrunComPlugin) Name() string {
	return "speedruncom"
}

func (self *SpeedrunComPlugin) Permissions() []string {
	return []string{"use_speedruncom_commands"}
}

func (self *SpeedrunComPlugin) Setup(bot *bot.Kabukibot) {
	self.config = speedruncomConfig{}
	self.dict = bot.Dictionary()

	err := bot.Configuration().PluginConfig(self.Name(), &self.config)
	if err != nil {
		bot.Logger().Warn("Could not load 'speedruncom' plugin configuration: %s", err)
	}

	go self.updater()
}

func (self *SpeedrunComPlugin) updater() {
	interval := time.Duration(self.config.Interval * int(time.Minute))

	for {
		for gameID, catList := range self.config.Mapping {
			game, err := srapi.GameByID(gameID, srapi.NoEmbeds)
			if err != nil {
				continue
			}

			leaderboards, err := game.Records(nil, "players,regions,platforms,category")
			if err != nil {
				continue
			}

			leaderboards.Walk(func(lb *srapi.Leaderboard) bool {
				if len(lb.Runs) == 0 {
					return true
				}

				cat, err := lb.Category(srapi.NoEmbeds)
				if err != nil {
					return true
				}

				dictKey, okay := catList[cat.ID]
				if !okay {
					return true
				}

				wr := lb.Runs[0]
				formatted := formatWorldRecord(&wr.Run, game, cat, nil, nil, nil)

				self.dict.Set(dictKey, formatted)

				return true
			})
		}

		time.Sleep(interval)
	}
}

func (self *SpeedrunComPlugin) CreateWorker(channel bot.Channel) bot.PluginWorker {
	return &speedruncomWorker{
		channel: channel.Name(),
		acl:     channel.ACL(),
	}
}

type speedruncomWorker struct {
	channel string
	acl     *bot.ACL
}

func (self *speedruncomWorker) Enable() {
	// nothing to do for us
}

func (self *speedruncomWorker) Disable() {
	// nothing to do for us
}

func (self *speedruncomWorker) Part() {
	// nothing to do for us
}

func (self *speedruncomWorker) Shutdown() {
	// nothing to do for us
}

func (self *speedruncomWorker) HandleTextMessage(msg *bot.TextMessage, sender bot.Sender) {
	if msg.IsProcessed() || msg.IsFromBot() {
		return
	}

	if msg.IsCommand("wr") {
		self.handleWorldRecordCommand(msg, sender)
		msg.SetProcessed()
	}
}

var cleanerRegexp = regexp.MustCompile(`[^a-zA-Z0-9]`)

func (self *speedruncomWorker) handleWorldRecordCommand(msg *bot.TextMessage, sender bot.Sender) {
	if !self.acl.IsAllowed(msg.User, "use_speedrun_commands") {
		return
	}

	args := msg.Arguments()

	if len(args) == 0 {
		sender.Respond("you have to give me a game abbreviation.")
		return
	}

	gameIdentifier := args[0]

	var category *srapi.Category

	// try to find the game
	game, err := srapi.GameByAbbreviation(gameIdentifier, "categories")
	if err != nil {
		sender.Respond("I could not find a game with the abbreviation \"" + gameIdentifier + "\".")
		return
	}

	// assume all further args form the category, like "All Missions" or "Any%";
	// we normalise the value to make it -- hopefully -- easier to find the correct category

	if len(args) > 1 {
		catIdentifier := cleanerRegexp.ReplaceAllString(strings.ToLower(strings.Join(args[1:], "")), "")

		categories, err := game.Categories(nil, nil, srapi.NoEmbeds)
		catNames := []string{}

		if err == nil {
			categories.Walk(func(cat *srapi.Category) bool {
				id := cleanerRegexp.ReplaceAllString(strings.ToLower(cat.Name), "")

				if cat.Type == "per-game" {
					catNames = append(catNames, cat.Name)
				}

				if id == catIdentifier {
					category = cat
				}

				return true
			})
		}

		if category == nil {
			sender.Respond("I could not find a category named \"" + strings.Join(args[1:], " ") + "\". Available categories are: " + bot.HumanJoin(catNames, ", "))
			return
		} else if category.Type != "per-game" {
			sender.Respond(category.Name + " is a IL category; cannot report records for now. Sorry.")
			return
		}
	}

	var lb *srapi.Leaderboard

	// fetch the leaderboard, if possible (only available for games with full-game categories by default)
	if category == nil {
		lb, err = game.PrimaryLeaderboard(&srapi.LeaderboardOptions{Top: 1}, "players,platforms,regions,category")
		if err != nil || lb == nil {
			sender.Respond(game.Names.International + " does not use full-game categories by default, so I don't know what category or level you are referring to.")
			return
		}

		category, err = lb.Category(srapi.NoEmbeds)
		if err != nil {
			sender.Respond("the data from speedrun.com is invalid, cannot procede. Sorry. Try again later or a with different game.")
			return
		}
	} else {
		lb, err = category.PrimaryLeaderboard(&srapi.LeaderboardOptions{Top: 1}, "players,platforms,regions")
		if err != nil || lb == nil {
			sender.Respond(game.Names.International + " does not have runs for its \"" + category.Name + "\" category.")
			return
		}
	}

	// the leaderboard could be empty
	if len(lb.Runs) == 0 {
		sender.Respond(game.Names.International + ": " + category.Name + " does not have any matching runs yet.")
		return
	}

	// show only the first WR
	firstRun := lb.Runs[0].Run
	formatted := formatWorldRecord(&firstRun, game, category, nil, nil, nil)

	sender.SendText(formatted)
}

func formatWorldRecord(run *srapi.Run, game *srapi.Game, cat *srapi.Category, players *srapi.PlayerCollection, region *srapi.Region, platform *srapi.Platform) string {
	var err *srapi.Error

	if game == nil {
		game, err = run.Game(srapi.NoEmbeds)
		if err != nil {
			return "Could not fetch game: " + err.Error()
		}
	}

	if cat == nil {
		cat, err = run.Category(srapi.NoEmbeds)
		if err != nil {
			return "Could not fetch category: " + err.Error()
		}
	}

	if players == nil {
		players, err = run.Players()
		if err != nil {
			return "Could not fetch players: " + err.Error()
		}
	}

	formatted := fmt.Sprintf("WR for %s [%s] is %s", game.Names.International, cat.Name, run.Times.Primary.Format())

	if run.Times.IngameTime.Duration > 0 {
		formatted += " (" + run.Times.IngameTime.Format() + " IGT)"
	}

	// collect player names
	names := []string{}

	players.Walk(func(p *srapi.Player) bool {
		names = append(names, p.Name())
		return true
	})

	formatted += " by " + bot.HumanJoin(names, ", ")

	if run.Date != nil {
		now := time.Now()
		duration := int(now.Sub(run.Date.Time).Hours() / 24)
		date := ""

		switch duration {
		case 0:
			date = "today"
		case 1:
			date = "yesterday"
		case -1:
			date = "tomorrow"
		default:
			if duration > 0 {
				date = fmt.Sprintf("%d days ago", duration)
			} else {
				date = fmt.Sprintf("in %d days", -duration)
			}
		}

		formatted += ", " + date
	}

	// append platform info
	if platform == nil {
		platform, err = run.Platform()
		if err != nil {
			return "Could not fetch platform."
		}
	}

	showRegion := true

	if platform != nil {
		formatted = formatted + " (played on " + platform.Name
		showRegion = platform.ID != "8zjwp7vo" // do not show on PC
	}

	// append region info

	if showRegion {
		if region == nil {
			region, err = run.Region()
			if err != nil {
				return "Could not fetch region."
			}
		}

		if region != nil {
			formatted += ", " + region.Name
		}
	}

	if platform != nil {
		formatted += ")"
	}

	return formatted + "."
}
<|endoftext|>


<commit_before>package air

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"path"
	"reflect"
	"runtime"
	"sync"
)

type (
	// Air is the top-level framework struct.
	Air struct {
		pregases    []Gas
		gases       []Gas
		maxParam    int
		contextPool *sync.Pool
		server      *server
		router      *router

		Config           *Config
		Logger           Logger
		Binder           Binder
		Renderer         Renderer
		HTTPErrorHandler HTTPErrorHandler
	}

	// Handler defines a function to serve HTTP requests.
	Handler func(*Context) error

	// Gas defines a function to process gases.
	Gas func(Handler) Handler

	// HTTPError represents an error that occurred while handling an HTTP request.
	HTTPError struct {
		Code    int
		Message string
	}

	// HTTPErrorHandler is a centralized HTTP error handler.
	HTTPErrorHandler func(error, *Context)

	// JSONMap is a map that organizes data in JSON format.
	JSONMap map[string]interface{}
)

// HTTP methods (which follows the REST principle)
const (
	GET    = "GET"
	POST   = "POST"
	PUT    = "PUT" // The Air advises you to forget the PATCH.
	DELETE = "DELETE"
)

// For easy for-range
var methods = [4]string{GET, POST, PUT, DELETE}

// MIME types
const (
	MIMEApplicationJSON       = "application/json; charset=utf-8"
	MIMEApplicationJavaScript = "application/javascript; charset=utf-8"
	MIMEApplicationXML        = "application/xml; charset=utf-8"
	MIMEApplicationYAML       = "application/x-yaml; charset=utf-8"
	MIMEApplicationForm       = "application/x-www-form-urlencoded"
	MIMEApplicationProtobuf   = "application/protobuf"
	MIMEApplicationMsgpack    = "application/msgpack"
	MIMETextHTML              = "text/html; charset=utf-8"
	MIMETextPlain             = "text/plain; charset=utf-8"
	MIMEMultipartForm         = "multipart/form-data"
	MIMEOctetStream           = "application/octet-stream"
)

// Headers
const (
	HeaderAcceptEncoding                = "Accept-Encoding"
	HeaderAllow                         = "Allow"
	HeaderAuthorization                 = "Authorization"
	HeaderContentDisposition            = "Content-Disposition"
	HeaderContentEncoding               = "Content-Encoding"
	HeaderContentLength                 = "Content-Length"
	HeaderContentType                   = "Content-Type"
	HeaderCookie                        = "Cookie"
	HeaderSetCookie                     = "Set-Cookie"
	HeaderIfModifiedSince               = "If-Modified-Since"
	HeaderLastModified                  = "Last-Modified"
	HeaderLocation                      = "Location"
	HeaderUpgrade                       = "Upgrade"
	HeaderVary                          = "Vary"
	HeaderWWWAuthenticate               = "WWW-Authenticate"
	HeaderXForwardedProto               = "X-Forwarded-Proto"
	HeaderXHTTPMethodOverride           = "X-HTTP-Method-Override"
	HeaderXForwardedFor                 = "X-Forwarded-For"
	HeaderXRealIP                       = "X-Real-IP"
	HeaderServer                        = "Server"
	HeaderOrigin                        = "Origin"
	HeaderAccessControlRequestMethod    = "Access-Control-Request-Method"
	HeaderAccessControlRequestHeaders   = "Access-Control-Request-Headers"
	HeaderAccessControlAllowOrigin      = "Access-Control-Allow-Origin"
	HeaderAccessControlAllowMethods     = "Access-Control-Allow-Methods"
	HeaderAccessControlAllowHeaders     = "Access-Control-Allow-Headers"
	HeaderAccessControlAllowCredentials = "Access-Control-Allow-Credentials"
	HeaderAccessControlExposeHeaders    = "Access-Control-Expose-Headers"
	HeaderAccessControlMaxAge           = "Access-Control-Max-Age"

	HeaderStrictTransportSecurity = "Strict-Transport-Security"
	HeaderXContentTypeOptions     = "X-Content-Type-Options"
	HeaderXXSSProtection          = "X-XSS-Protection"
	HeaderXFrameOptions           = "X-Frame-Options"
	HeaderContentSecurityPolicy   = "Content-Security-Policy"
	HeaderXCSRFToken              = "X-CSRF-Token"
)

// Errors
var (
	ErrUnauthorized                = NewHTTPError(http.StatusUnauthorized)          // 401
	ErrNotFound                    = NewHTTPError(http.StatusNotFound)              // 404
	ErrMethodNotAllowed            = NewHTTPError(http.StatusMethodNotAllowed)      // 405
	ErrStatusRequestEntityTooLarge = NewHTTPError(http.StatusRequestEntityTooLarge) // 413
	ErrUnsupportedMediaType        = NewHTTPError(http.StatusUnsupportedMediaType)  // 415

	ErrInternalServerError = NewHTTPError(http.StatusInternalServerError) // 500
	ErrBadGateway          = NewHTTPError(http.StatusBadGateway)          // 502
	ErrServiceUnavailable  = NewHTTPError(http.StatusServiceUnavailable)  // 503
	ErrGatewayTimeout      = NewHTTPError(http.StatusGatewayTimeout)      // 504

	ErrDisabledHTTP2       = errors.New("the HTTP/2 has been disabled")
	ErrInvalidRedirectCode = errors.New("invalid redirect status code")
)

// HTTP error handlers
var (
	NotFoundHandler = func(c *Context) error {
		return ErrNotFound
	}

	MethodNotAllowedHandler = func(c *Context) error {
		return ErrMethodNotAllowed
	}
)

// New returns a pointer of a new instance of the `Air`.
func New() *Air {
	a := &Air{}

	a.contextPool = &sync.Pool{
		New: func() interface{} {
			return newContext(a)
		},
	}
	a.router = newRouter(a)

	a.Config = newConfig()
	a.Logger = newLogger(a)
	a.Binder = newBinder()
	a.Renderer = newRenderer(a)
	a.HTTPErrorHandler = defaultHTTPErrorHandler

	return a
}

// Precontain adds the gases to the chain which is perform before the router.
func (a *Air) Precontain(gases ...Gas) {
	a.pregases = append(a.pregases, gases...)
}

// Contain adds the gases to the chain which is perform after the router.
func (a *Air) Contain(gases ...Gas) {
	a.gases = append(a.gases, gases...)
}

// GET registers a new GET route for the path with the matching h in the router with the optional
// route-level gases.
func (a *Air) GET(path string, h Handler, gases ...Gas) {
	a.add(GET, path, h, gases...)
}

// POST registers a new POST route for the path with the matching h in the router with the optional
// route-level gases.
func (a *Air) POST(path string, h Handler, gases ...Gas) {
	a.add(POST, path, h, gases...)
}

// PUT registers a new PUT route for the path with the matching h in the router with the optional
// route-level gases.
func (a *Air) PUT(path string, h Handler, gases ...Gas) {
	a.add(PUT, path, h, gases...)
}

// DELETE registers a new DELETE route for the path with the matching h in the router with the
// optional route-level gases.
func (a *Air) DELETE(path string, h Handler, gases ...Gas) {
	a.add(DELETE, path, h, gases...)
}

// Static registers a new route with the path prefix to serve the static files from the provided
// root directory.
func (a *Air) Static(prefix, root string) {
	a.GET(prefix+"*", func(c *Context) error {
		return c.File(path.Join(root, c.Param("*")))
	})
}

// File registers a new route with the path to serve a static file.
func (a *Air) File(path, file string) {
	a.GET(path, func(c *Context) error {
		return c.File(file)
	})
}

// add registers a new route for the path with the method and the matching h in the router with the
// optional route-level gases.
func (a *Air) add(method, path string, h Handler, gases ...Gas) {
	hn := handlerName(h)

	a.router.add(method, path, func(c *Context) error {
		hf := h
		for i := len(gases) - 1; i >= 0; i-- {
			hf = gases[i](hf)
		}
		return hf(c)
	})

	a.router.routes[method+path] = &route{
		method:  method,
		path:    path,
		handler: hn,
	}
}

// URL returns an URL generated from the h with the optional params.
func (a *Air) URL(h Handler, params ...interface{}) string {
	url := &bytes.Buffer{}
	hn := handlerName(h)
	ln := len(params)
	n := 0

	for _, r := range a.router.routes {
		if r.handler == hn {
			for i, l := 0, len(r.path); i < l; i++ {
				if r.path[i] == ':' && n < ln {
					for ; i < l && r.path[i] != '/'; i++ {
					}
					url.WriteString(fmt.Sprintf("%v", params[n]))
					n++
				}

				if i < l {
					url.WriteByte(r.path[i])
				}
			}

			break
		}
	}

	return url.String()
}

// Serve starts the HTTP server.
func (a *Air) Serve() {
	if a.Config.DebugMode {
		a.Config.LogEnabled = true
		a.Logger.Debug("serving in debug mode")
	}

	if err := a.Renderer.ParseTemplates(); err != nil {
		panic(err)
	}

	a.server = newServer(a)
	if err := a.server.serve(); err != nil {
		panic(err)
	}
}

// Close closes the HTTP server immediately.
func (a *Air) Close() error {
	return a.server.Close()
}

// Shutdown gracefully shuts down the HTTP server without interrupting any active connections.
func (a *Air) Shutdown(c *Context) error {
	return a.server.Shutdown(c.Context)
}

// handlerName returns the func name of the h.
func handlerName(h Handler) string {
	t := reflect.ValueOf(h).Type()
	if t.Kind() == reflect.Func {
		return runtime.FuncForPC(reflect.ValueOf(h).Pointer()).Name()
	}
	return t.String()
}

// WrapGas wraps the h into the `Gas`.
func WrapGas(h Handler) Gas {
	return func(next Handler) Handler {
		return func(c *Context) error {
			if err := h(c); err != nil {
				return err
			}
			return next(c)
		}
	}
}

// NewHTTPError returns a pointer of a new instance of the `HTTPError`.
func NewHTTPError(code int, msg ...interface{}) *HTTPError {
	he := &HTTPError{Code: code, Message: http.StatusText(code)}
	if len(msg) > 0 {
		he.Message = fmt.Sprint(msg...)
	}
	return he
}

// Error implements the `error#Error()`.
func (he *HTTPError) Error() string {
	return he.Message
}

// defaultHTTPErrorHandler invokes the default HTTP error handler.
func defaultHTTPErrorHandler(err error, c *Context) {
	he := ErrInternalServerError

	if che, ok := err.(*HTTPError); ok {
		he = che
	}

	if c.Air.Config.DebugMode {
		he.Message = err.Error()
	}

	if !c.Response.Written() {
		c.Response.WriteHeader(he.Code)
		c.String(he.Message)
	}

	c.Air.Logger.Error(err)
}
<commit_msg>feat: add the `WrapHandler()`<commit_after>package air

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"path"
	"reflect"
	"runtime"
	"sync"
)

type (
	// Air is the top-level framework struct.
	Air struct {
		pregases    []Gas
		gases       []Gas
		maxParam    int
		contextPool *sync.Pool
		server      *server
		router      *router

		Config           *Config
		Logger           Logger
		Binder           Binder
		Renderer         Renderer
		HTTPErrorHandler HTTPErrorHandler
	}

	// Handler defines a function to serve HTTP requests.
	Handler func(*Context) error

	// Gas defines a function to process gases.
	Gas func(Handler) Handler

	// HTTPError represents an error that occurred while handling an HTTP request.
	HTTPError struct {
		Code    int
		Message string
	}

	// HTTPErrorHandler is a centralized HTTP error handler.
	HTTPErrorHandler func(error, *Context)

	// JSONMap is a map that organizes data in JSON format.
	JSONMap map[string]interface{}
)

// HTTP methods (which follows the REST principle)
const (
	GET    = "GET"
	POST   = "POST"
	PUT    = "PUT" // The Air advises you to forget the PATCH.
	DELETE = "DELETE"
)

// For easy for-range
var methods = [4]string{GET, POST, PUT, DELETE}

// MIME types
const (
	MIMEApplicationJSON       = "application/json; charset=utf-8"
	MIMEApplicationJavaScript = "application/javascript; charset=utf-8"
	MIMEApplicationXML        = "application/xml; charset=utf-8"
	MIMEApplicationYAML       = "application/x-yaml; charset=utf-8"
	MIMEApplicationForm       = "application/x-www-form-urlencoded"
	MIMEApplicationProtobuf   = "application/protobuf"
	MIMEApplicationMsgpack    = "application/msgpack"
	MIMETextHTML              = "text/html; charset=utf-8"
	MIMETextPlain             = "text/plain; charset=utf-8"
	MIMEMultipartForm         = "multipart/form-data"
	MIMEOctetStream           = "application/octet-stream"
)

// Headers
const (
	HeaderAcceptEncoding                = "Accept-Encoding"
	HeaderAllow                         = "Allow"
	HeaderAuthorization                 = "Authorization"
	HeaderContentDisposition            = "Content-Disposition"
	HeaderContentEncoding               = "Content-Encoding"
	HeaderContentLength                 = "Content-Length"
	HeaderContentType                   = "Content-Type"
	HeaderCookie                        = "Cookie"
	HeaderSetCookie                     = "Set-Cookie"
	HeaderIfModifiedSince               = "If-Modified-Since"
	HeaderLastModified                  = "Last-Modified"
	HeaderLocation                      = "Location"
	HeaderUpgrade                       = "Upgrade"
	HeaderVary                          = "Vary"
	HeaderWWWAuthenticate               = "WWW-Authenticate"
	HeaderXForwardedProto               = "X-Forwarded-Proto"
	HeaderXHTTPMethodOverride           = "X-HTTP-Method-Override"
	HeaderXForwardedFor                 = "X-Forwarded-For"
	HeaderXRealIP                       = "X-Real-IP"
	HeaderServer                        = "Server"
	HeaderOrigin                        = "Origin"
	HeaderAccessControlRequestMethod    = "Access-Control-Request-Method"
	HeaderAccessControlRequestHeaders   = "Access-Control-Request-Headers"
	HeaderAccessControlAllowOrigin      = "Access-Control-Allow-Origin"
	HeaderAccessControlAllowMethods     = "Access-Control-Allow-Methods"
	HeaderAccessControlAllowHeaders     = "Access-Control-Allow-Headers"
	HeaderAccessControlAllowCredentials = "Access-Control-Allow-Credentials"
	HeaderAccessControlExposeHeaders    = "Access-Control-Expose-Headers"
	HeaderAccessControlMaxAge           = "Access-Control-Max-Age"

	HeaderStrictTransportSecurity = "Strict-Transport-Security"
	HeaderXContentTypeOptions     = "X-Content-Type-Options"
	HeaderXXSSProtection          = "X-XSS-Protection"
	HeaderXFrameOptions           = "X-Frame-Options"
	HeaderContentSecurityPolicy   = "Content-Security-Policy"
	HeaderXCSRFToken              = "X-CSRF-Token"
)

// Errors
var (
	ErrUnauthorized                = NewHTTPError(http.StatusUnauthorized)          // 401
	ErrNotFound                    = NewHTTPError(http.StatusNotFound)              // 404
	ErrMethodNotAllowed            = NewHTTPError(http.StatusMethodNotAllowed)      // 405
	ErrStatusRequestEntityTooLarge = NewHTTPError(http.StatusRequestEntityTooLarge) // 413
	ErrUnsupportedMediaType        = NewHTTPError(http.StatusUnsupportedMediaType)  // 415

	ErrInternalServerError = NewHTTPError(http.StatusInternalServerError) // 500
	ErrBadGateway          = NewHTTPError(http.StatusBadGateway)          // 502
	ErrServiceUnavailable  = NewHTTPError(http.StatusServiceUnavailable)  // 503
	ErrGatewayTimeout      = NewHTTPError(http.StatusGatewayTimeout)      // 504

	ErrDisabledHTTP2       = errors.New("the HTTP/2 has been disabled")
	ErrInvalidRedirectCode = errors.New("invalid redirect status code")
)

// HTTP error handlers
var (
	NotFoundHandler = func(c *Context) error {
		return ErrNotFound
	}

	MethodNotAllowedHandler = func(c *Context) error {
		return ErrMethodNotAllowed
	}
)

// New returns a pointer of a new instance of the `Air`.
func New() *Air {
	a := &Air{}

	a.contextPool = &sync.Pool{
		New: func() interface{} {
			return newContext(a)
		},
	}
	a.router = newRouter(a)

	a.Config = newConfig()
	a.Logger = newLogger(a)
	a.Binder = newBinder()
	a.Renderer = newRenderer(a)
	a.HTTPErrorHandler = defaultHTTPErrorHandler

	return a
}

// Precontain adds the gases to the chain which is perform before the router.
func (a *Air) Precontain(gases ...Gas) {
	a.pregases = append(a.pregases, gases...)
}

// Contain adds the gases to the chain which is perform after the router.
func (a *Air) Contain(gases ...Gas) {
	a.gases = append(a.gases, gases...)
}

// GET registers a new GET route for the path with the matching h in the router with the optional
// route-level gases.
func (a *Air) GET(path string, h Handler, gases ...Gas) {
	a.add(GET, path, h, gases...)
}

// POST registers a new POST route for the path with the matching h in the router with the optional
// route-level gases.
func (a *Air) POST(path string, h Handler, gases ...Gas) {
	a.add(POST, path, h, gases...)
}

// PUT registers a new PUT route for the path with the matching h in the router with the optional
// route-level gases.
func (a *Air) PUT(path string, h Handler, gases ...Gas) {
	a.add(PUT, path, h, gases...)
}

// DELETE registers a new DELETE route for the path with the matching h in the router with the
// optional route-level gases.
func (a *Air) DELETE(path string, h Handler, gases ...Gas) {
	a.add(DELETE, path, h, gases...)
}

// Static registers a new route with the path prefix to serve the static files from the provided
// root directory.
func (a *Air) Static(prefix, root string) {
	a.GET(prefix+"*", func(c *Context) error {
		return c.File(path.Join(root, c.Param("*")))
	})
}

// File registers a new route with the path to serve a static file.
func (a *Air) File(path, file string) {
	a.GET(path, func(c *Context) error {
		return c.File(file)
	})
}

// add registers a new route for the path with the method and the matching h in the router with the
// optional route-level gases.
func (a *Air) add(method, path string, h Handler, gases ...Gas) {
	hn := handlerName(h)

	a.router.add(method, path, func(c *Context) error {
		hf := h
		for i := len(gases) - 1; i >= 0; i-- {
			hf = gases[i](hf)
		}
		return hf(c)
	})

	a.router.routes[method+path] = &route{
		method:  method,
		path:    path,
		handler: hn,
	}
}

// URL returns an URL generated from the h with the optional params.
func (a *Air) URL(h Handler, params ...interface{}) string {
	url := &bytes.Buffer{}
	hn := handlerName(h)
	ln := len(params)
	n := 0

	for _, r := range a.router.routes {
		if r.handler == hn {
			for i, l := 0, len(r.path); i < l; i++ {
				if r.path[i] == ':' && n < ln {
					for ; i < l && r.path[i] != '/'; i++ {
					}
					url.WriteString(fmt.Sprintf("%v", params[n]))
					n++
				}

				if i < l {
					url.WriteByte(r.path[i])
				}
			}

			break
		}
	}

	return url.String()
}

// Serve starts the HTTP server.
func (a *Air) Serve() {
	if a.Config.DebugMode {
		a.Config.LogEnabled = true
		a.Logger.Debug("serving in debug mode")
	}

	if err := a.Renderer.ParseTemplates(); err != nil {
		panic(err)
	}

	a.server = newServer(a)
	if err := a.server.serve(); err != nil {
		panic(err)
	}
}

// Close closes the HTTP server immediately.
func (a *Air) Close() error {
	return a.server.Close()
}

// Shutdown gracefully shuts down the HTTP server without interrupting any active connections.
func (a *Air) Shutdown(c *Context) error {
	return a.server.Shutdown(c.Context)
}

// handlerName returns the func name of the h.
func handlerName(h Handler) string {
	t := reflect.ValueOf(h).Type()
	if t.Kind() == reflect.Func {
		return runtime.FuncForPC(reflect.ValueOf(h).Pointer()).Name()
	}
	return t.String()
}

// WrapHandler wraps the h into the `Handler`.
func WrapHandler(h http.Handler) Handler {
	return func(c *Context) error {
		h.ServeHTTP(c.Response.ResponseWriter, c.Request.Request)
		return nil
	}
}

// WrapGas wraps the h into the `Gas`.
func WrapGas(h Handler) Gas {
	return func(next Handler) Handler {
		return func(c *Context) error {
			if err := h(c); err != nil {
				return err
			}
			return next(c)
		}
	}
}

// NewHTTPError returns a pointer of a new instance of the `HTTPError`.
func NewHTTPError(code int, msg ...interface{}) *HTTPError {
	he := &HTTPError{Code: code, Message: http.StatusText(code)}
	if len(msg) > 0 {
		he.Message = fmt.Sprint(msg...)
	}
	return he
}

// Error implements the `error#Error()`.
func (he *HTTPError) Error() string {
	return he.Message
}

// defaultHTTPErrorHandler invokes the default HTTP error handler.
func defaultHTTPErrorHandler(err error, c *Context) {
	he := ErrInternalServerError

	if che, ok := err.(*HTTPError); ok {
		he = che
	}

	if c.Air.Config.DebugMode {
		he.Message = err.Error()
	}

	if !c.Response.Written() {
		c.Response.WriteHeader(he.Code)
		c.String(he.Message)
	}

	c.Air.Logger.Error(err)
}
<|endoftext|>