<commit_before>// Package gs provides a restic backend for Google Cloud Storage.
package gs

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/pkg/errors"
	"github.com/restic/restic/internal/backend"
	"github.com/restic/restic/internal/debug"
	"github.com/restic/restic/internal/restic"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/googleapi"
	storage "google.golang.org/api/storage/v1"
)

// Backend stores data in a GCS bucket.
//
// The service account used to access the bucket must have these permissions:
//  * storage.objects.create
//  * storage.objects.delete
//  * storage.objects.get
//  * storage.objects.list
type Backend struct {
	service      *storage.Service
	projectID    string
	sem          *backend.Semaphore
	bucketName   string
	prefix       string
	listMaxItems int
	backend.Layout
}

// Ensure that *Backend implements restic.Backend.
var _ restic.Backend = &Backend{}

func getStorageService(rt http.RoundTripper) (*storage.Service, error) {
	// create a new HTTP client
	httpClient := &http.Client{
		Transport: rt,
	}

	// create a new context with the HTTP client stored at the oauth2.HTTPClient key
	ctx := context.WithValue(context.Background(), oauth2.HTTPClient, httpClient)

	var ts oauth2.TokenSource
	if token := os.Getenv("GOOGLE_ACCESS_TOKEN"); token != "" {
		ts = oauth2.StaticTokenSource(&oauth2.Token{
			AccessToken: token,
			TokenType:   "Bearer",
		})
	} else {
		var err error
		ts, err = google.DefaultTokenSource(ctx, storage.DevstorageReadWriteScope)
		if err != nil {
			return nil, err
		}
	}

	client := oauth2.NewClient(ctx, ts)

	service, err := storage.New(client)
	if err != nil {
		return nil, err
	}

	return service, nil
}

const defaultListMaxItems = 1000

func open(cfg Config, rt http.RoundTripper) (*Backend, error) {
	debug.Log("open, config %#v", cfg)

	service, err := getStorageService(rt)
	if err != nil {
		return nil, errors.Wrap(err, "getStorageService")
	}

	sem, err := backend.NewSemaphore(cfg.Connections)
	if err != nil {
		return nil, err
	}

	be := &Backend{
		service:    service,
		projectID:  cfg.ProjectID,
		sem:        sem,
		bucketName: cfg.Bucket,
		prefix:     cfg.Prefix,
		Layout: &backend.DefaultLayout{
			Path: cfg.Prefix,
			Join: path.Join,
		},
		listMaxItems: defaultListMaxItems,
	}

	return be, nil
}

// Open opens the gs backend at the specified bucket.
func Open(cfg Config, rt http.RoundTripper) (restic.Backend, error) {
	return open(cfg, rt)
}

// Create opens the gs backend at the specified bucket and attempts to creates
// the bucket if it does not exist yet.
//
// The service account must have the "storage.buckets.create" permission to
// create a bucket the does not yet exist.
func Create(cfg Config, rt http.RoundTripper) (restic.Backend, error) {
	be, err := open(cfg, rt)
	if err != nil {
		return nil, errors.Wrap(err, "open")
	}

	// Try to determine if the bucket exists. If it does not, try to create it.
	//
	// A Get call has three typical error cases:
	//
	// * nil: Bucket exists and we have access to the metadata (returned).
	//
	// * 403: Bucket exists and we do not have access to the metadata. We
	// don't have storage.buckets.get permission to the bucket, but we may
	// still be able to access objects in the bucket.
	//
	// * 404: Bucket doesn't exist.
	//
	// Determining if the bucket is accessible is best-effort because the
	// 403 case is ambiguous.
	if _, err := be.service.Buckets.Get(be.bucketName).Do(); err != nil {
		gerr, ok := err.(*googleapi.Error)
		if !ok {
			// Don't know what to do with this error.
			return nil, errors.Wrap(err, "service.Buckets.Get")
		}

		switch gerr.Code {
		case 403:
			// Bucket exists, but we don't know if it is
			// accessible. Optimistically assume it is; if not,
			// future Backend calls will fail.
			debug.Log("Unable to determine if bucket %s is accessible (err %v). Continuing as if it is.", be.bucketName, err)
		case 404:
			// Bucket doesn't exist, try to create it.
			bucket := &storage.Bucket{
				Name: be.bucketName,
			}

			if _, err := be.service.Buckets.Insert(be.projectID, bucket).Do(); err != nil {
				// Always an error, as the bucket definitely
				// doesn't exist.
				return nil, errors.Wrap(err, "service.Buckets.Insert")
			}
		default:
			// Don't know what to do with this error.
			return nil, errors.Wrap(err, "service.Buckets.Get")
		}
	}

	return be, nil
}

// SetListMaxItems sets the number of list items to load per request.
func (be *Backend) SetListMaxItems(i int) {
	be.listMaxItems = i
}

// IsNotExist returns true if the error is caused by a not existing file.
func (be *Backend) IsNotExist(err error) bool {
	debug.Log("IsNotExist(%T, %#v)", err, err)

	if os.IsNotExist(err) {
		return true
	}

	if er, ok := err.(*googleapi.Error); ok {
		if er.Code == 404 {
			return true
		}
	}

	return false
}

// Join combines path components with slashes.
func (be *Backend) Join(p ...string) string {
	return path.Join(p...)
}

// Location returns this backend's location (the bucket name).
func (be *Backend) Location() string {
	return be.Join(be.bucketName, be.prefix)
}

// Path returns the path in the bucket that is used for this backend.
func (be *Backend) Path() string {
	return be.prefix
}

// Save stores data in the backend at the handle.
func (be *Backend) Save(ctx context.Context, h restic.Handle, rd restic.RewindReader) error {
	if err := h.Valid(); err != nil {
		return err
	}

	objName := be.Filename(h)

	debug.Log("Save %v at %v", h, objName)

	be.sem.GetToken()

	debug.Log("InsertObject(%v, %v)", be.bucketName, objName)

	// Set chunk size to zero to disable resumable uploads.
	//
	// With a non-zero chunk size (the default is
	// googleapi.DefaultUploadChunkSize, 8MB), Insert will buffer data from
	// rd in chunks of this size so it can upload these chunks in
	// individual requests.
	//
	// This chunking allows the library to automatically handle network
	// interruptions and re-upload only the last chunk rather than the full
	// file.
	//
	// Unfortunately, this buffering doesn't play nicely with
	// --limit-upload, which applies a rate limit to rd. This rate limit
	// ends up only limiting the read from rd into the buffer rather than
	// the network traffic itself. This results in poor network rate limit
	// behavior, where individual chunks are written to the network at full
	// bandwidth for several seconds, followed by several seconds of no
	// network traffic as the next chunk is read through the rate limiter.
	//
	// By disabling chunking, rd is passed further down the request stack,
	// where there is less (but some) buffering, which ultimately results
	// in better rate limiting behavior.
	//
	// restic typically writes small blobs (4MB-30MB), so the resumable
	// uploads are not providing significant benefit anyways.
	cs := googleapi.ChunkSize(0)

	info, err := be.service.Objects.Insert(be.bucketName,
		&storage.Object{
			Name: objName,
			Size: uint64(rd.Length()),
		}).Media(rd, cs).Do()

	be.sem.ReleaseToken()

	if err != nil {
		debug.Log("%v: err %#v: %v", objName, err, err)
		return errors.Wrap(err, "service.Objects.Insert")
	}

	debug.Log("%v -> %v bytes", objName, info.Size)
	return nil
}

// wrapReader wraps an io.ReadCloser to run an additional function on Close.
type wrapReader struct {
	io.ReadCloser
	f func()
}

func (wr wrapReader) Close() error {
	err := wr.ReadCloser.Close()
	wr.f()
	return err
}

// Load runs fn with a reader that yields the contents of the file at h at the
// given offset.
func (be *Backend) Load(ctx context.Context, h restic.Handle, length int, offset int64, fn func(rd io.Reader) error) error {
	return backend.DefaultLoad(ctx, h, length, offset, be.openReader, fn)
}

func (be *Backend) openReader(ctx context.Context, h restic.Handle, length int, offset int64) (io.ReadCloser, error) {
	debug.Log("Load %v, length %v, offset %v from %v", h, length, offset, be.Filename(h))
	if err := h.Valid(); err != nil {
		return nil, err
	}

	if offset < 0 {
		return nil, errors.New("offset is negative")
	}

	if length < 0 {
		return nil, errors.Errorf("invalid length %d", length)
	}

	objName := be.Filename(h)

	be.sem.GetToken()

	var byteRange string
	if length > 0 {
		byteRange = fmt.Sprintf("bytes=%d-%d", offset, offset+int64(length-1))
	} else {
		byteRange = fmt.Sprintf("bytes=%d-", offset)
	}

	req := be.service.Objects.Get(be.bucketName, objName)
	// https://cloud.google.com/storage/docs/json_api/v1/parameters#range
	req.Header().Set("Range", byteRange)
	res, err := req.Download()
	if err != nil {
		be.sem.ReleaseToken()
		return nil, err
	}

	closeRd := wrapReader{
		ReadCloser: res.Body,
		f: func() {
			debug.Log("Close()")
			be.sem.ReleaseToken()
		},
	}

	return closeRd, err
}

// Stat returns information about a blob.
func (be *Backend) Stat(ctx context.Context, h restic.Handle) (bi restic.FileInfo, err error) {
	debug.Log("%v", h)

	objName := be.Filename(h)

	be.sem.GetToken()
	obj, err := be.service.Objects.Get(be.bucketName, objName).Do()
	be.sem.ReleaseToken()

	if err != nil {
		debug.Log("GetObject() err %v", err)
		return restic.FileInfo{}, errors.Wrap(err, "service.Objects.Get")
	}

	return restic.FileInfo{Size: int64(obj.Size), Name: h.Name}, nil
}

// Test returns true if a blob of the given type and name exists in the backend.
func (be *Backend) Test(ctx context.Context, h restic.Handle) (bool, error) {
	found := false
	objName := be.Filename(h)

	be.sem.GetToken()
	_, err := be.service.Objects.Get(be.bucketName, objName).Do()
	be.sem.ReleaseToken()

	if err == nil {
		found = true
	}
	// If error, then not found
	return found, nil
}

// Remove removes the blob with the given name and type.
func (be *Backend) Remove(ctx context.Context, h restic.Handle) error {
	objName := be.Filename(h)

	be.sem.GetToken()
	err := be.service.Objects.Delete(be.bucketName, objName).Do()
	be.sem.ReleaseToken()

	if er, ok := err.(*googleapi.Error); ok {
		if er.Code == 404 {
			err = nil
		}
	}

	debug.Log("Remove(%v) at %v -> err %v", h, objName, err)
	return errors.Wrap(err, "client.RemoveObject")
}

// List runs fn for each file in the backend which has the type t. When an
// error occurs (or fn returns an error), List stops and returns it.
func (be *Backend) List(ctx context.Context, t restic.FileType, fn func(restic.FileInfo) error) error {
	debug.Log("listing %v", t)

	prefix, _ := be.Basedir(t)

	// make sure prefix ends with a slash
	if !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	listReq := be.service.Objects.List(be.bucketName).Context(ctx).Prefix(prefix).MaxResults(int64(be.listMaxItems))
	for {
		be.sem.GetToken()
		obj, err := listReq.Do()
		be.sem.ReleaseToken()

		if err != nil {
			return err
		}

		debug.Log("returned %v items", len(obj.Items))

		for _, item := range obj.Items {
			m := strings.TrimPrefix(item.Name, prefix)
			if m == "" {
				continue
			}

			if ctx.Err() != nil {
				return ctx.Err()
			}

			fi := restic.FileInfo{
				Name: path.Base(m),
				Size: int64(item.Size),
			}

			err := fn(fi)
			if err != nil {
				return err
			}

			if ctx.Err() != nil {
				return ctx.Err()
			}
		}

		if obj.NextPageToken == "" {
			break
		}
		listReq.PageToken(obj.NextPageToken)
	}

	return ctx.Err()
}

// Remove keys for a specified backend type.
func (be *Backend) removeKeys(ctx context.Context, t restic.FileType) error {
	return be.List(ctx, t, func(fi restic.FileInfo) error {
		return be.Remove(ctx, restic.Handle{Type: t, Name: fi.Name})
	})
}

// Delete removes all restic keys in the bucket. It will not remove the bucket itself.
func (be *Backend) Delete(ctx context.Context) error {
	alltypes := []restic.FileType{
		restic.PackFile,
		restic.KeyFile,
		restic.LockFile,
		restic.SnapshotFile,
		restic.IndexFile}

	for _, t := range alltypes {
		err := be.removeKeys(ctx, t)
		if err != nil {
			return nil
		}
	}

	return be.Remove(ctx, restic.Handle{Type: restic.ConfigFile})
}

// Close does nothing.
func (be *Backend) Close() error { return nil }
<commit_msg>Replace deprecated method in gs backend<commit_after>// Package gs provides a restic backend for Google Cloud Storage.
package gs

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/pkg/errors"
	"github.com/restic/restic/internal/backend"
	"github.com/restic/restic/internal/debug"
	"github.com/restic/restic/internal/restic"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/option"
	storage "google.golang.org/api/storage/v1"
)

// Backend stores data in a GCS bucket.
//
// The service account used to access the bucket must have these permissions:
//  * storage.objects.create
//  * storage.objects.delete
//  * storage.objects.get
//  * storage.objects.list
type Backend struct {
	service      *storage.Service
	projectID    string
	sem          *backend.Semaphore
	bucketName   string
	prefix       string
	listMaxItems int
	backend.Layout
}

// Ensure that *Backend implements restic.Backend.
var _ restic.Backend = &Backend{}

func getStorageService(rt http.RoundTripper) (*storage.Service, error) {
	// create a new HTTP client
	httpClient := &http.Client{
		Transport: rt,
	}

	// create a new context with the HTTP client stored at the oauth2.HTTPClient key
	ctx := context.WithValue(context.Background(), oauth2.HTTPClient, httpClient)

	var ts oauth2.TokenSource
	if token := os.Getenv("GOOGLE_ACCESS_TOKEN"); token != "" {
		ts = oauth2.StaticTokenSource(&oauth2.Token{
			AccessToken: token,
			TokenType:   "Bearer",
		})
	} else {
		var err error
		ts, err = google.DefaultTokenSource(ctx, storage.DevstorageReadWriteScope)
		if err != nil {
			return nil, err
		}
	}

	client := oauth2.NewClient(ctx, ts)

	service, err := storage.NewService(ctx, option.WithHTTPClient(client))
	if err != nil {
		return nil, err
	}

	return service, nil
}

const defaultListMaxItems = 1000

func open(cfg Config, rt http.RoundTripper) (*Backend, error) {
	debug.Log("open, config %#v", cfg)

	service, err := getStorageService(rt)
	if err != nil {
		return nil, errors.Wrap(err, "getStorageService")
	}

	sem, err := backend.NewSemaphore(cfg.Connections)
	if err != nil {
		return nil, err
	}

	be := &Backend{
		service:    service,
		projectID:  cfg.ProjectID,
		sem:        sem,
		bucketName: cfg.Bucket,
		prefix:     cfg.Prefix,
		Layout: &backend.DefaultLayout{
			Path: cfg.Prefix,
			Join: path.Join,
		},
		listMaxItems: defaultListMaxItems,
	}

	return be, nil
}

// Open opens the gs backend at the specified bucket.
func Open(cfg Config, rt http.RoundTripper) (restic.Backend, error) {
	return open(cfg, rt)
}

// Create opens the gs backend at the specified bucket and attempts to creates
// the bucket if it does not exist yet.
//
// The service account must have the "storage.buckets.create" permission to
// create a bucket the does not yet exist.
func Create(cfg Config, rt http.RoundTripper) (restic.Backend, error) {
	be, err := open(cfg, rt)
	if err != nil {
		return nil, errors.Wrap(err, "open")
	}

	// Try to determine if the bucket exists. If it does not, try to create it.
	//
	// A Get call has three typical error cases:
	//
	// * nil: Bucket exists and we have access to the metadata (returned).
	//
	// * 403: Bucket exists and we do not have access to the metadata. We
	// don't have storage.buckets.get permission to the bucket, but we may
	// still be able to access objects in the bucket.
	//
	// * 404: Bucket doesn't exist.
	//
	// Determining if the bucket is accessible is best-effort because the
	// 403 case is ambiguous.
	if _, err := be.service.Buckets.Get(be.bucketName).Do(); err != nil {
		gerr, ok := err.(*googleapi.Error)
		if !ok {
			// Don't know what to do with this error.
			return nil, errors.Wrap(err, "service.Buckets.Get")
		}

		switch gerr.Code {
		case 403:
			// Bucket exists, but we don't know if it is
			// accessible. Optimistically assume it is; if not,
			// future Backend calls will fail.
			debug.Log("Unable to determine if bucket %s is accessible (err %v). Continuing as if it is.", be.bucketName, err)
		case 404:
			// Bucket doesn't exist, try to create it.
			bucket := &storage.Bucket{
				Name: be.bucketName,
			}

			if _, err := be.service.Buckets.Insert(be.projectID, bucket).Do(); err != nil {
				// Always an error, as the bucket definitely
				// doesn't exist.
				return nil, errors.Wrap(err, "service.Buckets.Insert")
			}
		default:
			// Don't know what to do with this error.
			return nil, errors.Wrap(err, "service.Buckets.Get")
		}
	}

	return be, nil
}

// SetListMaxItems sets the number of list items to load per request.
func (be *Backend) SetListMaxItems(i int) {
	be.listMaxItems = i
}

// IsNotExist returns true if the error is caused by a not existing file.
func (be *Backend) IsNotExist(err error) bool {
	debug.Log("IsNotExist(%T, %#v)", err, err)

	if os.IsNotExist(err) {
		return true
	}

	if er, ok := err.(*googleapi.Error); ok {
		if er.Code == 404 {
			return true
		}
	}

	return false
}

// Join combines path components with slashes.
func (be *Backend) Join(p ...string) string {
	return path.Join(p...)
}

// Location returns this backend's location (the bucket name).
func (be *Backend) Location() string {
	return be.Join(be.bucketName, be.prefix)
}

// Path returns the path in the bucket that is used for this backend.
func (be *Backend) Path() string {
	return be.prefix
}

// Save stores data in the backend at the handle.
func (be *Backend) Save(ctx context.Context, h restic.Handle, rd restic.RewindReader) error {
	if err := h.Valid(); err != nil {
		return err
	}

	objName := be.Filename(h)

	debug.Log("Save %v at %v", h, objName)

	be.sem.GetToken()

	debug.Log("InsertObject(%v, %v)", be.bucketName, objName)

	// Set chunk size to zero to disable resumable uploads.
	//
	// With a non-zero chunk size (the default is
	// googleapi.DefaultUploadChunkSize, 8MB), Insert will buffer data from
	// rd in chunks of this size so it can upload these chunks in
	// individual requests.
	//
	// This chunking allows the library to automatically handle network
	// interruptions and re-upload only the last chunk rather than the full
	// file.
	//
	// Unfortunately, this buffering doesn't play nicely with
	// --limit-upload, which applies a rate limit to rd. This rate limit
	// ends up only limiting the read from rd into the buffer rather than
	// the network traffic itself. This results in poor network rate limit
	// behavior, where individual chunks are written to the network at full
	// bandwidth for several seconds, followed by several seconds of no
	// network traffic as the next chunk is read through the rate limiter.
	//
	// By disabling chunking, rd is passed further down the request stack,
	// where there is less (but some) buffering, which ultimately results
	// in better rate limiting behavior.
	//
	// restic typically writes small blobs (4MB-30MB), so the resumable
	// uploads are not providing significant benefit anyways.
	cs := googleapi.ChunkSize(0)

	info, err := be.service.Objects.Insert(be.bucketName,
		&storage.Object{
			Name: objName,
			Size: uint64(rd.Length()),
		}).Media(rd, cs).Do()

	be.sem.ReleaseToken()

	if err != nil {
		debug.Log("%v: err %#v: %v", objName, err, err)
		return errors.Wrap(err, "service.Objects.Insert")
	}

	debug.Log("%v -> %v bytes", objName, info.Size)
	return nil
}

// wrapReader wraps an io.ReadCloser to run an additional function on Close.
type wrapReader struct {
	io.ReadCloser
	f func()
}

func (wr wrapReader) Close() error {
	err := wr.ReadCloser.Close()
	wr.f()
	return err
}

// Load runs fn with a reader that yields the contents of the file at h at the
// given offset.
func (be *Backend) Load(ctx context.Context, h restic.Handle, length int, offset int64, fn func(rd io.Reader) error) error {
	return backend.DefaultLoad(ctx, h, length, offset, be.openReader, fn)
}

func (be *Backend) openReader(ctx context.Context, h restic.Handle, length int, offset int64) (io.ReadCloser, error) {
	debug.Log("Load %v, length %v, offset %v from %v", h, length, offset, be.Filename(h))
	if err := h.Valid(); err != nil {
		return nil, err
	}

	if offset < 0 {
		return nil, errors.New("offset is negative")
	}

	if length < 0 {
		return nil, errors.Errorf("invalid length %d", length)
	}

	objName := be.Filename(h)

	be.sem.GetToken()

	var byteRange string
	if length > 0 {
		byteRange = fmt.Sprintf("bytes=%d-%d", offset, offset+int64(length-1))
	} else {
		byteRange = fmt.Sprintf("bytes=%d-", offset)
	}

	req := be.service.Objects.Get(be.bucketName, objName)
	// https://cloud.google.com/storage/docs/json_api/v1/parameters#range
	req.Header().Set("Range", byteRange)
	res, err := req.Download()
	if err != nil {
		be.sem.ReleaseToken()
		return nil, err
	}

	closeRd := wrapReader{
		ReadCloser: res.Body,
		f: func() {
			debug.Log("Close()")
			be.sem.ReleaseToken()
		},
	}

	return closeRd, err
}

// Stat returns information about a blob.
func (be *Backend) Stat(ctx context.Context, h restic.Handle) (bi restic.FileInfo, err error) {
	debug.Log("%v", h)

	objName := be.Filename(h)

	be.sem.GetToken()
	obj, err := be.service.Objects.Get(be.bucketName, objName).Do()
	be.sem.ReleaseToken()

	if err != nil {
		debug.Log("GetObject() err %v", err)
		return restic.FileInfo{}, errors.Wrap(err, "service.Objects.Get")
	}

	return restic.FileInfo{Size: int64(obj.Size), Name: h.Name}, nil
}

// Test returns true if a blob of the given type and name exists in the backend.
func (be *Backend) Test(ctx context.Context, h restic.Handle) (bool, error) {
	found := false
	objName := be.Filename(h)

	be.sem.GetToken()
	_, err := be.service.Objects.Get(be.bucketName, objName).Do()
	be.sem.ReleaseToken()

	if err == nil {
		found = true
	}
	// If error, then not found
	return found, nil
}

// Remove removes the blob with the given name and type.
func (be *Backend) Remove(ctx context.Context, h restic.Handle) error {
	objName := be.Filename(h)

	be.sem.GetToken()
	err := be.service.Objects.Delete(be.bucketName, objName).Do()
	be.sem.ReleaseToken()

	if er, ok := err.(*googleapi.Error); ok {
		if er.Code == 404 {
			err = nil
		}
	}

	debug.Log("Remove(%v) at %v -> err %v", h, objName, err)
	return errors.Wrap(err, "client.RemoveObject")
}

// List runs fn for each file in the backend which has the type t. When an
// error occurs (or fn returns an error), List stops and returns it.
func (be *Backend) List(ctx context.Context, t restic.FileType, fn func(restic.FileInfo) error) error {
	debug.Log("listing %v", t)

	prefix, _ := be.Basedir(t)

	// make sure prefix ends with a slash
	if !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	listReq := be.service.Objects.List(be.bucketName).Context(ctx).Prefix(prefix).MaxResults(int64(be.listMaxItems))
	for {
		be.sem.GetToken()
		obj, err := listReq.Do()
		be.sem.ReleaseToken()

		if err != nil {
			return err
		}

		debug.Log("returned %v items", len(obj.Items))

		for _, item := range obj.Items {
			m := strings.TrimPrefix(item.Name, prefix)
			if m == "" {
				continue
			}

			if ctx.Err() != nil {
				return ctx.Err()
			}

			fi := restic.FileInfo{
				Name: path.Base(m),
				Size: int64(item.Size),
			}

			err := fn(fi)
			if err != nil {
				return err
			}

			if ctx.Err() != nil {
				return ctx.Err()
			}
		}

		if obj.NextPageToken == "" {
			break
		}
		listReq.PageToken(obj.NextPageToken)
	}

	return ctx.Err()
}

// Remove keys for a specified backend type.
func (be *Backend) removeKeys(ctx context.Context, t restic.FileType) error {
	return be.List(ctx, t, func(fi restic.FileInfo) error {
		return be.Remove(ctx, restic.Handle{Type: t, Name: fi.Name})
	})
}

// Delete removes all restic keys in the bucket. It will not remove the bucket itself.
func (be *Backend) Delete(ctx context.Context) error {
	alltypes := []restic.FileType{
		restic.PackFile,
		restic.KeyFile,
		restic.LockFile,
		restic.SnapshotFile,
		restic.IndexFile}

	for _, t := range alltypes {
		err := be.removeKeys(ctx, t)
		if err != nil {
			return nil
		}
	}

	return be.Remove(ctx, restic.Handle{Type: restic.ConfigFile})
}

// Close does nothing.
func (be *Backend) Close() error { return nil }
<|endoftext|>


<commit_before>/*
Copyright 2017 Mirantis

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 criproxy

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	"github.com/golang/glog"

	// TODO: switch to https://github.com/docker/docker/tree/master/client
	// Docker version used in k8s is too old for it
	dockerclient "github.com/docker/engine-api/client"
	dockertypes "github.com/docker/engine-api/types"
	dockercontainer "github.com/docker/engine-api/types/container"
	dockerfilters "github.com/docker/engine-api/types/filters"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/rest"
)

const (
	BusyboxImageName = "busybox:1.26.2"
	// TODO: use the same constant/setting in different parts of code
	proxyRuntimeEndpoint             = "/run/criproxy.sock"
	proxyStopTimeoutSeconds          = 5
	confFileMode                     = 0600
	confDirMode                      = 0700
	kubeletConfigPollInterval        = 1 * time.Second
	waitForCriProxySocketNumAttempts = 600
)

var kubeletSettingsForCriProxy map[string]interface{} = map[string]interface{}{
	"containerRuntime":      "remote",
	"remoteRuntimeEndpoint": proxyRuntimeEndpoint,
	"remoteImageEndpoint":   proxyRuntimeEndpoint,
	// NOTE: this setting is only needed by Virtlet to make flexvolume
	// handling simpler
	"enableControllerAttachDetach": false,
}

func loadJson(baseUrl, suffix string) (map[string]interface{}, error) {
	url := strings.TrimSuffix(baseUrl, "/") + suffix
	glog.V(1).Infof("Loading url %q", url)
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client := &http.Client{Transport: tr}
	res, err := client.Get(url)
	if err != nil {
		return nil, fmt.Errorf("trying to get %q: %v", url, err)
	}

	defer res.Body.Close()
	var r map[string]interface{}
	d := json.NewDecoder(res.Body)
	d.UseNumber() // avoid getting floats
	if err := d.Decode(&r); err != nil {
		return nil, fmt.Errorf("failed to unmarshal json from %q: %v", url, err)
	}
	return r, nil
}

func writeJson(data interface{}, file string) error {
	glog.V(1).Infof("Writing config file: %q", file)
	dir := path.Dir(file)
	if err := os.MkdirAll(dir, confDirMode); err != nil {
		return fmt.Errorf("failed to make conf dir: %v", err)
	}
	bs, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to marshal json: %v", err)
	}
	if err := ioutil.WriteFile(file, bs, confFileMode); err != nil {
		return fmt.Errorf("error writing %q: %v", file, err)
	}
	return nil
}

func kubeletConfigUpdated(kubeletCfg map[string]interface{}) bool {
	for k, v := range kubeletSettingsForCriProxy {
		if kubeletCfg[k] != v {
			return false
		}
	}
	return true
}

type BootstrapConfig struct {
	ApiServerHost string
	// TODO: remove this
	ConfigzBaseUrl  string
	ProxyPath       string
	ProxyArgs       []string
	ProxySocketPath string
	NodeInfo        *NodeInfo
}

type Bootstrap struct {
	config     *BootstrapConfig
	clientset  kubernetes.Interface
	kubeletCfg map[string]interface{}
}

// NewBootstrap creates a new Bootstrap object used for CRI proxy
// bootstrap using the specified BootstrapConfig. cs argument
// is used to pass a fake Clientset during tests, it should
// be nil when performing real bootstrap.
func NewBootstrap(config *BootstrapConfig, cs kubernetes.Interface) *Bootstrap {
	return &Bootstrap{config: config, clientset: cs}
}

func (b *Bootstrap) retrieveKubeletConfig() (map[string]interface{}, error) {
	// TODO: as we know node name, use k8s API for this
	cfg, err := loadJson(b.config.ConfigzBaseUrl, "/configz")
	if err != nil {
		return nil, err
	}
	var ok bool
	var kubeletCfg map[string]interface{}
	kubeletCfg, ok = cfg["componentconfig"].(map[string]interface{})
	if !ok {
		return nil, errors.New("couldn't get componentconfig from /configz")
	}
	return kubeletCfg, nil
}

func (b *Bootstrap) obtainKubeletConfig() error {
	kubeletCfg, err := b.retrieveKubeletConfig()
	if err != nil {
		return err
	}
	b.kubeletCfg = kubeletCfg
	return nil
}

func (b *Bootstrap) kubeletReadyAfterPatch() bool {
	kubeletCfg, err := b.retrieveKubeletConfig()
	if err != nil {
		glog.V(2).Infof("Can't retrieve kubelet config yet: %v", err)
		return false
	}
	return kubeletConfigUpdated(kubeletCfg)
}

func (b *Bootstrap) updateKubeletConfig() {
	for k, v := range kubeletSettingsForCriProxy {
		b.kubeletCfg[k] = v
	}
}

func (b *Bootstrap) buildConfigMap(nodeName string) *v1.ConfigMap {
	text, err := json.Marshal(b.kubeletCfg)
	if err != nil {
		log.Panicf("Couldn't marshal kubelet config: %v", err)
	}
	return &v1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "kubelet-" + nodeName,
			Namespace: "kube-system",
		},
		Data: map[string]string{
			"kubelet.config": string(text),
		},
	}
}

func (b *Bootstrap) putConfigMap(configMap *v1.ConfigMap) error {
	glog.V(1).Infof("Putting ConfigMap %q in namespace %q", configMap.Name, configMap.Namespace)
	_, err := b.clientset.Core().ConfigMaps("kube-system").Create(configMap)
	return err
}

func (b *Bootstrap) patchKubeletConfig() error {
	if kubeletConfigUpdated(b.kubeletCfg) {
		return fmt.Errorf("kubelet already configured for CRI, but no saved config")
	}
	b.updateKubeletConfig()

	glog.V(1).Infof("Node name: %q", b.config.NodeInfo.NodeName)
	if err := b.putConfigMap(b.buildConfigMap(b.config.NodeInfo.NodeName)); err != nil {
		return fmt.Errorf("failed to put ConfigMap: %v", err)
	}
	return nil
}

func (b *Bootstrap) installCriProxyContainer(dockerEndpoint, endpointToPass string) (string, error) {
	// CRI proxy container actually uses host namespaces to run.
	// It just uses docker's 'always' restart policy as a poor man's
	// substitute for a process manager.
	ctx := context.Background()

	client, err := dockerclient.NewClient(dockerEndpoint, "", nil, nil)
	if err != nil {
		return "", fmt.Errorf("failed to create Docker client: %v", err)
	}

	filterArgs := dockerfilters.NewArgs()
	filterArgs.Add("label", "criproxy")
	containers, err := client.ContainerList(ctx, dockertypes.ContainerListOptions{
		Filter: filterArgs,
	})
	if len(containers) > 0 {
		for _, container := range containers {
			glog.V(1).Infof("Removing old CRI proxy container %s", container.ID)
			if err := client.ContainerRemove(ctx, container.ID, dockertypes.ContainerRemoveOptions{
				Force: true,
			}); err != nil {
				return "", fmt.Errorf("failed to remove old container: %v", err)
			}
		}
	}

	if err := pullImage(ctx, client, BusyboxImageName, true); err != nil {
		return "", fmt.Errorf("failed to pull busybox image: %v", err)
	}

	containerName := fmt.Sprintf("criproxy-%d", time.Now().UnixNano())
	resp, err := client.ContainerCreate(ctx, &dockercontainer.Config{
		Image:  BusyboxImageName,
		Labels: map[string]string{"criproxy": "true"},
		Env:    []string{"DOCKER_HOST=" + endpointToPass},
		Cmd: append([]string{
			"nsenter",
			"--mount=/proc/1/ns/mnt",
			"--",
			b.config.ProxyPath,
		}, b.config.ProxyArgs...),
	}, &dockercontainer.HostConfig{
		Privileged:  true,
		NetworkMode: "host",
		UTSMode:     "host",
		PidMode:     "host",
		IpcMode:     "host",
		UsernsMode:  "host",
		RestartPolicy: dockercontainer.RestartPolicy{
			Name: "always",
		},
	}, nil, containerName)
	if err != nil {
		return "", fmt.Errorf("failed to create CRI proxy container: %v", err)
	}
	if err := client.ContainerStart(ctx, resp.ID); err != nil {
		client.ContainerRemove(ctx, resp.ID, dockertypes.ContainerRemoveOptions{
			Force: true,
		})
		return "", fmt.Errorf("failed to start CRI proxy container: %v", err)
	}
	glog.Infof("Started container %s", resp.ID)
	return resp.ID, nil
}

func (b *Bootstrap) initClientset() error {
	var err error
	var config *rest.Config
	if b.config.ApiServerHost == "" {
		config, err = rest.InClusterConfig()
		if err != nil {
			return fmt.Errorf("failed to get REST client config: %v", err)
		}
	} else {
		config = &rest.Config{Host: b.config.ApiServerHost}
	}
	b.clientset, err = kubernetes.NewForConfig(config)
	if err != nil {
		return fmt.Errorf("failed to create ClientSet: %v", err)
	}
	return nil
}

// EnsureCRIProxy checks whether kubelet configuration file exists
// and performs CRI proxy bootstrap procedure if it doesn't.
func (b *Bootstrap) EnsureCRIProxy() (bool, error) {
	if b.config.ConfigzBaseUrl == "" || b.config.ProxyPath == "" || b.config.ProxySocketPath == "" {
		return false, errors.New("invalid BootstrapConfig")
	}

	if !strings.HasPrefix(path.Clean(b.config.ProxySocketPath), "/run/") {
		return false, fmt.Errorf("proxy socket path %q not under /run", b.config.ProxySocketPath)
	}

	if err := b.obtainKubeletConfig(); err != nil {
		return false, err
	}

	dockerEndpoint := b.config.NodeInfo.DockerEndpoint
	glog.V(1).Infof("Using docker endpoint %q", dockerEndpoint)
	if _, err := b.installCriProxyContainer(dockerEndpoint, dockerEndpoint); err != nil {
		return false, err
	}

	if err := waitForSocket(b.config.ProxySocketPath, waitForCriProxySocketNumAttempts, nil); err != nil {
		return false, err
	}

	// We don't try to patch kubelet config before the container
	// is installed, because this process may die with kubelet
	// restart
	if err := b.initClientset(); err != nil {
		return false, err
	}
	if err := b.patchKubeletConfig(); err != nil {
		return false, err
	}

	for {
		time.Sleep(kubeletConfigPollInterval)
		if b.kubeletReadyAfterPatch() {
			break
		}
	}

	glog.V(1).Info("CRI proxy bootstrap complete")
	return true, nil
}
<commit_msg>Fix criproxy endpoint naming for bootstrap<commit_after>/*
Copyright 2017 Mirantis

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 criproxy

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	"github.com/golang/glog"

	// TODO: switch to https://github.com/docker/docker/tree/master/client
	// Docker version used in k8s is too old for it
	dockerclient "github.com/docker/engine-api/client"
	dockertypes "github.com/docker/engine-api/types"
	dockercontainer "github.com/docker/engine-api/types/container"
	dockerfilters "github.com/docker/engine-api/types/filters"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/rest"
)

const (
	BusyboxImageName = "busybox:1.26.2"
	// TODO: use the same constant/setting in different parts of code
	proxyRuntimeEndpoint             = "unix:///run/criproxy.sock"
	proxyStopTimeoutSeconds          = 5
	confFileMode                     = 0600
	confDirMode                      = 0700
	kubeletConfigPollInterval        = 1 * time.Second
	waitForCriProxySocketNumAttempts = 600
)

var kubeletSettingsForCriProxy map[string]interface{} = map[string]interface{}{
	"containerRuntime":      "remote",
	"remoteRuntimeEndpoint": proxyRuntimeEndpoint,
	"remoteImageEndpoint":   proxyRuntimeEndpoint,
	// NOTE: this setting is only needed by Virtlet to make flexvolume
	// handling simpler
	"enableControllerAttachDetach": false,
}

func loadJson(baseUrl, suffix string) (map[string]interface{}, error) {
	url := strings.TrimSuffix(baseUrl, "/") + suffix
	glog.V(1).Infof("Loading url %q", url)
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client := &http.Client{Transport: tr}
	res, err := client.Get(url)
	if err != nil {
		return nil, fmt.Errorf("trying to get %q: %v", url, err)
	}

	defer res.Body.Close()
	var r map[string]interface{}
	d := json.NewDecoder(res.Body)
	d.UseNumber() // avoid getting floats
	if err := d.Decode(&r); err != nil {
		return nil, fmt.Errorf("failed to unmarshal json from %q: %v", url, err)
	}
	return r, nil
}

func writeJson(data interface{}, file string) error {
	glog.V(1).Infof("Writing config file: %q", file)
	dir := path.Dir(file)
	if err := os.MkdirAll(dir, confDirMode); err != nil {
		return fmt.Errorf("failed to make conf dir: %v", err)
	}
	bs, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to marshal json: %v", err)
	}
	if err := ioutil.WriteFile(file, bs, confFileMode); err != nil {
		return fmt.Errorf("error writing %q: %v", file, err)
	}
	return nil
}

func kubeletConfigUpdated(kubeletCfg map[string]interface{}) bool {
	for k, v := range kubeletSettingsForCriProxy {
		if kubeletCfg[k] != v {
			return false
		}
	}
	return true
}

type BootstrapConfig struct {
	ApiServerHost string
	// TODO: remove this
	ConfigzBaseUrl  string
	ProxyPath       string
	ProxyArgs       []string
	ProxySocketPath string
	NodeInfo        *NodeInfo
}

type Bootstrap struct {
	config     *BootstrapConfig
	clientset  kubernetes.Interface
	kubeletCfg map[string]interface{}
}

// NewBootstrap creates a new Bootstrap object used for CRI proxy
// bootstrap using the specified BootstrapConfig. cs argument
// is used to pass a fake Clientset during tests, it should
// be nil when performing real bootstrap.
func NewBootstrap(config *BootstrapConfig, cs kubernetes.Interface) *Bootstrap {
	return &Bootstrap{config: config, clientset: cs}
}

func (b *Bootstrap) retrieveKubeletConfig() (map[string]interface{}, error) {
	// TODO: as we know node name, use k8s API for this
	cfg, err := loadJson(b.config.ConfigzBaseUrl, "/configz")
	if err != nil {
		return nil, err
	}
	var ok bool
	var kubeletCfg map[string]interface{}
	kubeletCfg, ok = cfg["componentconfig"].(map[string]interface{})
	if !ok {
		return nil, errors.New("couldn't get componentconfig from /configz")
	}
	return kubeletCfg, nil
}

func (b *Bootstrap) obtainKubeletConfig() error {
	kubeletCfg, err := b.retrieveKubeletConfig()
	if err != nil {
		return err
	}
	b.kubeletCfg = kubeletCfg
	return nil
}

func (b *Bootstrap) kubeletReadyAfterPatch() bool {
	kubeletCfg, err := b.retrieveKubeletConfig()
	if err != nil {
		glog.V(2).Infof("Can't retrieve kubelet config yet: %v", err)
		return false
	}
	return kubeletConfigUpdated(kubeletCfg)
}

func (b *Bootstrap) updateKubeletConfig() {
	for k, v := range kubeletSettingsForCriProxy {
		b.kubeletCfg[k] = v
	}
}

func (b *Bootstrap) buildConfigMap(nodeName string) *v1.ConfigMap {
	text, err := json.Marshal(b.kubeletCfg)
	if err != nil {
		log.Panicf("Couldn't marshal kubelet config: %v", err)
	}
	return &v1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "kubelet-" + nodeName,
			Namespace: "kube-system",
		},
		Data: map[string]string{
			"kubelet.config": string(text),
		},
	}
}

func (b *Bootstrap) putConfigMap(configMap *v1.ConfigMap) error {
	glog.V(1).Infof("Putting ConfigMap %q in namespace %q", configMap.Name, configMap.Namespace)
	_, err := b.clientset.Core().ConfigMaps("kube-system").Create(configMap)
	return err
}

func (b *Bootstrap) patchKubeletConfig() error {
	if kubeletConfigUpdated(b.kubeletCfg) {
		return fmt.Errorf("kubelet already configured for CRI, but no saved config")
	}
	b.updateKubeletConfig()

	glog.V(1).Infof("Node name: %q", b.config.NodeInfo.NodeName)
	if err := b.putConfigMap(b.buildConfigMap(b.config.NodeInfo.NodeName)); err != nil {
		return fmt.Errorf("failed to put ConfigMap: %v", err)
	}
	return nil
}

func (b *Bootstrap) installCriProxyContainer(dockerEndpoint, endpointToPass string) (string, error) {
	// CRI proxy container actually uses host namespaces to run.
	// It just uses docker's 'always' restart policy as a poor man's
	// substitute for a process manager.
	ctx := context.Background()

	client, err := dockerclient.NewClient(dockerEndpoint, "", nil, nil)
	if err != nil {
		return "", fmt.Errorf("failed to create Docker client: %v", err)
	}

	filterArgs := dockerfilters.NewArgs()
	filterArgs.Add("label", "criproxy")
	containers, err := client.ContainerList(ctx, dockertypes.ContainerListOptions{
		Filter: filterArgs,
	})
	if len(containers) > 0 {
		for _, container := range containers {
			glog.V(1).Infof("Removing old CRI proxy container %s", container.ID)
			if err := client.ContainerRemove(ctx, container.ID, dockertypes.ContainerRemoveOptions{
				Force: true,
			}); err != nil {
				return "", fmt.Errorf("failed to remove old container: %v", err)
			}
		}
	}

	if err := pullImage(ctx, client, BusyboxImageName, true); err != nil {
		return "", fmt.Errorf("failed to pull busybox image: %v", err)
	}

	containerName := fmt.Sprintf("criproxy-%d", time.Now().UnixNano())
	resp, err := client.ContainerCreate(ctx, &dockercontainer.Config{
		Image:  BusyboxImageName,
		Labels: map[string]string{"criproxy": "true"},
		Env:    []string{"DOCKER_HOST=" + endpointToPass},
		Cmd: append([]string{
			"nsenter",
			"--mount=/proc/1/ns/mnt",
			"--",
			b.config.ProxyPath,
		}, b.config.ProxyArgs...),
	}, &dockercontainer.HostConfig{
		Privileged:  true,
		NetworkMode: "host",
		UTSMode:     "host",
		PidMode:     "host",
		IpcMode:     "host",
		UsernsMode:  "host",
		RestartPolicy: dockercontainer.RestartPolicy{
			Name: "always",
		},
	}, nil, containerName)
	if err != nil {
		return "", fmt.Errorf("failed to create CRI proxy container: %v", err)
	}
	if err := client.ContainerStart(ctx, resp.ID); err != nil {
		client.ContainerRemove(ctx, resp.ID, dockertypes.ContainerRemoveOptions{
			Force: true,
		})
		return "", fmt.Errorf("failed to start CRI proxy container: %v", err)
	}
	glog.Infof("Started container %s", resp.ID)
	return resp.ID, nil
}

func (b *Bootstrap) initClientset() error {
	var err error
	var config *rest.Config
	if b.config.ApiServerHost == "" {
		config, err = rest.InClusterConfig()
		if err != nil {
			return fmt.Errorf("failed to get REST client config: %v", err)
		}
	} else {
		config = &rest.Config{Host: b.config.ApiServerHost}
	}
	b.clientset, err = kubernetes.NewForConfig(config)
	if err != nil {
		return fmt.Errorf("failed to create ClientSet: %v", err)
	}
	return nil
}

// EnsureCRIProxy checks whether kubelet configuration file exists
// and performs CRI proxy bootstrap procedure if it doesn't.
func (b *Bootstrap) EnsureCRIProxy() (bool, error) {
	if b.config.ConfigzBaseUrl == "" || b.config.ProxyPath == "" || b.config.ProxySocketPath == "" {
		return false, errors.New("invalid BootstrapConfig")
	}

	if !strings.HasPrefix(path.Clean(b.config.ProxySocketPath), "/run/") {
		return false, fmt.Errorf("proxy socket path %q not under /run", b.config.ProxySocketPath)
	}

	if err := b.obtainKubeletConfig(); err != nil {
		return false, err
	}

	dockerEndpoint := b.config.NodeInfo.DockerEndpoint
	glog.V(1).Infof("Using docker endpoint %q", dockerEndpoint)
	if _, err := b.installCriProxyContainer(dockerEndpoint, dockerEndpoint); err != nil {
		return false, err
	}

	if err := waitForSocket(b.config.ProxySocketPath, waitForCriProxySocketNumAttempts, nil); err != nil {
		return false, err
	}

	// We don't try to patch kubelet config before the container
	// is installed, because this process may die with kubelet
	// restart
	if err := b.initClientset(); err != nil {
		return false, err
	}
	if err := b.patchKubeletConfig(); err != nil {
		return false, err
	}

	for {
		time.Sleep(kubeletConfigPollInterval)
		if b.kubeletReadyAfterPatch() {
			break
		}
	}

	glog.V(1).Info("CRI proxy bootstrap complete")
	return true, nil
}
<|endoftext|>


<commit_before>package knuckles

import (
  "fmt"
  "io"
  "log"
  "net"
  "net/http"
  "net/url"
  "strconv"
  "strings"
  "time"
)

type HTTPProxyConfig struct {
  XForwardedFor         bool
  XRequestStart         bool
  XForwardedProto       string
  Store                 Store
  Addr                  string
  RedirectNoHostname    string
  RedirectNoBackend     string
  RedirectInternalError string
}

type HTTPProxy struct {
  Server   http.Server
  listener net.Listener
  Config   HTTPProxyConfig
}

func NewHTTPProxy(config HTTPProxyConfig) (*HTTPProxy, error) {
  h := &HTTPProxy{
    Config: config,
  }

  mux := http.NewServeMux()
  mux.Handle("/", h)
  h.Server.Handler = mux

  return h, nil
}

func (h *HTTPProxy) Start() error {
  var err error

  h.listener, err = net.Listen("tcp", h.Config.Addr)
  if err != nil {
    return err
  }

  return h.Server.Serve(h.listener)
}

func (h *HTTPProxy) Stop() error {
  return h.listener.Close()
}

func (h *HTTPProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  hostname := r.Host
  log.Println("Request for", hostname)

  if sep := strings.Index(hostname, ":"); sep >= 0 {
    hostname = hostname[:sep]
  }

  // tag before starting redis queries
  if h.Config.XRequestStart {
    r.Header.Set("X-Request-Start", requestStart())
  }

  if h.Config.XForwardedProto != "" {
    r.Header.Set("X-Forwarded-Proto", h.Config.XForwardedProto)
  }

  if h.Config.XForwardedFor {
    r.Header.Set("X-Forwarded-For", clientIP(r))
  }

  endpoint, err := h.Config.Store.EndpointForHostname(hostname)

  if err != nil {
    h.clientErr(w, r, err)
    return
  }

  r.URL.Host = endpoint.Addr()
  r.URL.Scheme = "http"

  connection := r.Header.Get("Connection")

  if strings.ToLower(connection) == "upgrade" {
    h.wsProxy(w, r)
  } else {
    h.simpleProxy(w, r)
  }
}

func (h *HTTPProxy) simpleProxy(w http.ResponseWriter, r *http.Request) {

  tr := &http.Transport{
    DisableKeepAlives: true,
  }

  resp, err := tr.RoundTrip(r)

  if err != nil {
    h.clientErr(w, r, err)
    return
  }

  defer resp.Body.Close()

  for name, values := range resp.Header {
    for _, val := range values {
      w.Header().Add(name, val)
    }
  }

  w.WriteHeader(resp.StatusCode)
  // TODO: check copy
  io.Copy(w, resp.Body)
}

func (h *HTTPProxy) wsProxy(w http.ResponseWriter, r *http.Request) {
  hj, ok := w.(http.Hijacker)

  if !ok {
    h.clientErr(w, r, ErrInvalidAction)
    return
  }

  client, _, err := hj.Hijack()
  if err != nil {
    h.clientErr(w, r, err)
    return
  }
  defer client.Close()

  server, err := net.Dial("tcp", r.URL.Host)
  if err != nil {
    h.clientErr(w, r, err)
    return
  }
  defer server.Close()

  err = r.Write(server)
  if err != nil {
    h.clientErr(w, r, err)
    return
  }

  go passBytes(client, server)
}

func passBytes(client, server net.Conn) {
  pass := func(from, to net.Conn, done chan error) {
    _, err := io.Copy(to, from)
    done <- err
  }

  done := make(chan error, 2)

  go pass(client, server, done)
  go pass(server, client, done)

  <-done
  client.Close()
  server.Close()
  <-done
  close(done)
}

func requestStart() string {
  return strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10)
}

func clientIP(r *http.Request) string {
  raddr := strings.Split(r.RemoteAddr, ":")

  if len(raddr) == 0 {
    return ""
  }

  return raddr[0]
}

func (h *HTTPProxy) clientErr(w http.ResponseWriter, r *http.Request, inputErr error) {
  var redirect string

  switch inputErr {
  case ErrNoBackend:
    redirect = h.Config.RedirectNoBackend
  case ErrDeadBackend:
    redirect = h.Config.RedirectNoBackend
  case ErrNoHostname:
    redirect = h.Config.RedirectNoHostname
  default:
    redirect = h.Config.RedirectInternalError
  }

  finalURL := fmt.Sprintf("%s?err=%s", redirect, url.QueryEscape(inputErr.Error()))
  http.Redirect(w, r, finalURL, http.StatusTemporaryRedirect)
}
<commit_msg>removing useless log<commit_after>package knuckles

import (
  "fmt"
  "io"
  "net"
  "net/http"
  "net/url"
  "strconv"
  "strings"
  "time"
)

type HTTPProxyConfig struct {
  XForwardedFor         bool
  XRequestStart         bool
  XForwardedProto       string
  Store                 Store
  Addr                  string
  RedirectNoHostname    string
  RedirectNoBackend     string
  RedirectInternalError string
}

type HTTPProxy struct {
  Server   http.Server
  listener net.Listener
  Config   HTTPProxyConfig
}

func NewHTTPProxy(config HTTPProxyConfig) (*HTTPProxy, error) {
  h := &HTTPProxy{
    Config: config,
  }

  mux := http.NewServeMux()
  mux.Handle("/", h)
  h.Server.Handler = mux

  return h, nil
}

func (h *HTTPProxy) Start() error {
  var err error

  h.listener, err = net.Listen("tcp", h.Config.Addr)
  if err != nil {
    return err
  }

  return h.Server.Serve(h.listener)
}

func (h *HTTPProxy) Stop() error {
  return h.listener.Close()
}

func (h *HTTPProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  hostname := r.Host

  if sep := strings.Index(hostname, ":"); sep >= 0 {
    hostname = hostname[:sep]
  }

  // tag before starting redis queries
  if h.Config.XRequestStart {
    r.Header.Set("X-Request-Start", requestStart())
  }

  if h.Config.XForwardedProto != "" {
    r.Header.Set("X-Forwarded-Proto", h.Config.XForwardedProto)
  }

  if h.Config.XForwardedFor {
    r.Header.Set("X-Forwarded-For", clientIP(r))
  }

  endpoint, err := h.Config.Store.EndpointForHostname(hostname)

  if err != nil {
    h.clientErr(w, r, err)
    return
  }

  r.URL.Host = endpoint.Addr()
  r.URL.Scheme = "http"

  connection := r.Header.Get("Connection")

  if strings.ToLower(connection) == "upgrade" {
    h.wsProxy(w, r)
  } else {
    h.simpleProxy(w, r)
  }
}

func (h *HTTPProxy) simpleProxy(w http.ResponseWriter, r *http.Request) {

  tr := &http.Transport{
    DisableKeepAlives: true,
  }

  resp, err := tr.RoundTrip(r)

  if err != nil {
    h.clientErr(w, r, err)
    return
  }

  defer resp.Body.Close()

  for name, values := range resp.Header {
    for _, val := range values {
      w.Header().Add(name, val)
    }
  }

  w.WriteHeader(resp.StatusCode)
  // TODO: check copy
  io.Copy(w, resp.Body)
}

func (h *HTTPProxy) wsProxy(w http.ResponseWriter, r *http.Request) {
  hj, ok := w.(http.Hijacker)

  if !ok {
    h.clientErr(w, r, ErrInvalidAction)
    return
  }

  client, _, err := hj.Hijack()
  if err != nil {
    h.clientErr(w, r, err)
    return
  }
  defer client.Close()

  server, err := net.Dial("tcp", r.URL.Host)
  if err != nil {
    h.clientErr(w, r, err)
    return
  }
  defer server.Close()

  err = r.Write(server)
  if err != nil {
    h.clientErr(w, r, err)
    return
  }

  go passBytes(client, server)
}

func passBytes(client, server net.Conn) {
  pass := func(from, to net.Conn, done chan error) {
    _, err := io.Copy(to, from)
    done <- err
  }

  done := make(chan error, 2)

  go pass(client, server, done)
  go pass(server, client, done)

  <-done
  client.Close()
  server.Close()
  <-done
  close(done)
}

func requestStart() string {
  return strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10)
}

func clientIP(r *http.Request) string {
  raddr := strings.Split(r.RemoteAddr, ":")

  if len(raddr) == 0 {
    return ""
  }

  return raddr[0]
}

func (h *HTTPProxy) clientErr(w http.ResponseWriter, r *http.Request, inputErr error) {
  var redirect string

  switch inputErr {
  case ErrNoBackend:
    redirect = h.Config.RedirectNoBackend
  case ErrDeadBackend:
    redirect = h.Config.RedirectNoBackend
  case ErrNoHostname:
    redirect = h.Config.RedirectNoHostname
  default:
    redirect = h.Config.RedirectInternalError
  }

  finalURL := fmt.Sprintf("%s?err=%s", redirect, url.QueryEscape(inputErr.Error()))
  http.Redirect(w, r, finalURL, http.StatusTemporaryRedirect)
}
<|endoftext|>


<commit_before>/*******************************************************************************
*
* Copyright 2018 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, 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 plugins

import (
	"github.com/gophercloud/gophercloud"
	"github.com/gophercloud/gophercloud/openstack"
	"github.com/gophercloud/gophercloud/openstack/identity/v3/tokens"
	"github.com/sapcc/go-bits/logg"
)

type cfmClient struct {
	*gophercloud.ServiceClient
	projectID string
}

var cfmProjectIDCache = map[string]string{}

func getProjectIDForToken(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts) (string, error) {
	//The CFM API is stupid and needs the caller to provide the scope of the
	//token redundantly in the X-Project-Id header.

	//cache the token -> project_id mapping
	token := provider.Token()
	if projectID, exists := cfmProjectIDCache[token]; exists {
		return projectID, nil
	}

	identityClient, err := openstack.NewIdentityV3(provider, eo)
	if err != nil {
		return "", err
	}
	project, err := tokens.Get(identityClient, token).ExtractProject()
	if err != nil {
		return "", err
	}
	//if our token changed mid-request because it expired and the 401 response
	//triggered Gophercloud to reauthenticate, then we cannot trust the response
	if provider.Token() != token {
		//restart this call
		return getProjectIDForToken(provider, eo)
	}

	cfmProjectIDCache[token] = project.ID
	return project.ID, nil
}

func newCFMClient(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts) (*cfmClient, error) {
	serviceType := "database"
	eoCFM := eo
	eoCFM.ApplyDefaults(serviceType)

	url, err := provider.EndpointLocator(eoCFM)
	if err != nil {
		return nil, err
	}

	ourProjectID, err := getProjectIDForToken(provider, eo)
	return &cfmClient{
		ServiceClient: &gophercloud.ServiceClient{
			ProviderClient: provider,
			Endpoint:       url,
			Type:           serviceType,
		},
		projectID: ourProjectID,
	}, nil
}

func (c *cfmClient) reqOpts(okCodes ...int) *gophercloud.RequestOpts {
	return &gophercloud.RequestOpts{
		OkCodes: okCodes,
		MoreHeaders: map[string]string{
			"X-Auth-Token":  "",
			"Authorization": "Token " + c.ProviderClient.Token(),
			"X-Project-ID":  c.projectID,
		},
		ErrorContext: cfmNotFoundError{},
	}
}

////////////////////////////////////////////////////////////////////////////////
// new-style quota/usage API

func (c *cfmClient) GetQuotaSet(projectID string) (result gophercloud.Result) {
	url := c.ServiceURL("v1.0", "quota-sets", projectID) + "/"
	_, result.Err = c.Get(url, &result.Body, c.reqOpts(200))
	return
}

func (c *cfmClient) CreateQuotaSet(projectID string, quotaBytes uint64) error {
	body := struct {
		StorageQuota struct {
			ProjectID string `json:"project_id"`
			SizeLimit uint64 `json:"size_limit"`
		} `json:"storage_quota"`
	}{}
	body.StorageQuota.ProjectID = projectID
	body.StorageQuota.SizeLimit = quotaBytes
	url := c.ServiceURL("v1.0", "quota-sets") + "/"
	_, err := c.Post(url, body, nil, c.reqOpts(202))
	return err
}

//UpdateQuotaSet may return cfmNotFoundError.
func (c *cfmClient) UpdateQuotaSet(projectID string, quotaBytes uint64) error {
	body := struct {
		StorageQuota struct {
			SizeLimit uint64 `json:"size_limit"`
		} `json:"storage_quota"`
	}{}
	body.StorageQuota.SizeLimit = quotaBytes
	url := c.ServiceURL("v1.0", "quota-sets", projectID) + "/"
	_, err := c.Put(url, body, nil, c.reqOpts(200))
	return err
}

//An error type that can be used in gophercloud's weird ErrorContext interface.
type cfmNotFoundError struct{}

//Error implements the builtin/error interface.
func (cfmNotFoundError) Error() string {
	return "not found"
}

//Error404 implements the gophercloud.Err404er interface.
func (cfmNotFoundError) Error404(gophercloud.ErrUnexpectedResponseCode) error {
	return cfmNotFoundError{}
}

////////////////////////////////////////////////////////////////////////////////
// old-style usage API

type cfmShareserver struct {
	Type        string
	ProjectUUID string
	DetailsURL  string
	//fields that are only filled by cfmGetShareserver, not by cfmListShareservers
	BytesUsed uint64
}

func (c *cfmClient) ListShareservers() ([]cfmShareserver, error) {
	url := c.ServiceURL("v1.0", "shareservers") + "/"
	var result gophercloud.Result
	_, err := c.Get(url, &result.Body, c.reqOpts(200))
	if err != nil {
		return nil, err
	}

	var data struct {
		Shareservers []struct {
			ID          string `json:"id"`
			Type        string `json:"type"`
			ProjectUUID string `json:"customer_id"`
		} `json:"shareservers"`
	}
	err = result.ExtractInto(&data)
	if err != nil {
		return nil, err
	}

	list := make([]cfmShareserver, len(data.Shareservers))
	for idx, srv := range data.Shareservers {
		list[idx] = cfmShareserver{
			Type:        srv.Type,
			ProjectUUID: srv.ProjectUUID,
			DetailsURL:  url + srv.ID + "/",
		}
	}
	return list, nil
}

func (c *cfmClient) GetShareserver(url string) (*cfmShareserver, error) {
	var result gophercloud.Result
	_, err := c.Get(url, &result.Body, c.reqOpts(200))
	if err != nil {
		return nil, err
	}

	var data struct {
		Shareserver struct {
			ID         string `json:"id"`
			Properties struct {
				SVMs []struct {
					Volumes []struct {
						Space struct {
							BytesUsed uint64 `json:"size"`
						} `json:"space"`
					} `json:"volumes"`
				} `json:"svms"`
			} `json:"properties"`
			Type        string `json:"type"`
			ProjectUUID string `json:"customer_id"`
		} `json:"shareserver"`
	}
	err = result.ExtractInto(&data)
	if err != nil {
		return nil, err
	}

	srv := data.Shareserver

	var totalBytesUsed uint64
	for _, svm := range srv.Properties.SVMs {
		for _, volume := range svm.Volumes {
			totalBytesUsed += volume.Space.BytesUsed
		}
	}

	logg.Info("CFM shareserver %s (type %s) in project %s has size = %d bytes",
		srv.ID, srv.Type, srv.ProjectUUID,
		totalBytesUsed,
	)
	return &cfmShareserver{
		Type:        srv.Type,
		ProjectUUID: srv.ProjectUUID,
		DetailsURL:  url + srv.ID + "/",
		BytesUsed:   totalBytesUsed,
	}, nil
}

////////////////////////////////////////////////////////////////////////////////
// capacity API

type cfmPool struct {
	HostName     string `json:"host"`
	Name         string `json:"name"`
	Type         string `json:"pool"`
	Capabilities struct {
		TotalCapacityBytes uint64 `json:"total_capacity"`
	} `json:"capabilities"`
}

func (c *cfmClient) ListPools() ([]cfmPool, error) {
	url := c.ServiceURL("v1.0", "scheduler-stats", "pools", "detail") + "/"
	var result gophercloud.Result
	_, err := c.Get(url, &result.Body, c.reqOpts(200))
	if err != nil {
		return nil, err
	}

	var data struct {
		Pools []cfmPool `json:"pools"`
	}
	err = result.ExtractInto(&data)
	return data.Pools, err
}
<commit_msg>use new AuthResult API to avoid duplicate Keystone requests<commit_after>/*******************************************************************************
*
* Copyright 2018 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, 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 plugins

import (
	"github.com/gophercloud/gophercloud"
	"github.com/gophercloud/gophercloud/openstack"
	"github.com/gophercloud/gophercloud/openstack/identity/v3/tokens"
	"github.com/sapcc/go-bits/logg"
)

type cfmClient struct {
	*gophercloud.ServiceClient
	projectID string
}

var cfmProjectIDCache = map[string]string{}

func getProjectIDForToken(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts) (string, error) {
	//The CFM API is stupid and needs the caller to provide the scope of the
	//token redundantly in the X-Project-Id header.

	//try to use the new AuthResult API to get the token without extra HTTP requests
	if result, ok := provider.GetAuthResult().(tokens.CreateResult); ok {
		project, err := result.ExtractProject()
		if err == nil {
			return project.ID, nil
		}
	}

	//cache the token -> project_id mapping
	token := provider.Token()
	if projectID, exists := cfmProjectIDCache[token]; exists {
		return projectID, nil
	}

	identityClient, err := openstack.NewIdentityV3(provider, eo)
	if err != nil {
		return "", err
	}
	project, err := tokens.Get(identityClient, token).ExtractProject()
	if err != nil {
		return "", err
	}
	//if our token changed mid-request because it expired and the 401 response
	//triggered Gophercloud to reauthenticate, then we cannot trust the response
	if provider.Token() != token {
		//restart this call
		return getProjectIDForToken(provider, eo)
	}

	cfmProjectIDCache[token] = project.ID
	return project.ID, nil
}

func newCFMClient(provider *gophercloud.ProviderClient, eo gophercloud.EndpointOpts) (*cfmClient, error) {
	serviceType := "database"
	eoCFM := eo
	eoCFM.ApplyDefaults(serviceType)

	url, err := provider.EndpointLocator(eoCFM)
	if err != nil {
		return nil, err
	}

	ourProjectID, err := getProjectIDForToken(provider, eo)
	return &cfmClient{
		ServiceClient: &gophercloud.ServiceClient{
			ProviderClient: provider,
			Endpoint:       url,
			Type:           serviceType,
		},
		projectID: ourProjectID,
	}, nil
}

func (c *cfmClient) reqOpts(okCodes ...int) *gophercloud.RequestOpts {
	return &gophercloud.RequestOpts{
		OkCodes: okCodes,
		MoreHeaders: map[string]string{
			"X-Auth-Token":  "",
			"Authorization": "Token " + c.ProviderClient.Token(),
			"X-Project-ID":  c.projectID,
		},
		ErrorContext: cfmNotFoundError{},
	}
}

////////////////////////////////////////////////////////////////////////////////
// new-style quota/usage API

func (c *cfmClient) GetQuotaSet(projectID string) (result gophercloud.Result) {
	url := c.ServiceURL("v1.0", "quota-sets", projectID) + "/"
	_, result.Err = c.Get(url, &result.Body, c.reqOpts(200))
	return
}

func (c *cfmClient) CreateQuotaSet(projectID string, quotaBytes uint64) error {
	body := struct {
		StorageQuota struct {
			ProjectID string `json:"project_id"`
			SizeLimit uint64 `json:"size_limit"`
		} `json:"storage_quota"`
	}{}
	body.StorageQuota.ProjectID = projectID
	body.StorageQuota.SizeLimit = quotaBytes
	url := c.ServiceURL("v1.0", "quota-sets") + "/"
	_, err := c.Post(url, body, nil, c.reqOpts(202))
	return err
}

//UpdateQuotaSet may return cfmNotFoundError.
func (c *cfmClient) UpdateQuotaSet(projectID string, quotaBytes uint64) error {
	body := struct {
		StorageQuota struct {
			SizeLimit uint64 `json:"size_limit"`
		} `json:"storage_quota"`
	}{}
	body.StorageQuota.SizeLimit = quotaBytes
	url := c.ServiceURL("v1.0", "quota-sets", projectID) + "/"
	_, err := c.Put(url, body, nil, c.reqOpts(200))
	return err
}

//An error type that can be used in gophercloud's weird ErrorContext interface.
type cfmNotFoundError struct{}

//Error implements the builtin/error interface.
func (cfmNotFoundError) Error() string {
	return "not found"
}

//Error404 implements the gophercloud.Err404er interface.
func (cfmNotFoundError) Error404(gophercloud.ErrUnexpectedResponseCode) error {
	return cfmNotFoundError{}
}

////////////////////////////////////////////////////////////////////////////////
// old-style usage API

type cfmShareserver struct {
	Type        string
	ProjectUUID string
	DetailsURL  string
	//fields that are only filled by cfmGetShareserver, not by cfmListShareservers
	BytesUsed uint64
}

func (c *cfmClient) ListShareservers() ([]cfmShareserver, error) {
	url := c.ServiceURL("v1.0", "shareservers") + "/"
	var result gophercloud.Result
	_, err := c.Get(url, &result.Body, c.reqOpts(200))
	if err != nil {
		return nil, err
	}

	var data struct {
		Shareservers []struct {
			ID          string `json:"id"`
			Type        string `json:"type"`
			ProjectUUID string `json:"customer_id"`
		} `json:"shareservers"`
	}
	err = result.ExtractInto(&data)
	if err != nil {
		return nil, err
	}

	list := make([]cfmShareserver, len(data.Shareservers))
	for idx, srv := range data.Shareservers {
		list[idx] = cfmShareserver{
			Type:        srv.Type,
			ProjectUUID: srv.ProjectUUID,
			DetailsURL:  url + srv.ID + "/",
		}
	}
	return list, nil
}

func (c *cfmClient) GetShareserver(url string) (*cfmShareserver, error) {
	var result gophercloud.Result
	_, err := c.Get(url, &result.Body, c.reqOpts(200))
	if err != nil {
		return nil, err
	}

	var data struct {
		Shareserver struct {
			ID         string `json:"id"`
			Properties struct {
				SVMs []struct {
					Volumes []struct {
						Space struct {
							BytesUsed uint64 `json:"size"`
						} `json:"space"`
					} `json:"volumes"`
				} `json:"svms"`
			} `json:"properties"`
			Type        string `json:"type"`
			ProjectUUID string `json:"customer_id"`
		} `json:"shareserver"`
	}
	err = result.ExtractInto(&data)
	if err != nil {
		return nil, err
	}

	srv := data.Shareserver

	var totalBytesUsed uint64
	for _, svm := range srv.Properties.SVMs {
		for _, volume := range svm.Volumes {
			totalBytesUsed += volume.Space.BytesUsed
		}
	}

	logg.Info("CFM shareserver %s (type %s) in project %s has size = %d bytes",
		srv.ID, srv.Type, srv.ProjectUUID,
		totalBytesUsed,
	)
	return &cfmShareserver{
		Type:        srv.Type,
		ProjectUUID: srv.ProjectUUID,
		DetailsURL:  url + srv.ID + "/",
		BytesUsed:   totalBytesUsed,
	}, nil
}

////////////////////////////////////////////////////////////////////////////////
// capacity API

type cfmPool struct {
	HostName     string `json:"host"`
	Name         string `json:"name"`
	Type         string `json:"pool"`
	Capabilities struct {
		TotalCapacityBytes uint64 `json:"total_capacity"`
	} `json:"capabilities"`
}

func (c *cfmClient) ListPools() ([]cfmPool, error) {
	url := c.ServiceURL("v1.0", "scheduler-stats", "pools", "detail") + "/"
	var result gophercloud.Result
	_, err := c.Get(url, &result.Body, c.reqOpts(200))
	if err != nil {
		return nil, err
	}

	var data struct {
		Pools []cfmPool `json:"pools"`
	}
	err = result.ExtractInto(&data)
	return data.Pools, err
}
<|endoftext|>


<commit_before>// Copyright 2016 The prometheus-operator 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 prometheus

import (
	"fmt"
	"regexp"
	"strings"

	yaml "gopkg.in/yaml.v2"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/coreos/prometheus-operator/pkg/client/monitoring/v1alpha1"
)

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

func sanitizeLabelName(name string) string {
	return invalidLabelCharRE.ReplaceAllString(name, "_")
}

func generateConfig(p *v1alpha1.Prometheus, mons map[string]*v1alpha1.ServiceMonitor) ([]byte, error) {
	cfg := map[string]interface{}{}

	cfg["global"] = map[string]string{
		"evaluation_interval": "30s",
		"scrape_interval":     "30s",
	}

	cfg["rule_files"] = []string{"/etc/prometheus/rules/*.rules"}

	var scrapeConfigs []interface{}
	for _, mon := range mons {
		for i, ep := range mon.Spec.Endpoints {
			scrapeConfigs = append(scrapeConfigs, generateServiceMonitorConfig(mon, ep, i))
		}
	}
	var alertmanagerConfigs []interface{}
	for _, am := range p.Spec.Alerting.Alertmanagers {
		alertmanagerConfigs = append(alertmanagerConfigs, generateAlertmanagerConfig(am))
	}

	cfg["scrape_configs"] = scrapeConfigs
	cfg["alerting"] = map[string]interface{}{
		"alertmanagers": alertmanagerConfigs,
	}

	return yaml.Marshal(cfg)
}

func generateServiceMonitorConfig(m *v1alpha1.ServiceMonitor, ep v1alpha1.Endpoint, i int) interface{} {
	cfg := map[string]interface{}{
		"job_name": fmt.Sprintf("%s/%s/%d", m.Namespace, m.Name, i),
		"kubernetes_sd_configs": []map[string]interface{}{
			{
				"role": "endpoints",
			},
		},
	}

	if ep.Interval != "" {
		cfg["scrape_interval"] = ep.Interval
	}
	if ep.Path != "" {
		cfg["metrics_path"] = ep.Path
	}
	if ep.Scheme != "" {
		cfg["scheme"] = ep.Scheme
	}
	if ep.TLSConfig != nil {
		tlsConfig := map[string]interface{}{
			"insecure_skip_verify": ep.TLSConfig.InsecureSkipVerify,
		}
		if ep.TLSConfig.CAFile != "" {
			tlsConfig["ca_file"] = ep.TLSConfig.CAFile
		}
		if ep.TLSConfig.CertFile != "" {
			tlsConfig["cert_file"] = ep.TLSConfig.CertFile
		}
		if ep.TLSConfig.KeyFile != "" {
			tlsConfig["key_file"] = ep.TLSConfig.KeyFile
		}
		if ep.TLSConfig.ServerName != "" {
			tlsConfig["server_name"] = ep.TLSConfig.ServerName
		}
		cfg["tls_config"] = tlsConfig
	}
	if ep.BearerTokenFile != "" {
		cfg["bearer_token_file"] = ep.BearerTokenFile
	}

	var relabelings []interface{}

	// Filter targets by services selected by the monitor.

	// Exact label matches.
	for k, v := range m.Spec.Selector.MatchLabels {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(k)},
			"regex":         v,
		})
	}
	// Set based label matching. We have to map the valid relations
	// `In`, `NotIn`, `Exists`, and `DoesNotExist`, into relabeling rules.
	for _, exp := range m.Spec.Selector.MatchExpressions {
		switch exp.Operator {
		case metav1.LabelSelectorOpIn:
			relabelings = append(relabelings, map[string]interface{}{
				"action":        "keep",
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(exp.Key)},
				"regex":         strings.Join(exp.Values, "|"),
			})
		case metav1.LabelSelectorOpNotIn:
			relabelings = append(relabelings, map[string]interface{}{
				"action":        "drop",
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(exp.Key)},
				"regex":         strings.Join(exp.Values, "|"),
			})
		case metav1.LabelSelectorOpExists:
			relabelings = append(relabelings, map[string]interface{}{
				"action":        "keep",
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(exp.Key)},
				"regex":         ".+",
			})
		case metav1.LabelSelectorOpDoesNotExist:
			relabelings = append(relabelings, map[string]interface{}{
				"action":        "drop",
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(exp.Key)},
				"regex":         ".+",
			})
		}
	}

	// Filter targets based on the namespace selection configuration.
	// By default we only discover services within the namespace of the
	// ServiceMonitor.
	// Selections allow extending this to all namespaces or to a subset
	// of them specified by label or name matching.
	//
	// Label selections are not supported yet as they require either supported
	// in the upstream SD integration or require out-of-band implementation
	// in the operator with configuration reload.
	//
	// There's no explicit nil for the selector, we decide for the default
	// case if it's all zero values.
	nsel := m.Spec.NamespaceSelector

	if !nsel.Any && len(nsel.MatchNames) == 0 {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_namespace"},
			"regex":         m.Namespace,
		})
	} else if len(nsel.MatchNames) > 0 {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_namespace"},
			"regex":         strings.Join(nsel.MatchNames, "|"),
		})
	}

	// Filter targets based on correct port for the endpoint.
	if ep.Port != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_endpoint_port_name"},
			"regex":         ep.Port,
		})
	} else if ep.TargetPort.StrVal != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_container_port_name"},
			"regex":         ep.TargetPort.String(),
		})
	} else if ep.TargetPort.IntVal != 0 {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_container_port_number"},
			"regex":         ep.TargetPort.String(),
		})
	}

	// Relabel namespace and pod and service labels into proper labels.
	relabelings = append(relabelings, []interface{}{
		map[string]interface{}{
			"source_labels": []string{"__meta_kubernetes_namespace"},
			"target_label":  "namespace",
		},
		map[string]interface{}{
			"action":      "labelmap",
			"regex":       "__meta_kubernetes_service_label_(.+)",
			"replacement": "svc_$1",
		},
		map[string]interface{}{
			"action":       "replace",
			"target_label": "__meta_kubernetes_pod_label_pod_template_hash",
			"replacement":  "",
		},
		map[string]interface{}{
			"action":      "labelmap",
			"regex":       "__meta_kubernetes_pod_label_(.+)",
			"replacement": "pod_$1",
		},
	}...)

	// By default, generate a safe job name from the service name and scraped port.
	// We also keep this around if a jobLabel is set in case the targets don't actually
	// have a value for it.
	if ep.Port != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"source_labels": []string{"__meta_kubernetes_service_name"},
			"target_label":  "job",
			"replacement":   "${1}-" + ep.Port,
		})
	} else if ep.TargetPort.String() != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"source_labels": []string{"__meta_kubernetes_service_name"},
			"target_label":  "job",
			"replacement":   "${1}-" + ep.TargetPort.String(),
		})
	}
	// Generate a job name with a base label. Same as above, just that we
	// get the base from the label, if present.
	if m.Spec.JobLabel != "" {
		if ep.Port != "" {
			relabelings = append(relabelings, map[string]interface{}{
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(m.Spec.JobLabel)},
				"target_label":  "job",
				"regex":         "(.+)",
				"replacement":   "${1}-" + ep.Port,
			})
		} else if ep.TargetPort.String() != "" {
			relabelings = append(relabelings, map[string]interface{}{
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(m.Spec.JobLabel)},
				"target_label":  "job",
				"regex":         "(.+)",
				"replacement":   "${1}-" + ep.TargetPort.String(),
			})
		}
	}

	cfg["relabel_configs"] = relabelings

	return cfg
}

func generateAlertmanagerConfig(am v1alpha1.AlertmanagerEndpoints) interface{} {
	if am.Scheme == "" {
		am.Scheme = "http"
	}
	cfg := map[string]interface{}{
		"kubernetes_sd_configs": []map[string]interface{}{
			{
				"role": "endpoints",
			},
		},
		"scheme": am.Scheme,
	}

	var relabelings []interface{}

	relabelings = append(relabelings, map[string]interface{}{
		"action":        "keep",
		"source_labels": []string{"__meta_kubernetes_service_name"},
		"regex":         am.Name,
	})
	relabelings = append(relabelings, map[string]interface{}{
		"action":        "keep",
		"source_labels": []string{"__meta_kubernetes_namespace"},
		"regex":         am.Namespace,
	})

	if am.Port.StrVal != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_endpoint_port_name"},
			"regex":         am.Port.String(),
		})
	} else if am.Port.IntVal != 0 {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_container_port_number"},
			"regex":         am.Port.String(),
		})
	}

	cfg["relabel_configs"] = relabelings

	return cfg
}
<commit_msg>prometheus: put endpoint in separate label<commit_after>// Copyright 2016 The prometheus-operator 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 prometheus

import (
	"fmt"
	"regexp"
	"strings"

	yaml "gopkg.in/yaml.v2"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/coreos/prometheus-operator/pkg/client/monitoring/v1alpha1"
)

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

func sanitizeLabelName(name string) string {
	return invalidLabelCharRE.ReplaceAllString(name, "_")
}

func generateConfig(p *v1alpha1.Prometheus, mons map[string]*v1alpha1.ServiceMonitor) ([]byte, error) {
	cfg := map[string]interface{}{}

	cfg["global"] = map[string]string{
		"evaluation_interval": "30s",
		"scrape_interval":     "30s",
	}

	cfg["rule_files"] = []string{"/etc/prometheus/rules/*.rules"}

	var scrapeConfigs []interface{}
	for _, mon := range mons {
		for i, ep := range mon.Spec.Endpoints {
			scrapeConfigs = append(scrapeConfigs, generateServiceMonitorConfig(mon, ep, i))
		}
	}
	var alertmanagerConfigs []interface{}
	for _, am := range p.Spec.Alerting.Alertmanagers {
		alertmanagerConfigs = append(alertmanagerConfigs, generateAlertmanagerConfig(am))
	}

	cfg["scrape_configs"] = scrapeConfigs
	cfg["alerting"] = map[string]interface{}{
		"alertmanagers": alertmanagerConfigs,
	}

	return yaml.Marshal(cfg)
}

func generateServiceMonitorConfig(m *v1alpha1.ServiceMonitor, ep v1alpha1.Endpoint, i int) interface{} {
	cfg := map[string]interface{}{
		"job_name": fmt.Sprintf("%s/%s/%d", m.Namespace, m.Name, i),
		"kubernetes_sd_configs": []map[string]interface{}{
			{
				"role": "endpoints",
			},
		},
	}

	if ep.Interval != "" {
		cfg["scrape_interval"] = ep.Interval
	}
	if ep.Path != "" {
		cfg["metrics_path"] = ep.Path
	}
	if ep.Scheme != "" {
		cfg["scheme"] = ep.Scheme
	}
	if ep.TLSConfig != nil {
		tlsConfig := map[string]interface{}{
			"insecure_skip_verify": ep.TLSConfig.InsecureSkipVerify,
		}
		if ep.TLSConfig.CAFile != "" {
			tlsConfig["ca_file"] = ep.TLSConfig.CAFile
		}
		if ep.TLSConfig.CertFile != "" {
			tlsConfig["cert_file"] = ep.TLSConfig.CertFile
		}
		if ep.TLSConfig.KeyFile != "" {
			tlsConfig["key_file"] = ep.TLSConfig.KeyFile
		}
		if ep.TLSConfig.ServerName != "" {
			tlsConfig["server_name"] = ep.TLSConfig.ServerName
		}
		cfg["tls_config"] = tlsConfig
	}
	if ep.BearerTokenFile != "" {
		cfg["bearer_token_file"] = ep.BearerTokenFile
	}

	var relabelings []interface{}

	// Filter targets by services selected by the monitor.

	// Exact label matches.
	for k, v := range m.Spec.Selector.MatchLabels {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(k)},
			"regex":         v,
		})
	}
	// Set based label matching. We have to map the valid relations
	// `In`, `NotIn`, `Exists`, and `DoesNotExist`, into relabeling rules.
	for _, exp := range m.Spec.Selector.MatchExpressions {
		switch exp.Operator {
		case metav1.LabelSelectorOpIn:
			relabelings = append(relabelings, map[string]interface{}{
				"action":        "keep",
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(exp.Key)},
				"regex":         strings.Join(exp.Values, "|"),
			})
		case metav1.LabelSelectorOpNotIn:
			relabelings = append(relabelings, map[string]interface{}{
				"action":        "drop",
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(exp.Key)},
				"regex":         strings.Join(exp.Values, "|"),
			})
		case metav1.LabelSelectorOpExists:
			relabelings = append(relabelings, map[string]interface{}{
				"action":        "keep",
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(exp.Key)},
				"regex":         ".+",
			})
		case metav1.LabelSelectorOpDoesNotExist:
			relabelings = append(relabelings, map[string]interface{}{
				"action":        "drop",
				"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(exp.Key)},
				"regex":         ".+",
			})
		}
	}

	// Filter targets based on the namespace selection configuration.
	// By default we only discover services within the namespace of the
	// ServiceMonitor.
	// Selections allow extending this to all namespaces or to a subset
	// of them specified by label or name matching.
	//
	// Label selections are not supported yet as they require either supported
	// in the upstream SD integration or require out-of-band implementation
	// in the operator with configuration reload.
	//
	// There's no explicit nil for the selector, we decide for the default
	// case if it's all zero values.
	nsel := m.Spec.NamespaceSelector

	if !nsel.Any && len(nsel.MatchNames) == 0 {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_namespace"},
			"regex":         m.Namespace,
		})
	} else if len(nsel.MatchNames) > 0 {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_namespace"},
			"regex":         strings.Join(nsel.MatchNames, "|"),
		})
	}

	// Filter targets based on correct port for the endpoint.
	if ep.Port != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_endpoint_port_name"},
			"regex":         ep.Port,
		})
	} else if ep.TargetPort.StrVal != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_container_port_name"},
			"regex":         ep.TargetPort.String(),
		})
	} else if ep.TargetPort.IntVal != 0 {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_container_port_number"},
			"regex":         ep.TargetPort.String(),
		})
	}

	// Relabel namespace and pod and service labels into proper labels.
	relabelings = append(relabelings, []interface{}{
		map[string]interface{}{
			"source_labels": []string{"__meta_kubernetes_namespace"},
			"target_label":  "namespace",
		},
		map[string]interface{}{
			"action":      "labelmap",
			"regex":       "__meta_kubernetes_service_label_(.+)",
			"replacement": "svc_$1",
		},
		map[string]interface{}{
			"action":       "replace",
			"target_label": "__meta_kubernetes_pod_label_pod_template_hash",
			"replacement":  "",
		},
		map[string]interface{}{
			"action":      "labelmap",
			"regex":       "__meta_kubernetes_pod_label_(.+)",
			"replacement": "pod_$1",
		},
	}...)

	// By default, generate a safe job name from the service name.  We also keep
	// this around if a jobLabel is set in case the targets don't actually have a
	// value for it. A single service may potentially have multiple metrics
	// endpoints, therefore the endpoints labels is filled with the ports name or
	// as a fallback the port number.

	relabelings = append(relabelings, map[string]interface{}{
		"source_labels": []string{"__meta_kubernetes_service_name"},
		"target_label":  "job",
		"replacement":   "${1}",
	})
	if m.Spec.JobLabel != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"source_labels": []string{"__meta_kubernetes_service_label_" + sanitizeLabelName(m.Spec.JobLabel)},
			"target_label":  "job",
			"regex":         "(.+)",
			"replacement":   "${1}",
		})
	}

	if ep.Port != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"target_label": "endpoint",
			"replacement":  ep.Port,
		})
	} else if ep.TargetPort.String() != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"target_label": "endpoint",
			"replacement":  ep.TargetPort.String(),
		})
	}

	cfg["relabel_configs"] = relabelings

	return cfg
}

func generateAlertmanagerConfig(am v1alpha1.AlertmanagerEndpoints) interface{} {
	if am.Scheme == "" {
		am.Scheme = "http"
	}
	cfg := map[string]interface{}{
		"kubernetes_sd_configs": []map[string]interface{}{
			{
				"role": "endpoints",
			},
		},
		"scheme": am.Scheme,
	}

	var relabelings []interface{}

	relabelings = append(relabelings, map[string]interface{}{
		"action":        "keep",
		"source_labels": []string{"__meta_kubernetes_service_name"},
		"regex":         am.Name,
	})
	relabelings = append(relabelings, map[string]interface{}{
		"action":        "keep",
		"source_labels": []string{"__meta_kubernetes_namespace"},
		"regex":         am.Namespace,
	})

	if am.Port.StrVal != "" {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_endpoint_port_name"},
			"regex":         am.Port.String(),
		})
	} else if am.Port.IntVal != 0 {
		relabelings = append(relabelings, map[string]interface{}{
			"action":        "keep",
			"source_labels": []string{"__meta_kubernetes_container_port_number"},
			"regex":         am.Port.String(),
		})
	}

	cfg["relabel_configs"] = relabelings

	return cfg
}
<|endoftext|>


<commit_before>package realtime

import (
	"encoding/json"
	"strings"

	"github.com/cozy/cozy-stack/pkg/logger"
	redis "github.com/go-redis/redis"
)

const eventsRedisKey = "realtime:events"

type redisHub struct {
	c     *redis.Client
	mem   *memHub
	local *topic
}

func newRedisHub(c *redis.Client) *redisHub {
	local := newTopic("*")
	mem := newMemHub()
	hub := &redisHub{c, mem, local}
	go hub.start()
	return hub
}

type jsonDoc struct {
	M    map[string]interface{}
	Type string
}

func (j jsonDoc) ID() string      { id, _ := j.M["_id"].(string); return id }
func (j jsonDoc) DocType() string { return j.Type }
func (j *jsonDoc) MarshalJSON() ([]byte, error) {
	if j.M == nil {
		return json.Marshal(map[string]string{"_type": j.Type})
	}
	j.M["_type"] = j.Type
	defer delete(j.M, "_type")
	return json.Marshal(j.M)
}

func toJSONDoc(d map[string]interface{}) *jsonDoc {
	if d == nil {
		return nil
	}
	doctype, _ := d["_type"].(string)
	delete(d, "_type")
	return &jsonDoc{d, doctype}
}

type jsonEvent struct {
	Domain string
	Verb   string
	Doc    *jsonDoc
	Old    *jsonDoc
}

func (j *jsonEvent) UnmarshalJSON(buf []byte) error {
	var m map[string]interface{}
	if err := json.Unmarshal(buf, &m); err != nil {
		return err
	}
	j.Domain, _ = m["domain"].(string)
	j.Verb, _ = m["verb"].(string)
	if doc, ok := m["doc"].(map[string]interface{}); ok {
		j.Doc = toJSONDoc(doc)
	}
	if old, ok := m["old"].(map[string]interface{}); ok {
		j.Old = toJSONDoc(old)
	}
	return nil
}

func (h *redisHub) start() {
	sub := h.c.Subscribe(eventsRedisKey)
	log := logger.WithNamespace("realtime-redis")
	for msg := range sub.Channel() {
		je := jsonEvent{}
		parts := strings.SplitN(msg.Payload, ",", 2)
		if len(parts) < 2 {
			log.Warnf("Invalid payload: %s", msg.Payload)
			continue
		}
		doctype := parts[0]
		buf := []byte(parts[1])
		if err := json.Unmarshal(buf, &je); err != nil {
			log.Warnf("Error on start: %s", err)
			continue
		}
		if je.Doc != nil {
			je.Doc.Type = doctype
		}
		if je.Old != nil {
			je.Old.Type = doctype
		}
		e := &Event{
			Domain: je.Domain,
			Verb:   je.Verb,
			Doc:    je.Doc,
			OldDoc: je.Old,
		}
		h.mem.Publish(e)
	}
}

func (h *redisHub) GetTopic(domain, doctype string) *topic {
	return nil
}

func (h *redisHub) Publish(e *Event) {
	h.local.broadcast <- e
	buf, err := json.Marshal(e)
	if err != nil {
		log := logger.WithNamespace("realtime-redis")
		log.Warnf("Error on publish: %s", err)
		return
	}
	h.c.Publish(eventsRedisKey, e.Doc.DocType()+","+string(buf))
}

func (h *redisHub) Subscriber(domain string) *DynamicSubscriber {
	return h.mem.Subscriber(domain)
}

func (h *redisHub) SubscribeLocalAll() *DynamicSubscriber {
	ds := newDynamicSubscriber(nil, "")
	ds.addTopic(h.local, "")
	return ds
}
<commit_msg>Fix concurrent access to map in jsonDoc marshalling<commit_after>package realtime

import (
	"encoding/json"
	"strings"

	"github.com/cozy/cozy-stack/pkg/logger"
	redis "github.com/go-redis/redis"
)

const eventsRedisKey = "realtime:events"

type redisHub struct {
	c     *redis.Client
	mem   *memHub
	local *topic
}

func newRedisHub(c *redis.Client) *redisHub {
	local := newTopic("*")
	mem := newMemHub()
	hub := &redisHub{c, mem, local}
	go hub.start()
	return hub
}

type jsonDoc struct {
	M    map[string]interface{}
	Type string
}

func (j jsonDoc) ID() string      { id, _ := j.M["_id"].(string); return id }
func (j jsonDoc) DocType() string { return j.Type }
func (j *jsonDoc) MarshalJSON() ([]byte, error) {
	m := map[string]interface{}{"_type": j.Type}
	if j.M != nil {
		for k, v := range j.M {
			m[k] = v
		}
	}
	return json.Marshal(m)
}

func toJSONDoc(d map[string]interface{}) *jsonDoc {
	if d == nil {
		return nil
	}
	doctype, _ := d["_type"].(string)
	delete(d, "_type")
	return &jsonDoc{d, doctype}
}

type jsonEvent struct {
	Domain string
	Verb   string
	Doc    *jsonDoc
	Old    *jsonDoc
}

func (j *jsonEvent) UnmarshalJSON(buf []byte) error {
	var m map[string]interface{}
	if err := json.Unmarshal(buf, &m); err != nil {
		return err
	}
	j.Domain, _ = m["domain"].(string)
	j.Verb, _ = m["verb"].(string)
	if doc, ok := m["doc"].(map[string]interface{}); ok {
		j.Doc = toJSONDoc(doc)
	}
	if old, ok := m["old"].(map[string]interface{}); ok {
		j.Old = toJSONDoc(old)
	}
	return nil
}

func (h *redisHub) start() {
	sub := h.c.Subscribe(eventsRedisKey)
	log := logger.WithNamespace("realtime-redis")
	for msg := range sub.Channel() {
		je := jsonEvent{}
		parts := strings.SplitN(msg.Payload, ",", 2)
		if len(parts) < 2 {
			log.Warnf("Invalid payload: %s", msg.Payload)
			continue
		}
		doctype := parts[0]
		buf := []byte(parts[1])
		if err := json.Unmarshal(buf, &je); err != nil {
			log.Warnf("Error on start: %s", err)
			continue
		}
		if je.Doc != nil {
			je.Doc.Type = doctype
		}
		if je.Old != nil {
			je.Old.Type = doctype
		}
		e := &Event{
			Domain: je.Domain,
			Verb:   je.Verb,
			Doc:    je.Doc,
			OldDoc: je.Old,
		}
		h.mem.Publish(e)
	}
}

func (h *redisHub) GetTopic(domain, doctype string) *topic {
	return nil
}

func (h *redisHub) Publish(e *Event) {
	h.local.broadcast <- e
	buf, err := json.Marshal(e)
	if err != nil {
		log := logger.WithNamespace("realtime-redis")
		log.Warnf("Error on publish: %s", err)
		return
	}
	h.c.Publish(eventsRedisKey, e.Doc.DocType()+","+string(buf))
}

func (h *redisHub) Subscriber(domain string) *DynamicSubscriber {
	return h.mem.Subscriber(domain)
}

func (h *redisHub) SubscribeLocalAll() *DynamicSubscriber {
	ds := newDynamicSubscriber(nil, "")
	ds.addTopic(h.local, "")
	return ds
}
<|endoftext|>