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

package app

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"math/big"
	"net/http"
	"net/http/httptest"
	"path"
	"testing"
	"time"
)

func TestCACacheVerify(t *testing.T) {
	var ca fakeCA
	srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		switch path.Base(r.URL.Path) {
		case "root.crt":
			rw.Write(ca.rootCert)
		case "intermediate.crt":
			rw.Write(ca.intermediateCert)
		case "self-signed-intermediate.crt":
			rw.Write(ca.selfSignedIntermediateCert)
		case "root.crl":
			rw.Write(ca.rootCRL)
		case "intermediate.crl":
			rw.Write(ca.intermediateCRL)
		case "self-signed-intermediate.crl":
			rw.Write(ca.selfSignedIntermediateCRL)
		default:
			http.Error(rw, "not found", http.StatusNotFound)
		}
	}))
	defer srv.Close()
	ca = initFakeCA(t, srv.URL)

	c := &caCache{
		rootCertURL: srv.URL + "/root.crt",
		interPrefix: srv.URL,
		certs:       make(map[string]*x509.Certificate),
		crls:        make(map[string]*cachedCRL),
	}

	t.Run("valid", func(t *testing.T) {
		t.Parallel()
		if err := c.verify(ca.validCert); err != nil {
			t.Errorf("verifying valid certificate: got %v, want nil", err)
		}
	})
	if err := c.verify(ca.validCert); err != nil {
		t.Errorf("verifying valid certificate: got %v, want nil", err)
	}
	for desc, invalidCert := range ca.invalidCerts {
		t.Run(desc, func(t *testing.T) {
			t.Parallel()
			if err := c.verify(invalidCert); err == nil {
				t.Errorf("verifying invalid certificate: got nil, want non-nil error")
			}
		})
	}
}

type fakeCA struct {
	rootCert, intermediateCert, selfSignedIntermediateCert []byte
	rootCRL, intermediateCRL, selfSignedIntermediateCRL    []byte
	validCert                                              *x509.Certificate
	invalidCerts                                           map[string]*x509.Certificate
}

func initFakeCA(t *testing.T, srvURL string) fakeCA {
	t.Helper()
	ca := fakeCA{invalidCerts: make(map[string]*x509.Certificate)}

	rootTmpl := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		IsCA:         true,
		BasicConstraintsValid: true,
	}
	rootCertDER, rootCert, rootKey := makeCert(t, rootTmpl, rootTmpl, nil)
	ca.rootCert = rootCertDER
	interTmpl := &x509.Certificate{
		SerialNumber: big.NewInt(2),
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		IsCA:         true,
		BasicConstraintsValid: true,
		CRLDistributionPoints: []string{srvURL + "/root.crl"},
	}
	interCertDER, interCert, interKey := makeCert(t, interTmpl, rootCert, rootKey)
	ca.intermediateCert = interCertDER

	_, ca.validCert, _ = makeCert(t, &x509.Certificate{
		SerialNumber:          big.NewInt(3),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{srvURL + "/intermediate.crt"},
		CRLDistributionPoints: []string{srvURL + "/intermediate.crl"},
	}, interCert, interKey)

	_, ca.invalidCerts["revoked"], _ = makeCert(t, &x509.Certificate{
		SerialNumber:          big.NewInt(4),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{srvURL + "/intermediate.crt"},
		CRLDistributionPoints: []string{srvURL + "/intermediate.crl"},
	}, interCert, interKey)
	_, ca.invalidCerts["no CRL"], _ = makeCert(t, &x509.Certificate{
		SerialNumber:          big.NewInt(4),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{srvURL + "/intermediate.crt"},
	}, interCert, interKey)
	selfSignedTmpl := &x509.Certificate{
		SerialNumber: big.NewInt(99),
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
	}
	_, ca.invalidCerts["self signed"], _ = makeCert(t, selfSignedTmpl, selfSignedTmpl, nil)
	_, ca.invalidCerts["wrong intermediate link"], _ = makeCert(t, &x509.Certificate{
		SerialNumber:          big.NewInt(3),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{"http://example.com/wrong-intermediate.crt"},
		CRLDistributionPoints: []string{srvURL + "/intermediate.crl"},
	}, interCert, interKey)

	selfSignedInterTmpl := &x509.Certificate{
		SerialNumber: big.NewInt(2),
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		IsCA:         true,
		BasicConstraintsValid: true,
		CRLDistributionPoints: []string{srvURL + "/root.crl"},
	}
	selfSignedInterCertDER, selfSignedInterCert, selfSignedInterKey := makeCert(t, selfSignedInterTmpl, selfSignedInterTmpl, nil)
	ca.selfSignedIntermediateCert = selfSignedInterCertDER
	wrongInterLeafTmpl := &x509.Certificate{
		SerialNumber:          big.NewInt(3),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{srvURL + "/self-signed-intermediate.crt"},
		CRLDistributionPoints: []string{srvURL + "/self-signed-intermediate.crl"},
	}
	_, ca.invalidCerts["self-signed intermediate"], _ = makeCert(t, wrongInterLeafTmpl, selfSignedInterCert, selfSignedInterKey)

	var err error
	ca.rootCRL, err = rootCert.CreateCRL(rand.Reader, rootKey, nil, time.Now(), time.Now())
	if err != nil {
		t.Fatal(err)
	}
	ca.intermediateCRL, err = interCert.CreateCRL(rand.Reader, interKey, []pkix.RevokedCertificate{{
		SerialNumber:   big.NewInt(4),
		RevocationTime: time.Now(),
	}}, time.Now(), time.Now())
	if err != nil {
		t.Fatal(err)
	}
	ca.selfSignedIntermediateCRL, err = selfSignedInterCert.CreateCRL(rand.Reader, selfSignedInterKey, nil, time.Now(), time.Now())
	if err != nil {
		t.Fatal(err)
	}

	return ca
}

func makeCert(t *testing.T, tmpl, parent *x509.Certificate, parentKey crypto.PrivateKey) ([]byte, *x509.Certificate, crypto.PrivateKey) {
	t.Helper()

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	if parentKey == nil {
		parentKey = key
	}
	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, parent, key.Public(), parentKey)
	if err != nil {
		t.Fatal(err)
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		t.Fatal(err)
	}

	return certDER, cert, key
}
<commit_msg>Refactor fakeCA in ca_cache_test.go to reuse in TPM validation tests<commit_after>/*
Copyright 2018 The Kubernetes Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"math/big"
	"net/http"
	"net/http/httptest"
	"path"
	"testing"
	"time"

	"k8s.io/cloud-provider-gcp/pkg/nodeidentity"
)

func TestCACacheVerify(t *testing.T) {
	ca, c, cleanup := initFakeCACache(t)
	defer cleanup()

	t.Run("valid", func(t *testing.T) {
		t.Parallel()
		if err := c.verify(ca.validCert); err != nil {
			t.Errorf("verifying valid certificate: got %v, want nil", err)
		}
	})
	if err := c.verify(ca.validCert); err != nil {
		t.Errorf("verifying valid certificate: got %v, want nil", err)
	}
	for desc, invalidCert := range ca.invalidCerts {
		t.Run(desc, func(t *testing.T) {
			t.Parallel()
			if err := c.verify(invalidCert); err == nil {
				t.Errorf("verifying invalid certificate: got nil, want non-nil error")
			}
		})
	}
}

func initFakeCACache(t *testing.T) (*fakeCA, *caCache, func()) {
	var ca *fakeCA
	srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		switch path.Base(r.URL.Path) {
		case "root.crt":
			rw.Write(ca.rootCert)
		case "intermediate.crt":
			rw.Write(ca.intermediateCertRaw)
		case "self-signed-intermediate.crt":
			rw.Write(ca.selfSignedIntermediateCert)
		case "root.crl":
			rw.Write(ca.rootCRL)
		case "intermediate.crl":
			rw.Write(ca.intermediateCRL)
		case "self-signed-intermediate.crl":
			rw.Write(ca.selfSignedIntermediateCRL)
		default:
			http.Error(rw, "not found", http.StatusNotFound)
		}
	}))

	ca = initFakeCA(t, srv.URL)
	c := &caCache{
		rootCertURL: srv.URL + "/root.crt",
		interPrefix: srv.URL,
		certs:       make(map[string]*x509.Certificate),
		crls:        make(map[string]*cachedCRL),
	}
	return ca, c, srv.Close
}

type fakeCA struct {
	rootCert, intermediateCertRaw, selfSignedIntermediateCert []byte
	rootCRL, intermediateCRL, selfSignedIntermediateCRL       []byte
	intermediateCert                                          *x509.Certificate
	intermediateCertKey                                       *rsa.PrivateKey
	validCert                                                 *x509.Certificate
	validCertKey                                              *rsa.PrivateKey
	invalidCerts                                              map[string]*x509.Certificate
	srvURL                                                    string
}

func initFakeCA(t *testing.T, srvURL string) *fakeCA {
	t.Helper()
	ca := &fakeCA{
		invalidCerts: make(map[string]*x509.Certificate),
		srvURL:       srvURL,
	}

	rootTmpl := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		IsCA:         true,
		BasicConstraintsValid: true,
	}
	rootCertDER, rootCert, rootKey := makeCert(t, rootTmpl, rootTmpl, nil)
	ca.rootCert = rootCertDER
	interTmpl := &x509.Certificate{
		SerialNumber: big.NewInt(2),
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		IsCA:         true,
		BasicConstraintsValid: true,
		CRLDistributionPoints: []string{srvURL + "/root.crl"},
	}
	ca.intermediateCertRaw, ca.intermediateCert, ca.intermediateCertKey = makeCert(t, interTmpl, rootCert, rootKey)

	ca.regenerateValidCert(t, nodeidentity.Identity{"z0", 1, "i0", 2, "p0"})

	_, ca.invalidCerts["revoked"], _ = makeCert(t, &x509.Certificate{
		SerialNumber:          big.NewInt(4),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{srvURL + "/intermediate.crt"},
		CRLDistributionPoints: []string{srvURL + "/intermediate.crl"},
	}, ca.intermediateCert, ca.intermediateCertKey)
	_, ca.invalidCerts["no CRL"], _ = makeCert(t, &x509.Certificate{
		SerialNumber:          big.NewInt(4),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{srvURL + "/intermediate.crt"},
	}, ca.intermediateCert, ca.intermediateCertKey)
	selfSignedTmpl := &x509.Certificate{
		SerialNumber: big.NewInt(99),
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
	}
	_, ca.invalidCerts["self signed"], _ = makeCert(t, selfSignedTmpl, selfSignedTmpl, nil)
	_, ca.invalidCerts["wrong intermediate link"], _ = makeCert(t, &x509.Certificate{
		SerialNumber:          big.NewInt(3),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{"http://example.com/wrong-intermediate.crt"},
		CRLDistributionPoints: []string{srvURL + "/intermediate.crl"},
	}, ca.intermediateCert, ca.intermediateCertKey)

	selfSignedInterTmpl := &x509.Certificate{
		SerialNumber: big.NewInt(2),
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		IsCA:         true,
		BasicConstraintsValid: true,
		CRLDistributionPoints: []string{srvURL + "/root.crl"},
	}
	selfSignedInterCertDER, selfSignedInterCert, selfSignedInterKey := makeCert(t, selfSignedInterTmpl, selfSignedInterTmpl, nil)
	ca.selfSignedIntermediateCert = selfSignedInterCertDER
	wrongInterLeafTmpl := &x509.Certificate{
		SerialNumber:          big.NewInt(3),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{srvURL + "/self-signed-intermediate.crt"},
		CRLDistributionPoints: []string{srvURL + "/self-signed-intermediate.crl"},
	}
	_, ca.invalidCerts["self-signed intermediate"], _ = makeCert(t, wrongInterLeafTmpl, selfSignedInterCert, selfSignedInterKey)

	var err error
	ca.rootCRL, err = rootCert.CreateCRL(insecureRand, rootKey, nil, time.Now(), time.Now())
	if err != nil {
		t.Fatal(err)
	}
	ca.intermediateCRL, err = ca.intermediateCert.CreateCRL(insecureRand, ca.intermediateCertKey, []pkix.RevokedCertificate{{
		SerialNumber:   big.NewInt(4),
		RevocationTime: time.Now(),
	}}, time.Now(), time.Now())
	if err != nil {
		t.Fatal(err)
	}
	ca.selfSignedIntermediateCRL, err = selfSignedInterCert.CreateCRL(insecureRand, selfSignedInterKey, nil, time.Now(), time.Now())
	if err != nil {
		t.Fatal(err)
	}

	return ca
}

func (ca *fakeCA) regenerateValidCert(t *testing.T, id nodeidentity.Identity) {
	vmID, err := id.ToASN1()
	if err != nil {
		t.Fatal(err)
	}
	_, ca.validCert, ca.validCertKey = makeCert(t, &x509.Certificate{
		SerialNumber:          big.NewInt(3),
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		IssuingCertificateURL: []string{ca.srvURL + "/intermediate.crt"},
		CRLDistributionPoints: []string{ca.srvURL + "/intermediate.crl"},
		ExtraExtensions: []pkix.Extension{{
			Id:    nodeidentity.CloudComputeInstanceIdentifierOID,
			Value: vmID,
		}},
	}, ca.intermediateCert, ca.intermediateCertKey)
}

func makeCert(t *testing.T, tmpl, parent *x509.Certificate, parentKey crypto.PrivateKey) ([]byte, *x509.Certificate, *rsa.PrivateKey) {
	t.Helper()

	key, err := rsa.GenerateKey(insecureRand, 2048)
	if err != nil {
		t.Fatal(err)
	}
	if parentKey == nil {
		parentKey = key
	}
	certDER, err := x509.CreateCertificate(insecureRand, tmpl, parent, key.Public(), parentKey)
	if err != nil {
		t.Fatal(err)
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		t.Fatal(err)
	}

	return certDER, cert, key
}
<|endoftext|>


<commit_before>package router

import (
	"bufio"
	"context"
	"errors"
	"github.com/gin-gonic/gin"
	"github.com/pterodactyl/wings/server"
	"golang.org/x/sync/errgroup"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
)

// Returns the contents of a file on the server.
func getServerFileContents(c *gin.Context) {
	s := GetServer(c.Param("server"))

	p, err := url.QueryUnescape(c.Query("file"))
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}
	p = "/" + strings.TrimLeft(p, "/")

	cleaned, err := s.Filesystem.SafePath(p)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
			"error": "The file requested could not be found.",
		})
		return
	}

	st, err := s.Filesystem.Stat(cleaned)
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	if st.Info.IsDir() {
		c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
			"error": "The requested resource was not found on the system.",
		})
		return
	}

	f, err := os.Open(cleaned)
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}
	defer f.Close()

	c.Header("X-Mime-Type", st.Mimetype)
	c.Header("Content-Length", strconv.Itoa(int(st.Info.Size())))

	// If a download parameter is included in the URL go ahead and attach the necessary headers
	// so that the file can be downloaded.
	if c.Query("download") != "" {
		c.Header("Content-Disposition", "attachment; filename="+st.Info.Name())
		c.Header("Content-Type", "application/octet-stream")
	}

	bufio.NewReader(f).WriteTo(c.Writer)
}

// Returns the contents of a directory for a server.
func getServerListDirectory(c *gin.Context) {
	s := GetServer(c.Param("server"))

	d, err := url.QueryUnescape(c.Query("directory"))
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	stats, err := s.Filesystem.ListDirectory(d)
	if err != nil {
		if err.Error() == "readdirent: not a directory" {
			c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
				"error": "The requested directory does not exist.",
			})
			return
		}

		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.JSON(http.StatusOK, stats)
}

type renameFile struct {
	To   string `json:"to"`
	From string `json:"from"`
}

// Renames (or moves) files for a server.
func putServerRenameFiles(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		Root  string       `json:"root"`
		Files []renameFile `json:"files"`
	}
	// BindJSON sends 400 if the request fails, all we need to do is return
	if err := c.BindJSON(&data); err != nil {
		return
	}

	if len(data.Files) == 0 {
		c.AbortWithStatusJSON(http.StatusUnprocessableEntity, gin.H{
			"error": "No files to move or rename were provided.",
		})
		return
	}

	g, ctx := errgroup.WithContext(context.Background())

	// Loop over the array of files passed in and perform the move or rename action against each.
	for _, p := range data.Files {
		pf := path.Join(data.Root, p.From)
		pt := path.Join(data.Root, p.To)

		g.Go(func() error {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
				if err := s.Filesystem.Rename(pf, pt); err != nil {
					// Return nil if the error is an is not exists.
					// NOTE: os.IsNotExist() does not work if the error is wrapped.
					if errors.Is(err, os.ErrNotExist) {
						return nil
					}

					return err
				}

				return nil
			}
		})
	}

	if err := g.Wait(); err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

// Copies a server file.
func postServerCopyFile(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		Location string `json:"location"`
	}
	// BindJSON sends 400 if the request fails, all we need to do is return
	if err := c.BindJSON(&data); err != nil {
		return
	}

	if err := s.Filesystem.Copy(data.Location); err != nil {
		// Check if the file does not exist.
		// NOTE: os.IsNotExist() does not work if the error is wrapped.
		if errors.Is(err, os.ErrNotExist) {
			c.Status(http.StatusNotFound)
			return
		}

		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

// Deletes files from a server.
func postServerDeleteFiles(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		Root  string   `json:"root"`
		Files []string `json:"files"`
	}

	if err := c.BindJSON(&data); err != nil {
		return
	}

	if len(data.Files) == 0 {
		c.AbortWithStatusJSON(http.StatusUnprocessableEntity, gin.H{
			"error": "No files were specified for deletion.",
		})
		return
	}

	g, ctx := errgroup.WithContext(context.Background())

	// Loop over the array of files passed in and delete them. If any of the file deletions
	// fail just abort the process entirely.
	for _, p := range data.Files {
		pi := path.Join(data.Root, p)

		g.Go(func() error {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
				return s.Filesystem.Delete(pi)
			}
		})
	}

	if err := g.Wait(); err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

// Writes the contents of the request to a file on a server.
func postServerWriteFile(c *gin.Context) {
	s := GetServer(c.Param("server"))

	f, err := url.QueryUnescape(c.Query("file"))
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}
	f = "/" + strings.TrimLeft(f, "/")

	if err := s.Filesystem.Writefile(f, c.Request.Body); err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

// Create a directory on a server.
func postServerCreateDirectory(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		Name string `json:"name"`
		Path string `json:"path"`
	}
	// BindJSON sends 400 if the request fails, all we need to do is return
	if err := c.BindJSON(&data); err != nil {
		return
	}

	if err := s.Filesystem.CreateDirectory(data.Name, data.Path); err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

func postServerCompressFiles(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		RootPath string   `json:"root"`
		Files    []string `json:"files"`
	}

	if err := c.BindJSON(&data); err != nil {
		return
	}

	if len(data.Files) == 0 {
		c.AbortWithStatusJSON(http.StatusUnprocessableEntity, gin.H{
			"error": "No files were passed through to be compressed.",
		})
		return
	}

	if !s.Filesystem.HasSpaceAvailable() {
		c.AbortWithStatusJSON(http.StatusConflict, gin.H{
			"error": "This server does not have enough available disk space to generate a compressed archive.",
		})
		return
	}

	f, err := s.Filesystem.CompressFiles(data.RootPath, data.Files)
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.JSON(http.StatusOK, &server.Stat{
		Info:     f,
		Mimetype: "application/tar+gzip",
	})
}

func postServerDecompressFiles(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		RootPath string `json:"root"`
		File     string `json:"file"`
	}

	if err := c.BindJSON(&data); err != nil {
		return
	}

	hasSpace, err := s.Filesystem.SpaceAvailableForDecompression(data.RootPath, data.File)
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	if !hasSpace {
		c.AbortWithStatusJSON(http.StatusConflict, gin.H{
			"error": "This server does not have enough available disk space to decompress this archive.",
		})
		return
	}

	if err := s.Filesystem.DecompressFile(data.RootPath, data.File); err != nil {
		// Check if the file does not exist.
		// NOTE: os.IsNotExist() does not work if the error is wrapped.
		if errors.Is(err, os.ErrNotExist) {
			c.Status(http.StatusNotFound)
			return
		}

		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}
<commit_msg>Update router_server_files.go<commit_after>package router

import (
	"bufio"
	"context"
	"github.com/gin-gonic/gin"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/server"
	"golang.org/x/sync/errgroup"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
)

// Returns the contents of a file on the server.
func getServerFileContents(c *gin.Context) {
	s := GetServer(c.Param("server"))

	p, err := url.QueryUnescape(c.Query("file"))
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}
	p = "/" + strings.TrimLeft(p, "/")

	cleaned, err := s.Filesystem.SafePath(p)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
			"error": "The file requested could not be found.",
		})
		return
	}

	st, err := s.Filesystem.Stat(cleaned)
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	if st.Info.IsDir() {
		c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
			"error": "The requested resource was not found on the system.",
		})
		return
	}

	f, err := os.Open(cleaned)
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}
	defer f.Close()

	c.Header("X-Mime-Type", st.Mimetype)
	c.Header("Content-Length", strconv.Itoa(int(st.Info.Size())))

	// If a download parameter is included in the URL go ahead and attach the necessary headers
	// so that the file can be downloaded.
	if c.Query("download") != "" {
		c.Header("Content-Disposition", "attachment; filename="+st.Info.Name())
		c.Header("Content-Type", "application/octet-stream")
	}

	bufio.NewReader(f).WriteTo(c.Writer)
}

// Returns the contents of a directory for a server.
func getServerListDirectory(c *gin.Context) {
	s := GetServer(c.Param("server"))

	d, err := url.QueryUnescape(c.Query("directory"))
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	stats, err := s.Filesystem.ListDirectory(d)
	if err != nil {
		if err.Error() == "readdirent: not a directory" {
			c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
				"error": "The requested directory does not exist.",
			})
			return
		}

		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.JSON(http.StatusOK, stats)
}

type renameFile struct {
	To   string `json:"to"`
	From string `json:"from"`
}

// Renames (or moves) files for a server.
func putServerRenameFiles(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		Root  string       `json:"root"`
		Files []renameFile `json:"files"`
	}
	// BindJSON sends 400 if the request fails, all we need to do is return
	if err := c.BindJSON(&data); err != nil {
		return
	}

	if len(data.Files) == 0 {
		c.AbortWithStatusJSON(http.StatusUnprocessableEntity, gin.H{
			"error": "No files to move or rename were provided.",
		})
		return
	}

	g, ctx := errgroup.WithContext(context.Background())

	// Loop over the array of files passed in and perform the move or rename action against each.
	for _, p := range data.Files {
		pf := path.Join(data.Root, p.From)
		pt := path.Join(data.Root, p.To)

		g.Go(func() error {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
				if err := s.Filesystem.Rename(pf, pt); err != nil {
					// Return nil if the error is an is not exists.
					// NOTE: os.IsNotExist() does not work if the error is wrapped.
					if errors.Is(err, os.ErrNotExist) {
						return nil
					}

					return err
				}

				return nil
			}
		})
	}

	if err := g.Wait(); err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

// Copies a server file.
func postServerCopyFile(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		Location string `json:"location"`
	}
	// BindJSON sends 400 if the request fails, all we need to do is return
	if err := c.BindJSON(&data); err != nil {
		return
	}

	if err := s.Filesystem.Copy(data.Location); err != nil {
		// Check if the file does not exist.
		// NOTE: os.IsNotExist() does not work if the error is wrapped.
		if errors.Is(err, os.ErrNotExist) {
			c.Status(http.StatusNotFound)
			return
		}

		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

// Deletes files from a server.
func postServerDeleteFiles(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		Root  string   `json:"root"`
		Files []string `json:"files"`
	}

	if err := c.BindJSON(&data); err != nil {
		return
	}

	if len(data.Files) == 0 {
		c.AbortWithStatusJSON(http.StatusUnprocessableEntity, gin.H{
			"error": "No files were specified for deletion.",
		})
		return
	}

	g, ctx := errgroup.WithContext(context.Background())

	// Loop over the array of files passed in and delete them. If any of the file deletions
	// fail just abort the process entirely.
	for _, p := range data.Files {
		pi := path.Join(data.Root, p)

		g.Go(func() error {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
				return s.Filesystem.Delete(pi)
			}
		})
	}

	if err := g.Wait(); err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

// Writes the contents of the request to a file on a server.
func postServerWriteFile(c *gin.Context) {
	s := GetServer(c.Param("server"))

	f, err := url.QueryUnescape(c.Query("file"))
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}
	f = "/" + strings.TrimLeft(f, "/")

	if err := s.Filesystem.Writefile(f, c.Request.Body); err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

// Create a directory on a server.
func postServerCreateDirectory(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		Name string `json:"name"`
		Path string `json:"path"`
	}
	// BindJSON sends 400 if the request fails, all we need to do is return
	if err := c.BindJSON(&data); err != nil {
		return
	}

	if err := s.Filesystem.CreateDirectory(data.Name, data.Path); err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}

func postServerCompressFiles(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		RootPath string   `json:"root"`
		Files    []string `json:"files"`
	}

	if err := c.BindJSON(&data); err != nil {
		return
	}

	if len(data.Files) == 0 {
		c.AbortWithStatusJSON(http.StatusUnprocessableEntity, gin.H{
			"error": "No files were passed through to be compressed.",
		})
		return
	}

	if !s.Filesystem.HasSpaceAvailable() {
		c.AbortWithStatusJSON(http.StatusConflict, gin.H{
			"error": "This server does not have enough available disk space to generate a compressed archive.",
		})
		return
	}

	f, err := s.Filesystem.CompressFiles(data.RootPath, data.Files)
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.JSON(http.StatusOK, &server.Stat{
		Info:     f,
		Mimetype: "application/tar+gzip",
	})
}

func postServerDecompressFiles(c *gin.Context) {
	s := GetServer(c.Param("server"))

	var data struct {
		RootPath string `json:"root"`
		File     string `json:"file"`
	}

	if err := c.BindJSON(&data); err != nil {
		return
	}

	hasSpace, err := s.Filesystem.SpaceAvailableForDecompression(data.RootPath, data.File)
	if err != nil {
		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	if !hasSpace {
		c.AbortWithStatusJSON(http.StatusConflict, gin.H{
			"error": "This server does not have enough available disk space to decompress this archive.",
		})
		return
	}

	if err := s.Filesystem.DecompressFile(data.RootPath, data.File); err != nil {
		// Check if the file does not exist.
		// NOTE: os.IsNotExist() does not work if the error is wrapped.
		if errors.Is(err, os.ErrNotExist) {
			c.Status(http.StatusNotFound)
			return
		}

		TrackedServerError(err, s).AbortWithServerError(c)
		return
	}

	c.Status(http.StatusNoContent)
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 The btcsuite developers
// Copyright (c) 2015-2016 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package chain

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrd/chaincfg/chainhash"
	dcrrpcclient "github.com/decred/dcrd/rpcclient"
)

var requiredChainServerAPI = semver{major: 3, minor: 1, patch: 0}

// RPCClient represents a persistent client connection to a decred RPC server
// for information regarding the current best block chain.
type RPCClient struct {
	*dcrrpcclient.Client
	connConfig  *dcrrpcclient.ConnConfig // Work around unexported field
	chainParams *chaincfg.Params

	enqueueNotification       chan interface{}
	dequeueNotification       chan interface{}
	enqueueVotingNotification chan interface{}
	dequeueVotingNotification chan interface{}

	quit    chan struct{}
	wg      sync.WaitGroup
	started bool
	quitMtx sync.Mutex
}

// NewRPCClient creates a client connection to the server described by the
// connect string.  If disableTLS is false, the remote RPC certificate must be
// provided in the certs slice.  The connection is not established immediately,
// but must be done using the Start method.  If the remote server does not
// operate on the same bitcoin network as described by the passed chain
// parameters, the connection will be disconnected.
func NewRPCClient(chainParams *chaincfg.Params, connect, user, pass string, certs []byte,
	disableTLS bool) (*RPCClient, error) {

	client := &RPCClient{
		connConfig: &dcrrpcclient.ConnConfig{
			Host:                 connect,
			Endpoint:             "ws",
			User:                 user,
			Pass:                 pass,
			Certificates:         certs,
			DisableAutoReconnect: true,
			DisableConnectOnNew:  true,
			DisableTLS:           disableTLS,
		},
		chainParams:               chainParams,
		enqueueNotification:       make(chan interface{}),
		dequeueNotification:       make(chan interface{}),
		enqueueVotingNotification: make(chan interface{}),
		dequeueVotingNotification: make(chan interface{}),
		quit: make(chan struct{}),
	}
	ntfnCallbacks := &dcrrpcclient.NotificationHandlers{
		OnBlockConnected:        client.onBlockConnected,
		OnBlockDisconnected:     client.onBlockDisconnected,
		OnRelevantTxAccepted:    client.onRelevantTxAccepted,
		OnReorganization:        client.onReorganization,
		OnWinningTickets:        client.onWinningTickets,
		OnSpentAndMissedTickets: client.onSpentAndMissedTickets,
		OnStakeDifficulty:       client.onStakeDifficulty,
	}
	rpcClient, err := dcrrpcclient.New(client.connConfig, ntfnCallbacks)
	if err != nil {
		return nil, err
	}
	client.Client = rpcClient
	return client, nil
}

// Start attempts to establish a client connection with the remote server.
// If successful, handler goroutines are started to process notifications
// sent by the server.  After a limited number of connection attempts, this
// function gives up, and therefore will not block forever waiting for the
// connection to be established to a server that may not exist.
func (c *RPCClient) Start(ctx context.Context, retry bool) (err error) {
	err = c.Client.Connect(ctx, retry)
	if err != nil {
		return err
	}

	defer func() {
		if err != nil {
			c.Disconnect()
		}
	}()

	// Verify that the server is running on the expected network.
	net, err := c.Client.GetCurrentNet()
	if err != nil {
		return err
	}
	if net != c.chainParams.Net {
		return errors.New("mismatched networks")
	}

	// Ensure the RPC server has a compatible API version.
	var serverAPI semver
	versions, err := c.Client.Version()
	if err == nil {
		versionResult := versions["dcrdjsonrpcapi"]
		serverAPI = semver{
			major: versionResult.Major,
			minor: versionResult.Minor,
			patch: versionResult.Patch,
		}
	}
	if !semverCompatible(requiredChainServerAPI, serverAPI) {
		return fmt.Errorf("consensus JSON-RPC server does not have a "+
			"compatible API version: advertises %v but require %v",
			serverAPI, requiredChainServerAPI)
	}

	c.quitMtx.Lock()
	c.started = true
	c.quitMtx.Unlock()

	c.wg.Add(2)
	go c.handler()
	go c.handlerVoting()
	return nil
}

// Stop disconnects the client and signals the shutdown of all goroutines
// started by Start.
func (c *RPCClient) Stop() {
	c.quitMtx.Lock()
	select {
	case <-c.quit:
	default:
		close(c.quit)
		c.Client.Shutdown()

		if !c.started {
			close(c.dequeueNotification)
			close(c.dequeueVotingNotification)
		}
	}
	c.quitMtx.Unlock()
}

// WaitForShutdown blocks until both the client has finished disconnecting
// and all handlers have exited.
func (c *RPCClient) WaitForShutdown() {
	c.Client.WaitForShutdown()
	c.wg.Wait()
}

// Notification types.  These are defined here and processed from from reading
// a notificationChan to avoid handling these notifications directly in
// dcrrpcclient callbacks, which isn't very Go-like and doesn't allow
// blocking client calls.
type (
	// blockConnected is a notification for a newly-attached block to the
	// best chain.
	blockConnected struct {
		blockHeader  []byte
		transactions [][]byte
	}

	// blockDisconnected is a notifcation that the block described by the header
	// was reorganized out of the best chain.
	blockDisconnected struct {
		blockHeader []byte
	}

	// relevantTxAccepted is a notification that a transaction accepted by
	// mempool passed the client's transaction filter.
	relevantTxAccepted struct {
		transaction []byte
	}

	// reorganization is a notification that a reorg has happen with the new
	// old and new tip included.
	reorganization struct {
		oldHash   *chainhash.Hash
		oldHeight int64
		newHash   *chainhash.Hash
		newHeight int64
	}

	// winningTickets is a notification with the winning tickets (and the
	// block they are in.
	winningTickets struct {
		blockHash   *chainhash.Hash
		blockHeight int64
		tickets     []*chainhash.Hash
	}

	// missedTickets is a notifcation for tickets that have been missed.
	missedTickets struct {
		blockHash   *chainhash.Hash
		blockHeight int64
		tickets     []*chainhash.Hash
	}

	// stakeDifficulty is a notification for the current stake difficulty.
	stakeDifficulty struct {
		blockHash   *chainhash.Hash
		blockHeight int64
		stakeDiff   int64
	}
)

// notifications returns a channel of parsed notifications sent by the remote
// decred RPC server.  This channel must be continually read or the process
// may abort for running out memory, as unread notifications are queued for
// later reads.
func (c *RPCClient) notifications() <-chan interface{} {
	return c.dequeueNotification
}

// notificationsVoting returns a channel of parsed voting notifications sent
// by the remote RPC server.  This channel must be continually read or the
// process may abort for running out memory, as unread notifications are
// queued for later reads.
func (c *RPCClient) notificationsVoting() <-chan interface{} {
	return c.dequeueVotingNotification
}

func (c *RPCClient) onBlockConnected(header []byte, transactions [][]byte) {
	select {
	case c.enqueueNotification <- blockConnected{
		blockHeader:  header,
		transactions: transactions,
	}:
	case <-c.quit:
	}
}

func (c *RPCClient) onBlockDisconnected(header []byte) {
	select {
	case c.enqueueNotification <- blockDisconnected{
		blockHeader: header,
	}:
	case <-c.quit:
	}
}

func (c *RPCClient) onRelevantTxAccepted(transaction []byte) {
	select {
	case c.enqueueNotification <- relevantTxAccepted{
		transaction: transaction,
	}:
	case <-c.quit:
	}
}

// onReorganization handles reorganization notifications and passes them
// downstream to the notifications queue.
func (c *RPCClient) onReorganization(oldHash *chainhash.Hash, oldHeight int32,
	newHash *chainhash.Hash, newHeight int32) {
	select {
	case c.enqueueNotification <- reorganization{
		oldHash,
		int64(oldHeight),
		newHash,
		int64(newHeight),
	}:
	case <-c.quit:
	}
}

// onWinningTickets handles winning tickets notifications data and passes it
// downstream to the notifications queue.
func (c *RPCClient) onWinningTickets(hash *chainhash.Hash, height int64, tickets []*chainhash.Hash) {
	select {
	case c.enqueueVotingNotification <- winningTickets{
		blockHash:   hash,
		blockHeight: height,
		tickets:     tickets,
	}:
	case <-c.quit:
	}
}

// onSpentAndMissedTickets handles missed tickets notifications data and passes
// it downstream to the notifications queue.
func (c *RPCClient) onSpentAndMissedTickets(blockHash *chainhash.Hash, height int64, sdiff int64, tickets map[chainhash.Hash]bool) {
	var missed []*chainhash.Hash

	// Copy the missed ticket hashes to a slice.
	for ticketHash, spent := range tickets {
		if !spent {
			ticketHash := ticketHash
			missed = append(missed, &ticketHash)
		}
	}

	if len(missed) == 0 {
		return
	}

	select {
	case c.enqueueNotification <- missedTickets{
		blockHash:   blockHash,
		blockHeight: height,
		tickets:     missed,
	}:
	case <-c.quit:
	}
}

// onStakeDifficulty handles stake difficulty notifications data and passes it
// downstream to the notification queue.
func (c *RPCClient) onStakeDifficulty(hash *chainhash.Hash,
	height int64,
	stakeDiff int64) {

	select {
	case c.enqueueNotification <- stakeDifficulty{
		hash,
		height,
		stakeDiff,
	}:
	case <-c.quit:
	}
}

// handler maintains a queue of notifications and the current state (best
// block) of the chain.
func (c *RPCClient) handler() {
	// TODO: Rather than leaving this as an unbounded queue for all types of
	// notifications, try dropping ones where a later enqueued notification
	// can fully invalidate one waiting to be processed.  For example,
	// blockconnected notifications for greater block heights can remove the
	// need to process earlier blockconnected notifications still waiting
	// here.

	var notifications []interface{}
	enqueue := c.enqueueNotification
	var dequeue chan interface{}
	var next interface{}
	pingChan := time.After(time.Minute)
	pingChanReset := make(chan (<-chan time.Time))
out:
	for {
		select {
		case n, ok := <-enqueue:
			if !ok {
				// If no notifications are queued for handling,
				// the queue is finished.
				if len(notifications) == 0 {
					break out
				}
				// nil channel so no more reads can occur.
				enqueue = nil
				continue
			}
			if len(notifications) == 0 {
				next = n
				dequeue = c.dequeueNotification
			}
			notifications = append(notifications, n)
			pingChan = time.After(time.Minute)

		case dequeue <- next:
			notifications[0] = nil
			notifications = notifications[1:]
			if len(notifications) != 0 {
				next = notifications[0]
			} else {
				// If no more notifications can be enqueued, the
				// queue is finished.
				if enqueue == nil {
					break out
				}
				dequeue = nil
			}

		case <-pingChan:
			// No notifications were received in the last 60s. Ensure the
			// connection is still active by making a new request to the server.
			//
			// This MUST wait for the response in a new goroutine so as to not
			// block channel sends enqueueing more notifications.  Doing so
			// would cause a deadlock and after the timeout expires, the client
			// would be shut down.
			//
			// TODO: A minute timeout is used to prevent the handler loop from
			// blocking here forever, but this is much larger than it needs to
			// be due to dcrd processing websocket requests synchronously (see
			// https://github.com/btcsuite/btcd/issues/504).  Decrease this to
			// something saner like 3s when the above issue is fixed.
			type sessionResult struct {
				err error
			}
			sessionResponse := make(chan sessionResult, 1)
			go func() {
				_, err := c.Session()
				sessionResponse <- sessionResult{err}
			}()
			go func() {
				select {
				case resp := <-sessionResponse:
					if resp.err != nil {
						log.Errorf("Failed to receive session "+
							"result: %v", resp.err)
						c.Stop()
					}
					pingChanReset <- time.After(time.Minute)

				case <-time.After(time.Minute):
					log.Errorf("Timeout waiting for session RPC")
					c.Stop()
				}
			}()

		case ch := <-pingChanReset:
			pingChan = ch

		case <-c.quit:
			break out
		}
	}

	c.Stop()
	close(c.dequeueNotification)
	c.wg.Done()
}

// handler maintains a queue of notifications and the current state (best
// block) of the chain.
func (c *RPCClient) handlerVoting() {
	var notifications []interface{}
	enqueue := c.enqueueVotingNotification
	var dequeue chan interface{}
	var next interface{}
out:
	for {
		select {
		case n, ok := <-enqueue:
			if !ok {
				// If no notifications are queued for handling,
				// the queue is finished.
				if len(notifications) == 0 {
					break out
				}
				// nil channel so no more reads can occur.
				enqueue = nil
				continue
			}
			if len(notifications) == 0 {
				next = n
				dequeue = c.dequeueVotingNotification
			}
			notifications = append(notifications, n)

		case dequeue <- next:
			notifications[0] = nil
			notifications = notifications[1:]
			if len(notifications) != 0 {
				next = notifications[0]
			} else {
				// If no more notifications can be enqueued, the
				// queue is finished.
				if enqueue == nil {
					break out
				}
				dequeue = nil
			}

		case <-c.quit:
			break out
		}
	}
	close(c.dequeueVotingNotification)
	c.wg.Done()
}

// POSTClient creates the equivalent HTTP POST dcrrpcclient.Client.
func (c *RPCClient) POSTClient() (*dcrrpcclient.Client, error) {
	configCopy := *c.connConfig
	configCopy.HTTPPostMode = true
	return dcrrpcclient.New(&configCopy, nil)
}
<commit_msg>chain: replace reference to bitcoin with decred (#1068)<commit_after>// Copyright (c) 2013-2015 The btcsuite developers
// Copyright (c) 2015-2016 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package chain

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrd/chaincfg/chainhash"
	dcrrpcclient "github.com/decred/dcrd/rpcclient"
)

var requiredChainServerAPI = semver{major: 3, minor: 1, patch: 0}

// RPCClient represents a persistent client connection to a decred RPC server
// for information regarding the current best block chain.
type RPCClient struct {
	*dcrrpcclient.Client
	connConfig  *dcrrpcclient.ConnConfig // Work around unexported field
	chainParams *chaincfg.Params

	enqueueNotification       chan interface{}
	dequeueNotification       chan interface{}
	enqueueVotingNotification chan interface{}
	dequeueVotingNotification chan interface{}

	quit    chan struct{}
	wg      sync.WaitGroup
	started bool
	quitMtx sync.Mutex
}

// NewRPCClient creates a client connection to the server described by the
// connect string.  If disableTLS is false, the remote RPC certificate must be
// provided in the certs slice.  The connection is not established immediately,
// but must be done using the Start method.  If the remote server does not
// operate on the same decred network as described by the passed chain
// parameters, the connection will be disconnected.
func NewRPCClient(chainParams *chaincfg.Params, connect, user, pass string, certs []byte,
	disableTLS bool) (*RPCClient, error) {

	client := &RPCClient{
		connConfig: &dcrrpcclient.ConnConfig{
			Host:                 connect,
			Endpoint:             "ws",
			User:                 user,
			Pass:                 pass,
			Certificates:         certs,
			DisableAutoReconnect: true,
			DisableConnectOnNew:  true,
			DisableTLS:           disableTLS,
		},
		chainParams:               chainParams,
		enqueueNotification:       make(chan interface{}),
		dequeueNotification:       make(chan interface{}),
		enqueueVotingNotification: make(chan interface{}),
		dequeueVotingNotification: make(chan interface{}),
		quit: make(chan struct{}),
	}
	ntfnCallbacks := &dcrrpcclient.NotificationHandlers{
		OnBlockConnected:        client.onBlockConnected,
		OnBlockDisconnected:     client.onBlockDisconnected,
		OnRelevantTxAccepted:    client.onRelevantTxAccepted,
		OnReorganization:        client.onReorganization,
		OnWinningTickets:        client.onWinningTickets,
		OnSpentAndMissedTickets: client.onSpentAndMissedTickets,
		OnStakeDifficulty:       client.onStakeDifficulty,
	}
	rpcClient, err := dcrrpcclient.New(client.connConfig, ntfnCallbacks)
	if err != nil {
		return nil, err
	}
	client.Client = rpcClient
	return client, nil
}

// Start attempts to establish a client connection with the remote server.
// If successful, handler goroutines are started to process notifications
// sent by the server.  After a limited number of connection attempts, this
// function gives up, and therefore will not block forever waiting for the
// connection to be established to a server that may not exist.
func (c *RPCClient) Start(ctx context.Context, retry bool) (err error) {
	err = c.Client.Connect(ctx, retry)
	if err != nil {
		return err
	}

	defer func() {
		if err != nil {
			c.Disconnect()
		}
	}()

	// Verify that the server is running on the expected network.
	net, err := c.Client.GetCurrentNet()
	if err != nil {
		return err
	}
	if net != c.chainParams.Net {
		return errors.New("mismatched networks")
	}

	// Ensure the RPC server has a compatible API version.
	var serverAPI semver
	versions, err := c.Client.Version()
	if err == nil {
		versionResult := versions["dcrdjsonrpcapi"]
		serverAPI = semver{
			major: versionResult.Major,
			minor: versionResult.Minor,
			patch: versionResult.Patch,
		}
	}
	if !semverCompatible(requiredChainServerAPI, serverAPI) {
		return fmt.Errorf("consensus JSON-RPC server does not have a "+
			"compatible API version: advertises %v but require %v",
			serverAPI, requiredChainServerAPI)
	}

	c.quitMtx.Lock()
	c.started = true
	c.quitMtx.Unlock()

	c.wg.Add(2)
	go c.handler()
	go c.handlerVoting()
	return nil
}

// Stop disconnects the client and signals the shutdown of all goroutines
// started by Start.
func (c *RPCClient) Stop() {
	c.quitMtx.Lock()
	select {
	case <-c.quit:
	default:
		close(c.quit)
		c.Client.Shutdown()

		if !c.started {
			close(c.dequeueNotification)
			close(c.dequeueVotingNotification)
		}
	}
	c.quitMtx.Unlock()
}

// WaitForShutdown blocks until both the client has finished disconnecting
// and all handlers have exited.
func (c *RPCClient) WaitForShutdown() {
	c.Client.WaitForShutdown()
	c.wg.Wait()
}

// Notification types.  These are defined here and processed from from reading
// a notificationChan to avoid handling these notifications directly in
// dcrrpcclient callbacks, which isn't very Go-like and doesn't allow
// blocking client calls.
type (
	// blockConnected is a notification for a newly-attached block to the
	// best chain.
	blockConnected struct {
		blockHeader  []byte
		transactions [][]byte
	}

	// blockDisconnected is a notifcation that the block described by the header
	// was reorganized out of the best chain.
	blockDisconnected struct {
		blockHeader []byte
	}

	// relevantTxAccepted is a notification that a transaction accepted by
	// mempool passed the client's transaction filter.
	relevantTxAccepted struct {
		transaction []byte
	}

	// reorganization is a notification that a reorg has happen with the new
	// old and new tip included.
	reorganization struct {
		oldHash   *chainhash.Hash
		oldHeight int64
		newHash   *chainhash.Hash
		newHeight int64
	}

	// winningTickets is a notification with the winning tickets (and the
	// block they are in.
	winningTickets struct {
		blockHash   *chainhash.Hash
		blockHeight int64
		tickets     []*chainhash.Hash
	}

	// missedTickets is a notifcation for tickets that have been missed.
	missedTickets struct {
		blockHash   *chainhash.Hash
		blockHeight int64
		tickets     []*chainhash.Hash
	}

	// stakeDifficulty is a notification for the current stake difficulty.
	stakeDifficulty struct {
		blockHash   *chainhash.Hash
		blockHeight int64
		stakeDiff   int64
	}
)

// notifications returns a channel of parsed notifications sent by the remote
// decred RPC server.  This channel must be continually read or the process
// may abort for running out memory, as unread notifications are queued for
// later reads.
func (c *RPCClient) notifications() <-chan interface{} {
	return c.dequeueNotification
}

// notificationsVoting returns a channel of parsed voting notifications sent
// by the remote RPC server.  This channel must be continually read or the
// process may abort for running out memory, as unread notifications are
// queued for later reads.
func (c *RPCClient) notificationsVoting() <-chan interface{} {
	return c.dequeueVotingNotification
}

func (c *RPCClient) onBlockConnected(header []byte, transactions [][]byte) {
	select {
	case c.enqueueNotification <- blockConnected{
		blockHeader:  header,
		transactions: transactions,
	}:
	case <-c.quit:
	}
}

func (c *RPCClient) onBlockDisconnected(header []byte) {
	select {
	case c.enqueueNotification <- blockDisconnected{
		blockHeader: header,
	}:
	case <-c.quit:
	}
}

func (c *RPCClient) onRelevantTxAccepted(transaction []byte) {
	select {
	case c.enqueueNotification <- relevantTxAccepted{
		transaction: transaction,
	}:
	case <-c.quit:
	}
}

// onReorganization handles reorganization notifications and passes them
// downstream to the notifications queue.
func (c *RPCClient) onReorganization(oldHash *chainhash.Hash, oldHeight int32,
	newHash *chainhash.Hash, newHeight int32) {
	select {
	case c.enqueueNotification <- reorganization{
		oldHash,
		int64(oldHeight),
		newHash,
		int64(newHeight),
	}:
	case <-c.quit:
	}
}

// onWinningTickets handles winning tickets notifications data and passes it
// downstream to the notifications queue.
func (c *RPCClient) onWinningTickets(hash *chainhash.Hash, height int64, tickets []*chainhash.Hash) {
	select {
	case c.enqueueVotingNotification <- winningTickets{
		blockHash:   hash,
		blockHeight: height,
		tickets:     tickets,
	}:
	case <-c.quit:
	}
}

// onSpentAndMissedTickets handles missed tickets notifications data and passes
// it downstream to the notifications queue.
func (c *RPCClient) onSpentAndMissedTickets(blockHash *chainhash.Hash, height int64, sdiff int64, tickets map[chainhash.Hash]bool) {
	var missed []*chainhash.Hash

	// Copy the missed ticket hashes to a slice.
	for ticketHash, spent := range tickets {
		if !spent {
			ticketHash := ticketHash
			missed = append(missed, &ticketHash)
		}
	}

	if len(missed) == 0 {
		return
	}

	select {
	case c.enqueueNotification <- missedTickets{
		blockHash:   blockHash,
		blockHeight: height,
		tickets:     missed,
	}:
	case <-c.quit:
	}
}

// onStakeDifficulty handles stake difficulty notifications data and passes it
// downstream to the notification queue.
func (c *RPCClient) onStakeDifficulty(hash *chainhash.Hash,
	height int64,
	stakeDiff int64) {

	select {
	case c.enqueueNotification <- stakeDifficulty{
		hash,
		height,
		stakeDiff,
	}:
	case <-c.quit:
	}
}

// handler maintains a queue of notifications and the current state (best
// block) of the chain.
func (c *RPCClient) handler() {
	// TODO: Rather than leaving this as an unbounded queue for all types of
	// notifications, try dropping ones where a later enqueued notification
	// can fully invalidate one waiting to be processed.  For example,
	// blockconnected notifications for greater block heights can remove the
	// need to process earlier blockconnected notifications still waiting
	// here.

	var notifications []interface{}
	enqueue := c.enqueueNotification
	var dequeue chan interface{}
	var next interface{}
	pingChan := time.After(time.Minute)
	pingChanReset := make(chan (<-chan time.Time))
out:
	for {
		select {
		case n, ok := <-enqueue:
			if !ok {
				// If no notifications are queued for handling,
				// the queue is finished.
				if len(notifications) == 0 {
					break out
				}
				// nil channel so no more reads can occur.
				enqueue = nil
				continue
			}
			if len(notifications) == 0 {
				next = n
				dequeue = c.dequeueNotification
			}
			notifications = append(notifications, n)
			pingChan = time.After(time.Minute)

		case dequeue <- next:
			notifications[0] = nil
			notifications = notifications[1:]
			if len(notifications) != 0 {
				next = notifications[0]
			} else {
				// If no more notifications can be enqueued, the
				// queue is finished.
				if enqueue == nil {
					break out
				}
				dequeue = nil
			}

		case <-pingChan:
			// No notifications were received in the last 60s. Ensure the
			// connection is still active by making a new request to the server.
			//
			// This MUST wait for the response in a new goroutine so as to not
			// block channel sends enqueueing more notifications.  Doing so
			// would cause a deadlock and after the timeout expires, the client
			// would be shut down.
			//
			// TODO: A minute timeout is used to prevent the handler loop from
			// blocking here forever, but this is much larger than it needs to
			// be due to dcrd processing websocket requests synchronously (see
			// https://github.com/btcsuite/btcd/issues/504).  Decrease this to
			// something saner like 3s when the above issue is fixed.
			type sessionResult struct {
				err error
			}
			sessionResponse := make(chan sessionResult, 1)
			go func() {
				_, err := c.Session()
				sessionResponse <- sessionResult{err}
			}()
			go func() {
				select {
				case resp := <-sessionResponse:
					if resp.err != nil {
						log.Errorf("Failed to receive session "+
							"result: %v", resp.err)
						c.Stop()
					}
					pingChanReset <- time.After(time.Minute)

				case <-time.After(time.Minute):
					log.Errorf("Timeout waiting for session RPC")
					c.Stop()
				}
			}()

		case ch := <-pingChanReset:
			pingChan = ch

		case <-c.quit:
			break out
		}
	}

	c.Stop()
	close(c.dequeueNotification)
	c.wg.Done()
}

// handler maintains a queue of notifications and the current state (best
// block) of the chain.
func (c *RPCClient) handlerVoting() {
	var notifications []interface{}
	enqueue := c.enqueueVotingNotification
	var dequeue chan interface{}
	var next interface{}
out:
	for {
		select {
		case n, ok := <-enqueue:
			if !ok {
				// If no notifications are queued for handling,
				// the queue is finished.
				if len(notifications) == 0 {
					break out
				}
				// nil channel so no more reads can occur.
				enqueue = nil
				continue
			}
			if len(notifications) == 0 {
				next = n
				dequeue = c.dequeueVotingNotification
			}
			notifications = append(notifications, n)

		case dequeue <- next:
			notifications[0] = nil
			notifications = notifications[1:]
			if len(notifications) != 0 {
				next = notifications[0]
			} else {
				// If no more notifications can be enqueued, the
				// queue is finished.
				if enqueue == nil {
					break out
				}
				dequeue = nil
			}

		case <-c.quit:
			break out
		}
	}
	close(c.dequeueVotingNotification)
	c.wg.Done()
}

// POSTClient creates the equivalent HTTP POST dcrrpcclient.Client.
func (c *RPCClient) POSTClient() (*dcrrpcclient.Client, error) {
	configCopy := *c.connConfig
	configCopy.HTTPPostMode = true
	return dcrrpcclient.New(&configCopy, nil)
}
<|endoftext|>


<commit_before>package socket

import (
	"net"
	"time"

	"github.com/davyxu/cellnet"
)

type socketConnector struct {
	*peerProfile
	*sessionMgr

	conn net.Conn

	autoReconnectSec int // 重连间隔时间, 0为不重连

	closeSignal chan bool

	working bool // 重入锁

	defaultSes cellnet.Session
}

// 自动重连间隔=0不重连
func (self *socketConnector) SetAutoReconnectSec(sec int) {
	self.autoReconnectSec = sec
}

func (self *socketConnector) Start(address string) cellnet.Peer {

	if self.working {
		return self
	}

	go self.connect(address)

	self.PostData(NewPeerEvent(Event_PeerStart, self))

	return self
}

func (self *socketConnector) connect(address string) {
	self.working = true

	for {

		// 开始连接
		cn, err := net.Dial("tcp", address)

		// 连不上
		if err != nil {

			log.Errorf("#connect failed(%s) %v", self.name, err.Error())

			// 没重连就退出
			if self.autoReconnectSec == 0 {
				break
			}

			// 有重连就等待
			time.Sleep(time.Duration(self.autoReconnectSec) * time.Second)

			// 继续连接
			continue
		}

		// 连上了, 记录连接
		self.conn = cn

		// 创建Session
		ses := newSession(NewPacketStream(cn), self.EventQueue, self)
		self.sessionMgr.Add(ses)
		self.defaultSes = ses

		log.Infof("#connected(%s) %s sid: %d", self.name, address, ses.id)

		// 内部断开回调
		ses.OnClose = func() {
			self.sessionMgr.Remove(ses)
			self.closeSignal <- true
		}

		// 抛出事件
		self.PostData(NewSessionEvent(Event_SessionConnected, ses, nil))

		if <-self.closeSignal {

			self.conn = nil

			// 没重连就退出
			if self.autoReconnectSec == 0 {
				break
			}

			// 有重连就等待
			time.Sleep(time.Duration(self.autoReconnectSec) * time.Second)

			// 继续连接
			continue

		}

	}

	self.working = false
}

func (self *socketConnector) Stop() {

	if self.conn != nil {

		self.PostData(NewPeerEvent(Event_PeerStop, self))

		self.conn.Close()
	}

}

func (self *socketConnector) DefaultSession() cellnet.Session {
	return self.defaultSes
}

func NewConnector(pipe cellnet.EventPipe) cellnet.Peer {
	self := &socketConnector{
		sessionMgr:  newSessionManager(),
		peerProfile: newPeerProfile(pipe.AddQueue()),
		closeSignal: make(chan bool),
	}

	self.PostData(NewPeerEvent(Event_PeerInit, self))

	return self
}
<commit_msg>连接/重连超过3次, 停止反复输出日志<commit_after>package socket

import (
	"net"
	"time"

	"github.com/davyxu/cellnet"
)

type socketConnector struct {
	*peerProfile
	*sessionMgr

	conn net.Conn

	autoReconnectSec int // 重连间隔时间, 0为不重连

	closeSignal chan bool

	working bool // 重入锁

	defaultSes cellnet.Session

	tryConnTimes int // 尝试连接次数
}

// 自动重连间隔=0不重连
func (self *socketConnector) SetAutoReconnectSec(sec int) {
	self.autoReconnectSec = sec
}

func (self *socketConnector) Start(address string) cellnet.Peer {

	if self.working {
		return self
	}

	go self.connect(address)

	self.PostData(NewPeerEvent(Event_PeerStart, self))

	return self
}

const reportConnectFailedLimitTimes = 3

func (self *socketConnector) connect(address string) {
	self.working = true

	for {

		self.tryConnTimes++

		// 开始连接
		cn, err := net.Dial("tcp", address)

		// 连不上
		if err != nil {

			if self.tryConnTimes <= reportConnectFailedLimitTimes {
				log.Errorf("#connect failed(%s) %v", self.name, err.Error())
			}

			if self.tryConnTimes == reportConnectFailedLimitTimes {
				log.Errorf("continue reconnecting, but mute log")
			}

			// 没重连就退出
			if self.autoReconnectSec == 0 {
				break
			}

			// 有重连就等待
			time.Sleep(time.Duration(self.autoReconnectSec) * time.Second)

			// 继续连接
			continue
		}

		self.tryConnTimes = 0

		// 连上了, 记录连接
		self.conn = cn

		// 创建Session
		ses := newSession(NewPacketStream(cn), self.EventQueue, self)
		self.sessionMgr.Add(ses)
		self.defaultSes = ses

		log.Infof("#connected(%s) %s sid: %d", self.name, address, ses.id)

		// 内部断开回调
		ses.OnClose = func() {
			self.sessionMgr.Remove(ses)
			self.closeSignal <- true
		}

		// 抛出事件
		self.PostData(NewSessionEvent(Event_SessionConnected, ses, nil))

		if <-self.closeSignal {

			self.conn = nil

			// 没重连就退出
			if self.autoReconnectSec == 0 {
				break
			}

			// 有重连就等待
			time.Sleep(time.Duration(self.autoReconnectSec) * time.Second)

			// 继续连接
			continue

		}

	}

	self.working = false
}

func (self *socketConnector) Stop() {

	if self.conn != nil {

		self.PostData(NewPeerEvent(Event_PeerStop, self))

		self.conn.Close()
	}

}

func (self *socketConnector) DefaultSession() cellnet.Session {
	return self.defaultSes
}

func NewConnector(pipe cellnet.EventPipe) cellnet.Peer {
	self := &socketConnector{
		sessionMgr:  newSessionManager(),
		peerProfile: newPeerProfile(pipe.AddQueue()),
		closeSignal: make(chan bool),
	}

	self.PostData(NewPeerEvent(Event_PeerInit, self))

	return self
}
<|endoftext|>


<commit_before>package signalfx

import (
	"math"
	"sync"
	"sync/atomic"

	"github.com/zvelo/go-signalfx/sfxproto"
)

var (
	// BucketMetricCount represents the count of datapoints seen
	BucketMetricCount = 1
	// BucketMetricMin represents the smallest datapoint seen
	BucketMetricMin = 2
	// BucketMetricMax represents the largest datapoint seen
	BucketMetricMax = 3
	// BucketMetricSum represents the sum of all datapoints seen
	BucketMetricSum = 4
	// BucketMetricSumOfSquares represents the sum of squares of all datapoints seen
	BucketMetricSumOfSquares = 5
)

// A Bucket trakcs groups of values, reporting metrics as gauges and
// resetting each time it reports. All operations on Buckets are goroutine safe.
type Bucket struct {
	metric             string
	dimensions         map[string]string
	count              uint64
	countMetric        Counter
	min                int64
	minMetric          Gauge
	max                int64
	maxMetric          Gauge
	sum                int64
	sumMetric          Gauge
	sumOfSquares       int64
	sumOfSquaresMetric Gauge
	mu                 sync.Mutex
	disabledMetrics    map[int]bool
}

func (b *Bucket) lock() {
	b.mu.Lock()
}

func (b *Bucket) unlock() {
	b.mu.Unlock()
}

// Clone makes a deep copy of a Bucket
func (b *Bucket) Clone() *Bucket {
	// use the mutex to keep all operations part of the same transaction
	b.lock()
	defer b.unlock()

	return &Bucket{
		metric:          b.metric,                                  // can't use Metric() since we already have a lock
		dimensions:      sfxproto.Dimensions(b.dimensions).Clone(), // can't use Dimensions() since we already have a lock
		count:           b.Count(),
		min:             b.Min(),
		max:             b.Max(),
		sum:             b.Sum(),
		sumOfSquares:    b.SumOfSquares(),
		disabledMetrics: b.disabledMetrics,
	}
}

// Disable disables the given metrics for this bucket.
// They will be collected, but not reported.
func (b *Bucket) Disable(metrics ...int) {
	b.lock()
	defer b.unlock()
	for _, metric := range metrics {
		b.disabledMetrics[metric] = true
	}
}

// Equal returns whether two buckets are exactly equal
func (b *Bucket) Equal(r *Bucket) bool {
	// lock the state of both buckets as this operation is effectively a
	// transaction on an exact state of a whole bucket

	b.lock()
	defer b.unlock()

	r.lock()
	defer r.unlock()

	if b.metric != r.metric {
		return false
	}

	if !sfxproto.Dimensions(b.dimensions).Equal(r.dimensions) {
		return false
	}

	if b.Count() != r.Count() {
		return false
	}

	if b.Min() != r.Min() {
		return false
	}

	if b.Max() != r.Max() {
		return false
	}

	if b.Sum() != r.Sum() {
		return false
	}

	if b.SumOfSquares() != r.SumOfSquares() {
		return false
	}

	return true
}

// Metric returns the metric name of the Bucket
func (b *Bucket) Metric() string {
	b.lock()
	defer b.unlock()

	return b.metric
}

// SetMetric sets the metric name of the Bucket
func (b *Bucket) SetMetric(name string) {
	b.lock()
	defer b.unlock()

	b.metric = name
}

// Dimensions returns a copy of the dimensions of the Bucket. Changes are not
// reflected inside the Bucket itself.
func (b *Bucket) Dimensions() map[string]string {
	b.lock()
	defer b.unlock()

	return sfxproto.Dimensions(b.dimensions).Clone()
}

// SetDimension adds or overwrites the dimension at key with value. If the key
// or value is empty, no changes are made
func (b *Bucket) SetDimension(key, value string) {
	if key == "" || value == "" {
		return
	}

	b.lock()
	defer b.unlock()

	b.dimensions[key] = value
}

// SetDimensions adds or overwrites multiple dimensions. Because the passed in
// dimensions can not be locked by this method, it is important that the caller
// ensures its state does not change for the duration of the operation.
func (b *Bucket) SetDimensions(dims map[string]string) {
	for key, value := range dims {
		b.SetDimension(key, value)
	}
}

// RemoveDimension removes one or more dimensions with the given keys
func (b *Bucket) RemoveDimension(keys ...string) {
	b.lock()
	defer b.unlock()

	for _, key := range keys {
		delete(b.dimensions, key)
	}
}

// Count returns the number of items added to the Bucket
func (b *Bucket) Count() uint64 {
	return atomic.LoadUint64(&b.count)
}

// Min returns the lowest item added to the Bucket
func (b *Bucket) Min() int64 {
	return atomic.LoadInt64(&b.min)
}

// Max returns the highest item added to the Bucket
func (b *Bucket) Max() int64 {
	return atomic.LoadInt64(&b.max)
}

// Sum returns the sum of all items added to the Bucket
func (b *Bucket) Sum() int64 {
	return atomic.LoadInt64(&b.sum)
}

// SumOfSquares returns the sum of the square of all items added to the Bucket
func (b *Bucket) SumOfSquares() int64 {
	return atomic.LoadInt64(&b.sumOfSquares)
}

// NewBucket creates a new Bucket. Because the passed in dimensions can not be
// locked by this method, it is important that the caller ensures its state does
// not change for the duration of the operation.
func NewBucket(metric string, dimensions map[string]string) *Bucket {
	return &Bucket{
		metric:          metric,
		dimensions:      sfxproto.Dimensions(dimensions).Clone(),
		min:             math.MaxInt64,
		max:             math.MinInt64,
		disabledMetrics: make(map[int]bool, 0),
	}
}

// Add an item to the Bucket, later reporting the result in the next report
// cycle.
func (b *Bucket) Add(val int64) {
	// wrap all of these changes in the mutex lock since they need to occur as a
	// transaction, not a set of atomic operations
	b.lock()
	defer b.unlock()

	// still use atomic though, so that the atomic "getters" will never read
	// from an inconsistent state
	_ = atomic.AddUint64(&b.count, 1)
	atomic.AddInt64(&b.sum, val)
	atomic.AddInt64(&b.sumOfSquares, val*val)

	if b.Min() > val {
		atomic.StoreInt64(&b.min, val)
	}

	if b.Max() < val {
		atomic.StoreInt64(&b.max, val)
	}
}

func (b *Bucket) dimFor(defaultDims map[string]string, rollup string) map[string]string {
	b.lock()
	defer b.unlock()

	dims := sfxproto.Dimensions(defaultDims).Append(b.dimensions)
	dims["rollup"] = rollup

	return dims
}

// DataPoints returns a DataPoints object with DataPoint values for
// Count, Sum, SumOfSquares, Min and Max (if set). Note that this
// resets all values.  If no values have been added to the bucket
// since the last report, it returns 0 for count, sum and
// sum-of-squares, omitting max and min.  If the count is higher than
// may be represented in an int64, then the count will be omitted.
func (b *Bucket) DataPoints() []DataPoint {
	dps := make([]DataPoint, 0, 5)
	cnt := atomic.SwapUint64(&b.count, 0)
	min := atomic.SwapInt64(&b.min, math.MaxInt64)
	max := atomic.SwapInt64(&b.max, math.MinInt64)
	sum := atomic.SwapInt64(&b.sum, 0)
	sos := atomic.SwapInt64(&b.sumOfSquares, 0)
	if cnt != 0 {
		if !b.disabledMetrics[BucketMetricMin] {
			dp := DataPoint{
				Metric:     b.metric,
				Dimensions: b.dimensions,
				Type:       GaugeType,
				Value:      min,
			}
			dps = append(dps, dp)
		}
		if !b.disabledMetrics[BucketMetricMax] {
			dp := DataPoint{
				Metric:     b.metric,
				Dimensions: b.dimensions,
				Type:       GaugeType,
				Value:      max,
			}
			dps = append(dps, dp)
		}
	}
	if !b.disabledMetrics[BucketMetricCount] && cnt <= math.MaxInt64 {
		dp := DataPoint{
			Metric:     b.metric,
			Dimensions: b.dimensions,
			Type:       CounterType,
			Value:      int64(cnt),
		}
		dps = append(dps, dp)
	}
	if !b.disabledMetrics[BucketMetricSum] {
		dp := DataPoint{
			Metric:     b.metric,
			Dimensions: b.dimensions,
			Type:       GaugeType,
			Value:      sum,
		}
		dps = append(dps, dp)
	}
	if !b.disabledMetrics[BucketMetricSumOfSquares] {
		dp := DataPoint{
			Metric:     b.metric,
			Dimensions: b.dimensions,
			Type:       GaugeType,
			Value:      sos,
		}
		dps = append(dps, dp)
	}

	return dps
}
<commit_msg>Changed constants in bucket to be actual consts<commit_after>package signalfx

import (
	"math"
	"sync"
	"sync/atomic"

	"github.com/zvelo/go-signalfx/sfxproto"
)

const (
	// BucketMetricCount represents the count of datapoints seen
	BucketMetricCount = iota
	// BucketMetricMin represents the smallest datapoint seen
	BucketMetricMin = iota
	// BucketMetricMax represents the largest datapoint seen
	BucketMetricMax = iota
	// BucketMetricSum represents the sum of all datapoints seen
	BucketMetricSum = iota
	// BucketMetricSumOfSquares represents the sum of squares of all datapoints seen
	BucketMetricSumOfSquares = iota
)

// A Bucket trakcs groups of values, reporting metrics as gauges and
// resetting each time it reports. All operations on Buckets are goroutine safe.
type Bucket struct {
	metric             string
	dimensions         map[string]string
	count              uint64
	countMetric        Counter
	min                int64
	minMetric          Gauge
	max                int64
	maxMetric          Gauge
	sum                int64
	sumMetric          Gauge
	sumOfSquares       int64
	sumOfSquaresMetric Gauge
	mu                 sync.Mutex
	disabledMetrics    map[int]bool
}

func (b *Bucket) lock() {
	b.mu.Lock()
}

func (b *Bucket) unlock() {
	b.mu.Unlock()
}

// Clone makes a deep copy of a Bucket
func (b *Bucket) Clone() *Bucket {
	// use the mutex to keep all operations part of the same transaction
	b.lock()
	defer b.unlock()

	return &Bucket{
		metric:          b.metric,                                  // can't use Metric() since we already have a lock
		dimensions:      sfxproto.Dimensions(b.dimensions).Clone(), // can't use Dimensions() since we already have a lock
		count:           b.Count(),
		min:             b.Min(),
		max:             b.Max(),
		sum:             b.Sum(),
		sumOfSquares:    b.SumOfSquares(),
		disabledMetrics: b.disabledMetrics,
	}
}

// Disable disables the given metrics for this bucket.
// They will be collected, but not reported.
func (b *Bucket) Disable(metrics ...int) {
	b.lock()
	defer b.unlock()
	for _, metric := range metrics {
		b.disabledMetrics[metric] = true
	}
}

// Equal returns whether two buckets are exactly equal
func (b *Bucket) Equal(r *Bucket) bool {
	// lock the state of both buckets as this operation is effectively a
	// transaction on an exact state of a whole bucket

	b.lock()
	defer b.unlock()

	r.lock()
	defer r.unlock()

	if b.metric != r.metric {
		return false
	}

	if !sfxproto.Dimensions(b.dimensions).Equal(r.dimensions) {
		return false
	}

	if b.Count() != r.Count() {
		return false
	}

	if b.Min() != r.Min() {
		return false
	}

	if b.Max() != r.Max() {
		return false
	}

	if b.Sum() != r.Sum() {
		return false
	}

	if b.SumOfSquares() != r.SumOfSquares() {
		return false
	}

	return true
}

// Metric returns the metric name of the Bucket
func (b *Bucket) Metric() string {
	b.lock()
	defer b.unlock()

	return b.metric
}

// SetMetric sets the metric name of the Bucket
func (b *Bucket) SetMetric(name string) {
	b.lock()
	defer b.unlock()

	b.metric = name
}

// Dimensions returns a copy of the dimensions of the Bucket. Changes are not
// reflected inside the Bucket itself.
func (b *Bucket) Dimensions() map[string]string {
	b.lock()
	defer b.unlock()

	return sfxproto.Dimensions(b.dimensions).Clone()
}

// SetDimension adds or overwrites the dimension at key with value. If the key
// or value is empty, no changes are made
func (b *Bucket) SetDimension(key, value string) {
	if key == "" || value == "" {
		return
	}

	b.lock()
	defer b.unlock()

	b.dimensions[key] = value
}

// SetDimensions adds or overwrites multiple dimensions. Because the passed in
// dimensions can not be locked by this method, it is important that the caller
// ensures its state does not change for the duration of the operation.
func (b *Bucket) SetDimensions(dims map[string]string) {
	for key, value := range dims {
		b.SetDimension(key, value)
	}
}

// RemoveDimension removes one or more dimensions with the given keys
func (b *Bucket) RemoveDimension(keys ...string) {
	b.lock()
	defer b.unlock()

	for _, key := range keys {
		delete(b.dimensions, key)
	}
}

// Count returns the number of items added to the Bucket
func (b *Bucket) Count() uint64 {
	return atomic.LoadUint64(&b.count)
}

// Min returns the lowest item added to the Bucket
func (b *Bucket) Min() int64 {
	return atomic.LoadInt64(&b.min)
}

// Max returns the highest item added to the Bucket
func (b *Bucket) Max() int64 {
	return atomic.LoadInt64(&b.max)
}

// Sum returns the sum of all items added to the Bucket
func (b *Bucket) Sum() int64 {
	return atomic.LoadInt64(&b.sum)
}

// SumOfSquares returns the sum of the square of all items added to the Bucket
func (b *Bucket) SumOfSquares() int64 {
	return atomic.LoadInt64(&b.sumOfSquares)
}

// NewBucket creates a new Bucket. Because the passed in dimensions can not be
// locked by this method, it is important that the caller ensures its state does
// not change for the duration of the operation.
func NewBucket(metric string, dimensions map[string]string) *Bucket {
	return &Bucket{
		metric:          metric,
		dimensions:      sfxproto.Dimensions(dimensions).Clone(),
		min:             math.MaxInt64,
		max:             math.MinInt64,
		disabledMetrics: make(map[int]bool, 0),
	}
}

// Add an item to the Bucket, later reporting the result in the next report
// cycle.
func (b *Bucket) Add(val int64) {
	// wrap all of these changes in the mutex lock since they need to occur as a
	// transaction, not a set of atomic operations
	b.lock()
	defer b.unlock()

	// still use atomic though, so that the atomic "getters" will never read
	// from an inconsistent state
	_ = atomic.AddUint64(&b.count, 1)
	atomic.AddInt64(&b.sum, val)
	atomic.AddInt64(&b.sumOfSquares, val*val)

	if b.Min() > val {
		atomic.StoreInt64(&b.min, val)
	}

	if b.Max() < val {
		atomic.StoreInt64(&b.max, val)
	}
}

func (b *Bucket) dimFor(defaultDims map[string]string, rollup string) map[string]string {
	b.lock()
	defer b.unlock()

	dims := sfxproto.Dimensions(defaultDims).Append(b.dimensions)
	dims["rollup"] = rollup

	return dims
}

// DataPoints returns a DataPoints object with DataPoint values for
// Count, Sum, SumOfSquares, Min and Max (if set). Note that this
// resets all values.  If no values have been added to the bucket
// since the last report, it returns 0 for count, sum and
// sum-of-squares, omitting max and min.  If the count is higher than
// may be represented in an int64, then the count will be omitted.
func (b *Bucket) DataPoints() []DataPoint {
	dps := make([]DataPoint, 0, 5)
	cnt := atomic.SwapUint64(&b.count, 0)
	min := atomic.SwapInt64(&b.min, math.MaxInt64)
	max := atomic.SwapInt64(&b.max, math.MinInt64)
	sum := atomic.SwapInt64(&b.sum, 0)
	sos := atomic.SwapInt64(&b.sumOfSquares, 0)
	if cnt != 0 {
		if !b.disabledMetrics[BucketMetricMin] {
			dp := DataPoint{
				Metric:     b.metric,
				Dimensions: b.dimensions,
				Type:       GaugeType,
				Value:      min,
			}
			dps = append(dps, dp)
		}
		if !b.disabledMetrics[BucketMetricMax] {
			dp := DataPoint{
				Metric:     b.metric,
				Dimensions: b.dimensions,
				Type:       GaugeType,
				Value:      max,
			}
			dps = append(dps, dp)
		}
	}
	if !b.disabledMetrics[BucketMetricCount] && cnt <= math.MaxInt64 {
		dp := DataPoint{
			Metric:     b.metric,
			Dimensions: b.dimensions,
			Type:       CounterType,
			Value:      int64(cnt),
		}
		dps = append(dps, dp)
	}
	if !b.disabledMetrics[BucketMetricSum] {
		dp := DataPoint{
			Metric:     b.metric,
			Dimensions: b.dimensions,
			Type:       GaugeType,
			Value:      sum,
		}
		dps = append(dps, dp)
	}
	if !b.disabledMetrics[BucketMetricSumOfSquares] {
		dp := DataPoint{
			Metric:     b.metric,
			Dimensions: b.dimensions,
			Type:       GaugeType,
			Value:      sos,
		}
		dps = append(dps, dp)
	}

	return dps
}
<|endoftext|>


<commit_before>package riakpbc

// ListBuckets lists all buckets.
func (c *Client) ListBuckets() (*RpbListBucketsResp, error) {
	opts := []byte{}

	response, err := c.ReqResp(opts, "RpbListBucketsReq", true)
	if err != nil {
		return nil, err
	}

	return response.(*RpbListBucketsResp), nil
}

// ListKeysRequest prepares a ListKeys request.
func (c *Client) NewListKeysRequest(bucket string) *RpbListKeysReq {
	return &RpbListKeysReq{
		Bucket: []byte(bucket),
	}
}

func (c *Client) listKeys(opts *RpbListKeysReq, bucket string) ([][]byte, error) {
	if opts == nil {
		opts = c.NewListKeysRequest(bucket)
	}

	response, err := c.ReqMultiResp(opts, "RpbListKeysReq")
	if err != nil {
		return nil, err
	}

	keys := response.([][]byte)

	return keys, nil
}

// ListKeys lists all keys from bucket.
func (c *Client) ListKeys(bucket string) ([][]byte, error) {
	return c.listKeys(nil, bucket)
}

// NewGetBucketRequest prepares a GetBucket request.
func (c *Client) NewGetBucketRequest(bucket string) *RpbGetBucketReq {
	return &RpbGetBucketReq{
		Bucket: []byte(bucket),
	}
}

func (c *Client) getBucket(opts *RpbGetBucketReq, bucket string) (*RpbGetBucketResp, error) {
	if opts == nil {
		opts = c.NewGetBucketRequest(bucket)
	}

	response, err := c.ReqResp(opts, "RpbGetBucketReq", false)
	if err != nil {
		return nil, err
	}

	return response.(*RpbGetBucketResp), nil
}

// GetBucket gets the bucket info.
func (c *Client) GetBucket(bucket string) (*RpbGetBucketResp, error) {
	return c.getBucket(nil, bucket)
}

// NewSetBucketRequest prepares a SetBucket request.
func (c *Client) NewSetBucketRequest(bucket string, nval *uint32, allowmult *bool) *RpbSetBucketReq {
	return &RpbSetBucketReq{
		Bucket: []byte(bucket),
		Props: &RpbBucketProps{
			NVal:      nval,
			AllowMult: allowmult,
		},
	}
}

func (c *Client) setBucket(opts *RpbSetBucketReq, bucket string, nval *uint32, allowmult *bool) ([]byte, error) {
	if opts == nil {
		opts = c.NewSetBucketRequest(bucket, nval, allowmult)
	}

	response, err := c.ReqResp(opts, "RpbSetBucketReq", false)
	if err != nil {
		return nil, err
	}

	return response.([]byte), nil
}

// SetBucket sets the bucket info.
func (c *Client) SetBucket(bucket string, nval *uint32, allowmult *bool) ([]byte, error) {
	return c.setBucket(nil, bucket, nval, allowmult)
}
<commit_msg>Fix function name in comment<commit_after>package riakpbc

// ListBuckets lists all buckets.
func (c *Client) ListBuckets() (*RpbListBucketsResp, error) {
	opts := []byte{}

	response, err := c.ReqResp(opts, "RpbListBucketsReq", true)
	if err != nil {
		return nil, err
	}

	return response.(*RpbListBucketsResp), nil
}

// NewListKeysRequest prepares a ListKeys request.
func (c *Client) NewListKeysRequest(bucket string) *RpbListKeysReq {
	return &RpbListKeysReq{
		Bucket: []byte(bucket),
	}
}

func (c *Client) listKeys(opts *RpbListKeysReq, bucket string) ([][]byte, error) {
	if opts == nil {
		opts = c.NewListKeysRequest(bucket)
	}

	response, err := c.ReqMultiResp(opts, "RpbListKeysReq")
	if err != nil {
		return nil, err
	}

	keys := response.([][]byte)

	return keys, nil
}

// ListKeys lists all keys from bucket.
func (c *Client) ListKeys(bucket string) ([][]byte, error) {
	return c.listKeys(nil, bucket)
}

// NewGetBucketRequest prepares a GetBucket request.
func (c *Client) NewGetBucketRequest(bucket string) *RpbGetBucketReq {
	return &RpbGetBucketReq{
		Bucket: []byte(bucket),
	}
}

func (c *Client) getBucket(opts *RpbGetBucketReq, bucket string) (*RpbGetBucketResp, error) {
	if opts == nil {
		opts = c.NewGetBucketRequest(bucket)
	}

	response, err := c.ReqResp(opts, "RpbGetBucketReq", false)
	if err != nil {
		return nil, err
	}

	return response.(*RpbGetBucketResp), nil
}

// GetBucket gets the bucket info.
func (c *Client) GetBucket(bucket string) (*RpbGetBucketResp, error) {
	return c.getBucket(nil, bucket)
}

// NewSetBucketRequest prepares a SetBucket request.
func (c *Client) NewSetBucketRequest(bucket string, nval *uint32, allowmult *bool) *RpbSetBucketReq {
	return &RpbSetBucketReq{
		Bucket: []byte(bucket),
		Props: &RpbBucketProps{
			NVal:      nval,
			AllowMult: allowmult,
		},
	}
}

func (c *Client) setBucket(opts *RpbSetBucketReq, bucket string, nval *uint32, allowmult *bool) ([]byte, error) {
	if opts == nil {
		opts = c.NewSetBucketRequest(bucket, nval, allowmult)
	}

	response, err := c.ReqResp(opts, "RpbSetBucketReq", false)
	if err != nil {
		return nil, err
	}

	return response.([]byte), nil
}

// SetBucket sets the bucket info.
func (c *Client) SetBucket(bucket string, nval *uint32, allowmult *bool) ([]byte, error) {
	return c.setBucket(nil, bucket, nval, allowmult)
}
<|endoftext|>


<commit_before>/*
* @Author: souravray
* @Date:   2014-10-15 02:23:23
* @Last Modified by:   souravray
* @Last Modified time: 2015-02-14 23:17:35
 */

package polybolos

import (
	"sync/atomic"
	"time"
)

type bucket struct {
	tokens     int32
	usedTokens int32
	capacity   int32
	period     time.Duration
	stop       chan bool
}

func newBucket(capacity int32, rate int32) (b *bucket, err error) {

	b = &bucket{capacity: capacity}
	period := time.Duration(1e9 / int64(rate))
	if period > 0 {
		b.period = period
	} else {
		b.period = 1
	}
	return
}

func (b *bucket) setupUsedTokens(delta int32) {
	for {
		usedTokens := atomic.LoadInt32(&b.usedTokens)
		if !atomic.CompareAndSwapInt32(&b.usedTokens, usedTokens, usedTokens+delta) {
			continue
		} else {
			break
		}
	}
	return
}

func (b *bucket) setdownUsedTokens(delta int32) {
	for {
		if usedTokens := atomic.LoadInt32(&b.usedTokens); usedTokens < delta {
			if !atomic.CompareAndSwapInt32(&b.usedTokens, usedTokens, 0) {
				continue
			} else {
				break
			}
		} else {
			if !atomic.CompareAndSwapInt32(&b.usedTokens, usedTokens, usedTokens-delta) {
				continue
			} else {
				break
			}
		}
	}
	return
}

func (b *bucket) Put() (success bool) {
	for {
		tokens := atomic.LoadInt32(&b.tokens)
		usedTokens := atomic.LoadInt32(&b.usedTokens)
		if tokens+usedTokens < b.capacity {
			if !atomic.CompareAndSwapInt32(&b.tokens, tokens, tokens+1) {
				continue
			} else {
				break
			}
		} else {
			break
		}
	}
	return
}

func (b *bucket) Take(n int32) <-chan int32 {
	c := make(chan int32)
	go func(c chan int32, b *bucket, n int32) {
		var tokens int32
		defer close(c)
	TryToTake:
		for {
			if tokens = atomic.LoadInt32(&b.tokens); tokens == 0 {
				break
			} else if n <= tokens {
				if !atomic.CompareAndSwapInt32(&b.tokens, tokens, tokens-n) {
					continue
				} else {
					break
				}
			} else {
				if !atomic.CompareAndSwapInt32(&b.tokens, tokens, 0) {
					continue
				} else {
					break
				}
			}
		}

		if tokens > 0 {
			b.setupUsedTokens(tokens)
			c <- tokens
		} else {
			time.Sleep(time.Duration(n * int32(b.period.Nanoseconds())))
			goto TryToTake
		}
	}(c, b, n)
	return c
}

func (b *bucket) Spend() {
	b.setdownUsedTokens(1)
	return
}

func (b *bucket) Fill() {
	b.stop = make(chan bool, 0)
	go func(b *bucket) {
		defer close(b.stop)
		ticker := time.NewTicker(b.period)
		for _ = range ticker.C {
			select {
			case <-b.stop:
				ticker.Stop()
				return
			default:
				go b.Put()
			}
		}
	}(b)
}

func (b *bucket) Close() {
	b.stop <- true
	return
}
<commit_msg>workaround for minimize skipping ticks when fill rate is more than 100 per sec<commit_after>/*
* @Author: souravray
* @Date:   2014-10-15 02:23:23
* @Last Modified by:   souravray
* @Last Modified time: 2015-02-14 23:40:12
 */

package polybolos

import (
	"math"
	"sync/atomic"
	"time"
)

type bucket struct {
	tokens     int32
	usedTokens int32
	capacity   int32
	period     time.Duration
	fillSize   int32
	stop       chan bool
}

func newBucket(capacity int32, rate int32) (b *bucket, err error) {

	b = &bucket{capacity: capacity}

	if rate > 100 {
		b.fillSize = int32(math.Ceil(float64(rate / 10)))
		rate = 10
	} else {
		b.fillSize = 1
	}

	period := time.Duration(1e9 / int64(rate))
	if period > 0 {
		b.period = period
	} else {
		b.period = 1
	}
	return
}

func (b *bucket) setupUsedTokens(delta int32) {
	for {
		usedTokens := atomic.LoadInt32(&b.usedTokens)
		if !atomic.CompareAndSwapInt32(&b.usedTokens, usedTokens, usedTokens+delta) {
			continue
		} else {
			break
		}
	}
	return
}

func (b *bucket) setdownUsedTokens(delta int32) {
	for {
		if usedTokens := atomic.LoadInt32(&b.usedTokens); usedTokens < delta {
			if !atomic.CompareAndSwapInt32(&b.usedTokens, usedTokens, 0) {
				continue
			} else {
				break
			}
		} else {
			if !atomic.CompareAndSwapInt32(&b.usedTokens, usedTokens, usedTokens-delta) {
				continue
			} else {
				break
			}
		}
	}
	return
}

func (b *bucket) Put(n int32) (success bool) {
	for {
		tokens := atomic.LoadInt32(&b.tokens)
		usedTokens := atomic.LoadInt32(&b.usedTokens)
		if tokens+usedTokens < b.capacity {
			if !atomic.CompareAndSwapInt32(&b.tokens, tokens, tokens+n) {
				continue
			} else {
				break
			}
		} else {
			break
		}
	}
	return
}

func (b *bucket) Take(n int32) <-chan int32 {
	c := make(chan int32)
	go func(c chan int32, b *bucket, n int32) {
		var tokens int32
		defer close(c)
	TryToTake:
		for {
			if tokens = atomic.LoadInt32(&b.tokens); tokens == 0 {
				break
			} else if n <= tokens {
				if !atomic.CompareAndSwapInt32(&b.tokens, tokens, tokens-n) {
					continue
				} else {
					break
				}
			} else {
				if !atomic.CompareAndSwapInt32(&b.tokens, tokens, 0) {
					continue
				} else {
					break
				}
			}
		}

		if tokens > 0 {
			b.setupUsedTokens(tokens)
			c <- tokens
		} else {
			time.Sleep(time.Duration(n * int32(b.period.Nanoseconds())))
			goto TryToTake
		}
	}(c, b, n)
	return c
}

func (b *bucket) Spend() {
	b.setdownUsedTokens(1)
	return
}

func (b *bucket) Fill() {
	b.stop = make(chan bool, 0)
	go func(b *bucket) {
		defer close(b.stop)
		ticker := time.NewTicker(b.period)
		for _ = range ticker.C {
			select {
			case <-b.stop:
				ticker.Stop()
				return
			default:
				go b.Put(b.fillSize)
			}
		}
	}(b)
}

func (b *bucket) Close() {
	b.stop <- true
	return
}
<|endoftext|>