<commit_before>// Copyright 2017 Vector Creations Ltd
//
// 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 routing

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"strings"

	"github.com/matrix-org/dendrite/clientapi/jsonerror"
	"github.com/matrix-org/dendrite/mediaapi/fileutils"
	"github.com/matrix-org/dendrite/mediaapi/storage"
	"github.com/matrix-org/dendrite/mediaapi/thumbnailer"
	"github.com/matrix-org/dendrite/mediaapi/types"
	"github.com/matrix-org/dendrite/setup/config"
	userapi "github.com/matrix-org/dendrite/userapi/api"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
	log "github.com/sirupsen/logrus"
)

// uploadRequest metadata included in or derivable from an upload request
// https://matrix.org/docs/spec/client_server/r0.2.0.html#post-matrix-media-r0-upload
// NOTE: The members come from HTTP request metadata such as headers, query parameters or can be derived from such
type uploadRequest struct {
	MediaMetadata *types.MediaMetadata
	Logger        *log.Entry
}

// uploadResponse defines the format of the JSON response
// https://matrix.org/docs/spec/client_server/r0.2.0.html#post-matrix-media-r0-upload
type uploadResponse struct {
	ContentURI string `json:"content_uri"`
}

// Upload implements POST /upload
// This endpoint involves uploading potentially significant amounts of data to the homeserver.
// This implementation supports a configurable maximum file size limit in bytes. If a user tries to upload more than this, they will receive an error that their upload is too large.
// Uploaded files are processed piece-wise to avoid DoS attacks which would starve the server of memory.
// TODO: We should time out requests if they have not received any data within a configured timeout period.
func Upload(req *http.Request, cfg *config.MediaAPI, dev *userapi.Device, db storage.Database, activeThumbnailGeneration *types.ActiveThumbnailGeneration) util.JSONResponse {
	r, resErr := parseAndValidateRequest(req, cfg, dev)
	if resErr != nil {
		return *resErr
	}

	if resErr = r.doUpload(req.Context(), req.Body, cfg, db, activeThumbnailGeneration); resErr != nil {
		return *resErr
	}

	return util.JSONResponse{
		Code: http.StatusOK,
		JSON: uploadResponse{
			ContentURI: fmt.Sprintf("mxc://%s/%s", cfg.Matrix.ServerName, r.MediaMetadata.MediaID),
		},
	}
}

// parseAndValidateRequest parses the incoming upload request to validate and extract
// all the metadata about the media being uploaded.
// Returns either an uploadRequest or an error formatted as a util.JSONResponse
func parseAndValidateRequest(req *http.Request, cfg *config.MediaAPI, dev *userapi.Device) (*uploadRequest, *util.JSONResponse) {
	r := &uploadRequest{
		MediaMetadata: &types.MediaMetadata{
			Origin:        cfg.Matrix.ServerName,
			FileSizeBytes: types.FileSizeBytes(req.ContentLength),
			ContentType:   types.ContentType(req.Header.Get("Content-Type")),
			UploadName:    types.Filename(url.PathEscape(req.FormValue("filename"))),
			UserID:        types.MatrixUserID(dev.UserID),
		},
		Logger: util.GetLogger(req.Context()).WithField("Origin", cfg.Matrix.ServerName),
	}

	if resErr := r.Validate(*cfg.MaxFileSizeBytes); resErr != nil {
		return nil, resErr
	}

	return r, nil
}

func (r *uploadRequest) generateMediaID(ctx context.Context, db storage.Database) (types.MediaID, error) {
	for {
		// First try generating a meda ID. We'll do this by
		// generating some random bytes and then hex-encoding.
		mediaIDBytes := make([]byte, 32)
		_, err := rand.Read(mediaIDBytes)
		if err != nil {
			return "", fmt.Errorf("rand.Read: %w", err)
		}
		mediaID := types.MediaID(hex.EncodeToString(mediaIDBytes))
		// Then we will check if this media ID already exists in
		// our database. If it does then we had best generate a
		// new one.
		existingMetadata, err := db.GetMediaMetadata(ctx, mediaID, r.MediaMetadata.Origin)
		if err != nil {
			return "", fmt.Errorf("db.GetMediaMetadata: %w", err)
		}
		if existingMetadata != nil {
			// The media ID was already used - repeat the process
			// and generate a new one instead.
			continue
		}
		// The media ID was not already used - let's return that.
		return mediaID, nil
	}
}

func (r *uploadRequest) doUpload(
	ctx context.Context,
	reqReader io.Reader,
	cfg *config.MediaAPI,
	db storage.Database,
	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
) *util.JSONResponse {
	r.Logger.WithFields(log.Fields{
		"UploadName":    r.MediaMetadata.UploadName,
		"FileSizeBytes": r.MediaMetadata.FileSizeBytes,
		"ContentType":   r.MediaMetadata.ContentType,
	}).Info("Uploading file")

	// The file data is hashed and the hash is used as the MediaID. The hash is useful as a
	// method of deduplicating files to save storage, as well as a way to conduct
	// integrity checks on the file data in the repository.
	// Data is truncated to maxFileSizeBytes. Content-Length was reported as 0 < Content-Length <= maxFileSizeBytes so this is OK.
	//
	// TODO: This has a bad API shape where you either need to call:
	//   fileutils.RemoveDir(tmpDir, r.Logger)
	// or call:
	//   r.storeFileAndMetadata(ctx, tmpDir, ...)
	// before you return from doUpload else we will leak a temp file. We could make this nicer with a `WithTransaction` style of
	// nested function to guarantee either storage or cleanup.
	hash, bytesWritten, tmpDir, err := fileutils.WriteTempFile(ctx, reqReader, cfg.AbsBasePath)
	if err != nil {
		r.Logger.WithError(err).WithFields(log.Fields{
			"MaxFileSizeBytes": *cfg.MaxFileSizeBytes,
		}).Warn("Error while transferring file")
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("Failed to upload"),
		}
	}

	// Look up the media by the file hash. If we already have the file but under a
	// different media ID then we won't upload the file again - instead we'll just
	// add a new metadata entry that refers to the same file.
	existingMetadata, err := db.GetMediaMetadataByHash(
		ctx, hash, r.MediaMetadata.Origin,
	)
	if err != nil {
		fileutils.RemoveDir(tmpDir, r.Logger)
		r.Logger.WithError(err).Error("Error querying the database by hash.")
		resErr := jsonerror.InternalServerError()
		return &resErr
	}
	if existingMetadata != nil {
		// The file already exists, delete the uploaded temporary file.
		defer fileutils.RemoveDir(tmpDir, r.Logger)
		// The file already exists. Make a new media ID up for it.
		mediaID, merr := r.generateMediaID(ctx, db)
		if merr != nil {
			r.Logger.WithError(merr).Error("Failed to generate media ID for existing file")
			resErr := jsonerror.InternalServerError()
			return &resErr
		}

		// Then amend the upload metadata.
		r.MediaMetadata = &types.MediaMetadata{
			MediaID:           mediaID,
			Origin:            r.MediaMetadata.Origin,
			ContentType:       r.MediaMetadata.ContentType,
			FileSizeBytes:     r.MediaMetadata.FileSizeBytes,
			CreationTimestamp: r.MediaMetadata.CreationTimestamp,
			UploadName:        r.MediaMetadata.UploadName,
			Base64Hash:        hash,
			UserID:            r.MediaMetadata.UserID,
		}
	} else {
		// The file doesn't exist. Update the request metadata.
		r.MediaMetadata.FileSizeBytes = bytesWritten
		r.MediaMetadata.Base64Hash = hash
		r.MediaMetadata.MediaID, err = r.generateMediaID(ctx, db)
		if err != nil {
			fileutils.RemoveDir(tmpDir, r.Logger)
			r.Logger.WithError(err).Error("Failed to generate media ID for new upload")
			resErr := jsonerror.InternalServerError()
			return &resErr
		}
	}

	r.Logger = r.Logger.WithField("media_id", r.MediaMetadata.MediaID)
	r.Logger.WithFields(log.Fields{
		"Base64Hash":    r.MediaMetadata.Base64Hash,
		"UploadName":    r.MediaMetadata.UploadName,
		"FileSizeBytes": r.MediaMetadata.FileSizeBytes,
		"ContentType":   r.MediaMetadata.ContentType,
	}).Info("File uploaded")

	return r.storeFileAndMetadata(
		ctx, tmpDir, cfg.AbsBasePath, db, cfg.ThumbnailSizes,
		activeThumbnailGeneration, cfg.MaxThumbnailGenerators,
	)
}

// Validate validates the uploadRequest fields
func (r *uploadRequest) Validate(maxFileSizeBytes config.FileSizeBytes) *util.JSONResponse {
	if r.MediaMetadata.FileSizeBytes < 1 {
		return &util.JSONResponse{
			Code: http.StatusLengthRequired,
			JSON: jsonerror.Unknown("HTTP Content-Length request header must be greater than zero."),
		}
	}
	if maxFileSizeBytes > 0 && r.MediaMetadata.FileSizeBytes > types.FileSizeBytes(maxFileSizeBytes) {
		return &util.JSONResponse{
			Code: http.StatusRequestEntityTooLarge,
			JSON: jsonerror.Unknown(fmt.Sprintf("HTTP Content-Length is greater than the maximum allowed upload size (%v).", maxFileSizeBytes)),
		}
	}
	// TODO: Check if the Content-Type is a valid type?
	if r.MediaMetadata.ContentType == "" {
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("HTTP Content-Type request header must be set."),
		}
	}
	if strings.HasPrefix(string(r.MediaMetadata.UploadName), "~") {
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("File name must not begin with '~'."),
		}
	}
	// TODO: Validate filename - what are the valid characters?
	if r.MediaMetadata.UserID != "" {
		// TODO: We should put user ID parsing code into gomatrixserverlib and use that instead
		//       (see https://github.com/matrix-org/gomatrixserverlib/blob/3394e7c7003312043208aa73727d2256eea3d1f6/eventcontent.go#L347 )
		//       It should be a struct (with pointers into a single string to avoid copying) and
		//       we should update all refs to use UserID types rather than strings.
		// https://github.com/matrix-org/synapse/blob/v0.19.2/synapse/types.py#L92
		if _, _, err := gomatrixserverlib.SplitID('@', string(r.MediaMetadata.UserID)); err != nil {
			return &util.JSONResponse{
				Code: http.StatusBadRequest,
				JSON: jsonerror.BadJSON("user id must be in the form @localpart:domain"),
			}
		}
	}
	return nil
}

// storeFileAndMetadata moves the temporary file to its final path based on metadata and stores the metadata in the database
// See getPathFromMediaMetadata in fileutils for details of the final path.
// The order of operations is important as it avoids metadata entering the database before the file
// is ready, and if we fail to move the file, it never gets added to the database.
// Returns a util.JSONResponse error and cleans up directories in case of error.
func (r *uploadRequest) storeFileAndMetadata(
	ctx context.Context,
	tmpDir types.Path,
	absBasePath config.Path,
	db storage.Database,
	thumbnailSizes []config.ThumbnailSize,
	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
	maxThumbnailGenerators int,
) *util.JSONResponse {
	finalPath, duplicate, err := fileutils.MoveFileWithHashCheck(tmpDir, r.MediaMetadata, absBasePath, r.Logger)
	if err != nil {
		r.Logger.WithError(err).Error("Failed to move file.")
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("Failed to upload"),
		}
	}
	if duplicate {
		r.Logger.WithField("dst", finalPath).Info("File was stored previously - discarding duplicate")
	}

	if err = db.StoreMediaMetadata(ctx, r.MediaMetadata); err != nil {
		r.Logger.WithError(err).Warn("Failed to store metadata")
		// If the file is a duplicate (has the same hash as an existing file) then
		// there is valid metadata in the database for that file. As such we only
		// remove the file if it is not a duplicate.
		if !duplicate {
			fileutils.RemoveDir(types.Path(path.Dir(string(finalPath))), r.Logger)
		}
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("Failed to upload"),
		}
	}

	go func() {
		busy, err := thumbnailer.GenerateThumbnails(
			context.Background(), finalPath, thumbnailSizes, r.MediaMetadata,
			activeThumbnailGeneration, maxThumbnailGenerators, db, r.Logger,
		)
		if err != nil {
			r.Logger.WithError(err).Warn("Error generating thumbnails")
		}
		if busy {
			r.Logger.Warn("Maximum number of active thumbnail generators reached. Skipping pre-generation.")
		}
	}()

	return nil
}
<commit_msg>Remove the Content-Length requirement on upload request (#1831)<commit_after>// Copyright 2017 Vector Creations Ltd
//
// 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 routing

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"strings"

	"github.com/matrix-org/dendrite/clientapi/jsonerror"
	"github.com/matrix-org/dendrite/mediaapi/fileutils"
	"github.com/matrix-org/dendrite/mediaapi/storage"
	"github.com/matrix-org/dendrite/mediaapi/thumbnailer"
	"github.com/matrix-org/dendrite/mediaapi/types"
	"github.com/matrix-org/dendrite/setup/config"
	userapi "github.com/matrix-org/dendrite/userapi/api"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
	log "github.com/sirupsen/logrus"
)

// uploadRequest metadata included in or derivable from an upload request
// https://matrix.org/docs/spec/client_server/r0.2.0.html#post-matrix-media-r0-upload
// NOTE: The members come from HTTP request metadata such as headers, query parameters or can be derived from such
type uploadRequest struct {
	MediaMetadata *types.MediaMetadata
	Logger        *log.Entry
}

// uploadResponse defines the format of the JSON response
// https://matrix.org/docs/spec/client_server/r0.2.0.html#post-matrix-media-r0-upload
type uploadResponse struct {
	ContentURI string `json:"content_uri"`
}

// Upload implements POST /upload
// This endpoint involves uploading potentially significant amounts of data to the homeserver.
// This implementation supports a configurable maximum file size limit in bytes. If a user tries to upload more than this, they will receive an error that their upload is too large.
// Uploaded files are processed piece-wise to avoid DoS attacks which would starve the server of memory.
// TODO: We should time out requests if they have not received any data within a configured timeout period.
func Upload(req *http.Request, cfg *config.MediaAPI, dev *userapi.Device, db storage.Database, activeThumbnailGeneration *types.ActiveThumbnailGeneration) util.JSONResponse {
	r, resErr := parseAndValidateRequest(req, cfg, dev)
	if resErr != nil {
		return *resErr
	}

	if resErr = r.doUpload(req.Context(), req.Body, cfg, db, activeThumbnailGeneration); resErr != nil {
		return *resErr
	}

	return util.JSONResponse{
		Code: http.StatusOK,
		JSON: uploadResponse{
			ContentURI: fmt.Sprintf("mxc://%s/%s", cfg.Matrix.ServerName, r.MediaMetadata.MediaID),
		},
	}
}

// parseAndValidateRequest parses the incoming upload request to validate and extract
// all the metadata about the media being uploaded.
// Returns either an uploadRequest or an error formatted as a util.JSONResponse
func parseAndValidateRequest(req *http.Request, cfg *config.MediaAPI, dev *userapi.Device) (*uploadRequest, *util.JSONResponse) {
	r := &uploadRequest{
		MediaMetadata: &types.MediaMetadata{
			Origin:        cfg.Matrix.ServerName,
			FileSizeBytes: types.FileSizeBytes(req.ContentLength),
			ContentType:   types.ContentType(req.Header.Get("Content-Type")),
			UploadName:    types.Filename(url.PathEscape(req.FormValue("filename"))),
			UserID:        types.MatrixUserID(dev.UserID),
		},
		Logger: util.GetLogger(req.Context()).WithField("Origin", cfg.Matrix.ServerName),
	}

	if resErr := r.Validate(*cfg.MaxFileSizeBytes); resErr != nil {
		return nil, resErr
	}

	return r, nil
}

func (r *uploadRequest) generateMediaID(ctx context.Context, db storage.Database) (types.MediaID, error) {
	for {
		// First try generating a meda ID. We'll do this by
		// generating some random bytes and then hex-encoding.
		mediaIDBytes := make([]byte, 32)
		_, err := rand.Read(mediaIDBytes)
		if err != nil {
			return "", fmt.Errorf("rand.Read: %w", err)
		}
		mediaID := types.MediaID(hex.EncodeToString(mediaIDBytes))
		// Then we will check if this media ID already exists in
		// our database. If it does then we had best generate a
		// new one.
		existingMetadata, err := db.GetMediaMetadata(ctx, mediaID, r.MediaMetadata.Origin)
		if err != nil {
			return "", fmt.Errorf("db.GetMediaMetadata: %w", err)
		}
		if existingMetadata != nil {
			// The media ID was already used - repeat the process
			// and generate a new one instead.
			continue
		}
		// The media ID was not already used - let's return that.
		return mediaID, nil
	}
}

func (r *uploadRequest) doUpload(
	ctx context.Context,
	reqReader io.Reader,
	cfg *config.MediaAPI,
	db storage.Database,
	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
) *util.JSONResponse {
	r.Logger.WithFields(log.Fields{
		"UploadName":    r.MediaMetadata.UploadName,
		"FileSizeBytes": r.MediaMetadata.FileSizeBytes,
		"ContentType":   r.MediaMetadata.ContentType,
	}).Info("Uploading file")

	// The file data is hashed and the hash is used as the MediaID. The hash is useful as a
	// method of deduplicating files to save storage, as well as a way to conduct
	// integrity checks on the file data in the repository.
	// Data is truncated to maxFileSizeBytes. Content-Length was reported as 0 < Content-Length <= maxFileSizeBytes so this is OK.
	//
	// TODO: This has a bad API shape where you either need to call:
	//   fileutils.RemoveDir(tmpDir, r.Logger)
	// or call:
	//   r.storeFileAndMetadata(ctx, tmpDir, ...)
	// before you return from doUpload else we will leak a temp file. We could make this nicer with a `WithTransaction` style of
	// nested function to guarantee either storage or cleanup.
	hash, bytesWritten, tmpDir, err := fileutils.WriteTempFile(ctx, reqReader, cfg.AbsBasePath)
	if err != nil {
		r.Logger.WithError(err).WithFields(log.Fields{
			"MaxFileSizeBytes": *cfg.MaxFileSizeBytes,
		}).Warn("Error while transferring file")
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("Failed to upload"),
		}
	}

	// Check if temp file size exceeds max file size configuration
	if bytesWritten > types.FileSizeBytes(*cfg.MaxFileSizeBytes) {
		fileutils.RemoveDir(tmpDir, r.Logger) // delete temp file
		return requestEntityTooLargeJSONResponse(*cfg.MaxFileSizeBytes)
	}

	// Look up the media by the file hash. If we already have the file but under a
	// different media ID then we won't upload the file again - instead we'll just
	// add a new metadata entry that refers to the same file.
	existingMetadata, err := db.GetMediaMetadataByHash(
		ctx, hash, r.MediaMetadata.Origin,
	)
	if err != nil {
		fileutils.RemoveDir(tmpDir, r.Logger)
		r.Logger.WithError(err).Error("Error querying the database by hash.")
		resErr := jsonerror.InternalServerError()
		return &resErr
	}
	if existingMetadata != nil {
		// The file already exists, delete the uploaded temporary file.
		defer fileutils.RemoveDir(tmpDir, r.Logger)
		// The file already exists. Make a new media ID up for it.
		mediaID, merr := r.generateMediaID(ctx, db)
		if merr != nil {
			r.Logger.WithError(merr).Error("Failed to generate media ID for existing file")
			resErr := jsonerror.InternalServerError()
			return &resErr
		}

		// Then amend the upload metadata.
		r.MediaMetadata = &types.MediaMetadata{
			MediaID:           mediaID,
			Origin:            r.MediaMetadata.Origin,
			ContentType:       r.MediaMetadata.ContentType,
			FileSizeBytes:     r.MediaMetadata.FileSizeBytes,
			CreationTimestamp: r.MediaMetadata.CreationTimestamp,
			UploadName:        r.MediaMetadata.UploadName,
			Base64Hash:        hash,
			UserID:            r.MediaMetadata.UserID,
		}
	} else {
		// The file doesn't exist. Update the request metadata.
		r.MediaMetadata.FileSizeBytes = bytesWritten
		r.MediaMetadata.Base64Hash = hash
		r.MediaMetadata.MediaID, err = r.generateMediaID(ctx, db)
		if err != nil {
			fileutils.RemoveDir(tmpDir, r.Logger)
			r.Logger.WithError(err).Error("Failed to generate media ID for new upload")
			resErr := jsonerror.InternalServerError()
			return &resErr
		}
	}

	r.Logger = r.Logger.WithField("media_id", r.MediaMetadata.MediaID)
	r.Logger.WithFields(log.Fields{
		"Base64Hash":    r.MediaMetadata.Base64Hash,
		"UploadName":    r.MediaMetadata.UploadName,
		"FileSizeBytes": r.MediaMetadata.FileSizeBytes,
		"ContentType":   r.MediaMetadata.ContentType,
	}).Info("File uploaded")

	return r.storeFileAndMetadata(
		ctx, tmpDir, cfg.AbsBasePath, db, cfg.ThumbnailSizes,
		activeThumbnailGeneration, cfg.MaxThumbnailGenerators,
	)
}

func requestEntityTooLargeJSONResponse(maxFileSizeBytes config.FileSizeBytes) *util.JSONResponse {
	return &util.JSONResponse{
		Code: http.StatusRequestEntityTooLarge,
		JSON: jsonerror.Unknown(fmt.Sprintf("HTTP Content-Length is greater than the maximum allowed upload size (%v).", maxFileSizeBytes)),
	}
}

// Validate validates the uploadRequest fields
func (r *uploadRequest) Validate(maxFileSizeBytes config.FileSizeBytes) *util.JSONResponse {
	if maxFileSizeBytes > 0 && r.MediaMetadata.FileSizeBytes > types.FileSizeBytes(maxFileSizeBytes) {
		return requestEntityTooLargeJSONResponse(maxFileSizeBytes)
	}
	// TODO: Check if the Content-Type is a valid type?
	if r.MediaMetadata.ContentType == "" {
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("HTTP Content-Type request header must be set."),
		}
	}
	if strings.HasPrefix(string(r.MediaMetadata.UploadName), "~") {
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("File name must not begin with '~'."),
		}
	}
	// TODO: Validate filename - what are the valid characters?
	if r.MediaMetadata.UserID != "" {
		// TODO: We should put user ID parsing code into gomatrixserverlib and use that instead
		//       (see https://github.com/matrix-org/gomatrixserverlib/blob/3394e7c7003312043208aa73727d2256eea3d1f6/eventcontent.go#L347 )
		//       It should be a struct (with pointers into a single string to avoid copying) and
		//       we should update all refs to use UserID types rather than strings.
		// https://github.com/matrix-org/synapse/blob/v0.19.2/synapse/types.py#L92
		if _, _, err := gomatrixserverlib.SplitID('@', string(r.MediaMetadata.UserID)); err != nil {
			return &util.JSONResponse{
				Code: http.StatusBadRequest,
				JSON: jsonerror.BadJSON("user id must be in the form @localpart:domain"),
			}
		}
	}
	return nil
}

// storeFileAndMetadata moves the temporary file to its final path based on metadata and stores the metadata in the database
// See getPathFromMediaMetadata in fileutils for details of the final path.
// The order of operations is important as it avoids metadata entering the database before the file
// is ready, and if we fail to move the file, it never gets added to the database.
// Returns a util.JSONResponse error and cleans up directories in case of error.
func (r *uploadRequest) storeFileAndMetadata(
	ctx context.Context,
	tmpDir types.Path,
	absBasePath config.Path,
	db storage.Database,
	thumbnailSizes []config.ThumbnailSize,
	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
	maxThumbnailGenerators int,
) *util.JSONResponse {
	finalPath, duplicate, err := fileutils.MoveFileWithHashCheck(tmpDir, r.MediaMetadata, absBasePath, r.Logger)
	if err != nil {
		r.Logger.WithError(err).Error("Failed to move file.")
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("Failed to upload"),
		}
	}
	if duplicate {
		r.Logger.WithField("dst", finalPath).Info("File was stored previously - discarding duplicate")
	}

	if err = db.StoreMediaMetadata(ctx, r.MediaMetadata); err != nil {
		r.Logger.WithError(err).Warn("Failed to store metadata")
		// If the file is a duplicate (has the same hash as an existing file) then
		// there is valid metadata in the database for that file. As such we only
		// remove the file if it is not a duplicate.
		if !duplicate {
			fileutils.RemoveDir(types.Path(path.Dir(string(finalPath))), r.Logger)
		}
		return &util.JSONResponse{
			Code: http.StatusBadRequest,
			JSON: jsonerror.Unknown("Failed to upload"),
		}
	}

	go func() {
		busy, err := thumbnailer.GenerateThumbnails(
			context.Background(), finalPath, thumbnailSizes, r.MediaMetadata,
			activeThumbnailGeneration, maxThumbnailGenerators, db, r.Logger,
		)
		if err != nil {
			r.Logger.WithError(err).Warn("Error generating thumbnails")
		}
		if busy {
			r.Logger.Warn("Maximum number of active thumbnail generators reached. Skipping pre-generation.")
		}
	}()

	return nil
}
<|endoftext|>


<commit_before>package bcast

/*
   bcast package for Go. Broadcasting on a set of channels.

   Copyright © 2013 Alexander I.Grafov <grafov@gmail.com>.
   All rights reserved.
   Use of this source code is governed by a BSD-style
   license that can be found in the LICENSE file.
*/

import (
	"container/heap"
	"errors"
	"sync"
	"time"
)

// Message is an internal structure to pack messages together with
// info about sender.
type Message struct {
	sender  *Member
	payload interface{}
	clock   int
}

// Member represents member of a Broadcast group.
type Member struct {
	group        *Group
	Read         chan interface{}
	clock        int
	messageQueue PriorityQueue
	send         chan Message
	close        chan bool
}

// Group provides a mechanism for the broadcast of messages to a
// collection of channels.
type Group struct {
	in         chan Message
	close      chan bool
	members    []*Member
	clock      int
	memberLock sync.Mutex
	clockLock  sync.Mutex
}

// NewGroup creates a new broadcast group.
func NewGroup() *Group {
	in := make(chan Message)
	close := make(chan bool)
	return &Group{in: in, close: close, clock: 0}
}

// MemberCount returns the number of members in the Broadcast Group.
func (g *Group) MemberCount() int {
	return len(g.Members())
}

// Members returns a slice of Members that are currently in the Group.
func (g *Group) Members() []*Member {
	g.memberLock.Lock()
	res := g.members[:]
	g.memberLock.Unlock()
	return res
}

// Join returns a new member object and handles the creation of its
// output channel.
func (g *Group) Join() *Member {
	memberChannel := make(chan interface{})
	return g.Add(memberChannel)
}

// Leave removes the provided member from the group
func (g *Group) Leave(leaving *Member) error {
	g.memberLock.Lock()
	memberIndex := -1
	for index, member := range g.members {
		if member == leaving {
			memberIndex = index
			break
		}
	}
	if memberIndex == -1 {
		g.memberLock.Unlock()
		return errors.New("Could not find provided memeber for removal")
	}
	g.members = append(g.members[:memberIndex], g.members[memberIndex+1:]...)
	leaving.close <- true // TODO: need to handle the case where there
	// is still stuff in this Members priorityQueue
	g.memberLock.Unlock()
	return nil
}

// Add adds a member to the group for the provided interface channel.
func (g *Group) Add(memberChannel chan interface{}) *Member {
	g.memberLock.Lock()
	g.clockLock.Lock()
	member := &Member{
		group:        g,
		Read:         memberChannel,
		clock:        g.clock,
		messageQueue: PriorityQueue{},
		send:         make(chan Message),
		close:        make(chan bool),
	}
	go member.listen()
	g.members = append(g.members, member)
	g.clockLock.Unlock()
	g.memberLock.Unlock()
	return member
}

// Close terminates the group immediately.
func (g *Group) Close() {
	g.close <- true
}

// Broadcast messages received from one group member to others.
// If incoming messages not arrived during `timeout` then function returns.
func (g *Group) Broadcast(timeout time.Duration) {
	var timeoutChannel <-chan time.Time
	if timeout != 0 {
		timeoutChannel = time.After(timeout)
	}
	for {
		select {
		case received := <-g.in:
			g.memberLock.Lock()
			g.clockLock.Lock()
			members := g.members[:]
			received.clock = g.clock
			g.clock++
			g.clockLock.Unlock()
			g.memberLock.Unlock()
			for _, member := range members {
				// This is done in a goroutine because if it
				// weren't it would be a blocking call
				go func(member *Member, received Message) {
					member.send <- received
				}(member, received)
			}
		case <-timeoutChannel:
			if timeout > 0 {
				return
			}
		case <-g.close:
			return
		}
	}
}

// Send broadcasts a message to every one of a Group's members.
func (g *Group) Send(val interface{}) {
	g.in <- Message{sender: nil, payload: val}
}

// Close removes the member it is called on from its broadcast group.
func (m *Member) Close() {
	m.group.Leave(m)
}

// Send broadcasts a message from one Member to the channels of all
// the other members in its group.
func (m *Member) Send(val interface{}) {
	m.group.in <- Message{sender: m, payload: val}
}

// Recv reads one value from the member's Read channel
func (m *Member) Recv() interface{} {
	return <-m.Read
}

func (m *Member) listen() {
	for {
		select {
		case message := <-m.send:
			m.handleMessage(&message)
		case <-m.close:
			return
		}
	}
}

func (m *Member) handleMessage(message *Message) {
	if !m.trySend(message) {
		heap.Push(&m.messageQueue, &Item{
			priority: message.clock,
			value:    message,
		})
		return
	}
	if m.messageQueue.Len() > 0 {
		nextMessage := m.messageQueue[0].value.(*Message)
		for m.trySend(nextMessage) {
			heap.Pop(&m.messageQueue)
			if m.messageQueue.Len() > 0 {
				nextMessage = m.messageQueue[0].value.(*Message)
			} else {
				break
			}
		}
	}
}

func (m *Member) trySend(message *Message) bool {
	shouldSend := message.clock == m.clock
	if shouldSend {
		if message.sender != m {
			m.Read <- message.payload
		}
		m.clock++
	}
	return shouldSend
}
<commit_msg>close the read channel on close<commit_after>package bcast

/*
   bcast package for Go. Broadcasting on a set of channels.

   Copyright © 2013 Alexander I.Grafov <grafov@gmail.com>.
   All rights reserved.
   Use of this source code is governed by a BSD-style
   license that can be found in the LICENSE file.
*/

import (
	"container/heap"
	"errors"
	"sync"
	"time"
)

// Message is an internal structure to pack messages together with
// info about sender.
type Message struct {
	sender  *Member
	payload interface{}
	clock   int
}

// Member represents member of a Broadcast group.
type Member struct {
	group        *Group
	Read         chan interface{}
	clock        int
	messageQueue PriorityQueue
	send         chan Message
	close        chan bool
}

// Group provides a mechanism for the broadcast of messages to a
// collection of channels.
type Group struct {
	in         chan Message
	close      chan bool
	members    []*Member
	clock      int
	memberLock sync.Mutex
	clockLock  sync.Mutex
}

// NewGroup creates a new broadcast group.
func NewGroup() *Group {
	in := make(chan Message)
	close := make(chan bool)
	return &Group{in: in, close: close, clock: 0}
}

// MemberCount returns the number of members in the Broadcast Group.
func (g *Group) MemberCount() int {
	return len(g.Members())
}

// Members returns a slice of Members that are currently in the Group.
func (g *Group) Members() []*Member {
	g.memberLock.Lock()
	res := g.members[:]
	g.memberLock.Unlock()
	return res
}

// Join returns a new member object and handles the creation of its
// output channel.
func (g *Group) Join() *Member {
	memberChannel := make(chan interface{})
	return g.Add(memberChannel)
}

// Leave removes the provided member from the group
func (g *Group) Leave(leaving *Member) error {
	g.memberLock.Lock()
	memberIndex := -1
	for index, member := range g.members {
		if member == leaving {
			memberIndex = index
			break
		}
	}
	if memberIndex == -1 {
		g.memberLock.Unlock()
		return errors.New("Could not find provided memeber for removal")
	}
	g.members = append(g.members[:memberIndex], g.members[memberIndex+1:]...)
	leaving.close <- true // TODO: need to handle the case where there
	// is still stuff in this Members priorityQueue
	g.memberLock.Unlock()
	return nil
}

// Add adds a member to the group for the provided interface channel.
func (g *Group) Add(memberChannel chan interface{}) *Member {
	g.memberLock.Lock()
	g.clockLock.Lock()
	member := &Member{
		group:        g,
		Read:         memberChannel,
		clock:        g.clock,
		messageQueue: PriorityQueue{},
		send:         make(chan Message),
		close:        make(chan bool),
	}
	go member.listen()
	g.members = append(g.members, member)
	g.clockLock.Unlock()
	g.memberLock.Unlock()
	return member
}

// Close terminates the group immediately.
func (g *Group) Close() {
	g.close <- true
}

// Broadcast messages received from one group member to others.
// If incoming messages not arrived during `timeout` then function returns.
func (g *Group) Broadcast(timeout time.Duration) {
	var timeoutChannel <-chan time.Time
	if timeout != 0 {
		timeoutChannel = time.After(timeout)
	}
	for {
		select {
		case received := <-g.in:
			g.memberLock.Lock()
			g.clockLock.Lock()
			members := g.members[:]
			received.clock = g.clock
			g.clock++
			g.clockLock.Unlock()
			g.memberLock.Unlock()
			for _, member := range members {
				// This is done in a goroutine because if it
				// weren't it would be a blocking call
				go func(member *Member, received Message) {
					member.send <- received
				}(member, received)
			}
		case <-timeoutChannel:
			if timeout > 0 {
				return
			}
		case <-g.close:
			return
		}
	}
}

// Send broadcasts a message to every one of a Group's members.
func (g *Group) Send(val interface{}) {
	g.in <- Message{sender: nil, payload: val}
}

// Close removes the member it is called on from its broadcast group.
func (m *Member) Close() {
	m.group.Leave(m)
}

// Send broadcasts a message from one Member to the channels of all
// the other members in its group.
func (m *Member) Send(val interface{}) {
	m.group.in <- Message{sender: m, payload: val}
}

// Recv reads one value from the member's Read channel
func (m *Member) Recv() interface{} {
	return <-m.Read
}

func (m *Member) listen() {
	for {
		select {
		case message := <-m.send:
			m.handleMessage(&message)
		case <-m.close:
			close(m.Read)
			return
		}
	}
}

func (m *Member) handleMessage(message *Message) {
	if !m.trySend(message) {
		heap.Push(&m.messageQueue, &Item{
			priority: message.clock,
			value:    message,
		})
		return
	}
	if m.messageQueue.Len() > 0 {
		nextMessage := m.messageQueue[0].value.(*Message)
		for m.trySend(nextMessage) {
			heap.Pop(&m.messageQueue)
			if m.messageQueue.Len() > 0 {
				nextMessage = m.messageQueue[0].value.(*Message)
			} else {
				break
			}
		}
	}
}

func (m *Member) trySend(message *Message) bool {
	shouldSend := message.clock == m.clock
	if shouldSend {
		if message.sender != m {
			m.Read <- message.payload
		}
		m.clock++
	}
	return shouldSend
}
<|endoftext|>


<commit_before>// Copyright 2016 Openprovider Authors. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

/*
Package rates 0.1.3
This package helps to manage exchange rates from any provider

Example 1: Get all exchange rates for the ECB Provider

    package main

    import (
        "fmt"

        "github.com/openprovider/rates"
        "github.com/openprovider/rates/providers"
    )

    func main() {
        registry := rates.Registry{
            // any collection of providers which implement rates.Provider interface
            providers.NewECBProvider(new(rates.Options)),
        }
        service := rates.New(registry)
        rates, errors := service.FetchLast()
        if len(errors) != 0 {
            fmt.Println(errors)
        }
        fmt.Println("European Central Bank exchange rates for today")
        for index, rate := range rates {
            fmt.Printf("%d. %s - %v\r\n", index+1, rate.Currency, rate.Value)
        }
    }

Example 2: Get exchange rates for EUR, USD, CHF, HKD

    package main

    import (
        "fmt"

        "github.com/openprovider/rates"
        "github.com/openprovider/rates/providers"
    )

    func main() {
        registry := rates.Registry{
            // any collection of providers which implement rates.Provider interface
            providers.NewECBProvider(
                &rates.Options{
                    Currencies: []string{
                        providers.EUR,
                        providers.USD,
                        providers.CHF,
                        providers.HKD,
                    },
                },
            ),
        }
        service := rates.New(registry)
        rates, errors := service.FetchLast()
        if len(errors) != 0 {
            fmt.Println(errors)
        }
        fmt.Println("European Central Bank exchange rates for today")
        for index, rate := range rates {
            fmt.Printf("%d. %s - %v\r\n", index+1, rate.Currency, rate.Value)
        }
    }

Exchange Rates Provider
*/
package rates

import (
	"time"

	"golang.org/x/text/currency"
)

// Rate represent date and currency exchange rates
type Rate struct {
	ID       uint64        `json:"id,omitempty"`
	Date     string        `json:"date"`
	Currency string        `json:"currency"`
	Time     time.Time     `json:"-"`
	Base     currency.Unit `json:"-"`
	Unit     currency.Unit `json:"-"`
	Value    interface{}   `json:"value"`
}

// Options is some specific things for the specific provider
// It should configure the provider to manage currencies
type Options struct {
	// API key/token
	Token string
	// List of the currencies which need to get from the provider
	// If it is empty, should get all of existing currencies from the provider
	Currencies []string
	// Flexible settings list
	Settings map[string]interface{}
}

// Provider holds methods for providers which implement this interface
type Provider interface {
	FetchLast() (rates []Rate, errors []error)
	FetchHistory() (rates []Rate, errors []error)
}

// Registry contains registered providers
type Registry []Provider

// New service which contains registered providers
func New(providers ...Provider) Provider {
	var registry Registry
	for _, provider := range providers {
		registry = append(registry, provider)
	}
	return registry
}

// FetchLast returns exchange rates from all registered providers on last day
func (registry Registry) FetchLast() (rates []Rate, errors []error) {
	for _, provider := range registry {
		r, errs := provider.FetchLast()
		rates = append(rates, r...)
		errors = append(errors, errs...)
	}
	return
}

// FetchHistory returns exchange rates from all registered providers from history
func (registry Registry) FetchHistory() (rates []Rate, errors []error) {
	for _, provider := range registry {
		r, errs := provider.FetchHistory()
		rates = append(rates, r...)
		errors = append(errors, errs...)
	}
	return
}
<commit_msg>Bumped version number to 0.2.0<commit_after>// Copyright 2016 Openprovider Authors. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

/*
Package rates 0.2.0
This package helps to manage exchange rates from any provider

Example 1: Get all exchange rates for the ECB Provider

    package main

    import (
        "fmt"

        "github.com/openprovider/rates"
        "github.com/openprovider/rates/providers"
    )

    func main() {
        registry := rates.Registry{
            // any collection of providers which implement rates.Provider interface
            providers.NewECBProvider(new(rates.Options)),
        }
        service := rates.New(registry)
        rates, errors := service.FetchLast()
        if len(errors) != 0 {
            fmt.Println(errors)
        }
        fmt.Println("European Central Bank exchange rates for today")
        for index, rate := range rates {
            fmt.Printf("%d. %s - %v\r\n", index+1, rate.Currency, rate.Value)
        }
    }

Example 2: Get exchange rates for EUR, USD, CHF, HKD

    package main

    import (
        "fmt"

        "github.com/openprovider/rates"
        "github.com/openprovider/rates/providers"
    )

    func main() {
        registry := rates.Registry{
            // any collection of providers which implement rates.Provider interface
            providers.NewECBProvider(
                &rates.Options{
                    Currencies: []string{
                        providers.EUR,
                        providers.USD,
                        providers.CHF,
                        providers.HKD,
                    },
                },
            ),
        }
        service := rates.New(registry)
        rates, errors := service.FetchLast()
        if len(errors) != 0 {
            fmt.Println(errors)
        }
        fmt.Println("European Central Bank exchange rates for today")
        for index, rate := range rates {
            fmt.Printf("%d. %s - %v\r\n", index+1, rate.Currency, rate.Value)
        }
    }

Exchange Rates Provider
*/
package rates

import (
	"time"

	"golang.org/x/text/currency"
)

// Rate represent date and currency exchange rates
type Rate struct {
	ID       uint64        `json:"id,omitempty"`
	Date     string        `json:"date"`
	Currency string        `json:"currency"`
	Time     time.Time     `json:"-"`
	Base     currency.Unit `json:"-"`
	Unit     currency.Unit `json:"-"`
	Value    interface{}   `json:"value"`
}

// Options is some specific things for the specific provider
// It should configure the provider to manage currencies
type Options struct {
	// API key/token
	Token string
	// List of the currencies which need to get from the provider
	// If it is empty, should get all of existing currencies from the provider
	Currencies []string
	// Flexible settings list
	Settings map[string]interface{}
}

// Provider holds methods for providers which implement this interface
type Provider interface {
	FetchLast() (rates []Rate, errors []error)
	FetchHistory() (rates []Rate, errors []error)
}

// Registry contains registered providers
type Registry []Provider

// New service which contains registered providers
func New(providers ...Provider) Provider {
	var registry Registry
	for _, provider := range providers {
		registry = append(registry, provider)
	}
	return registry
}

// FetchLast returns exchange rates from all registered providers on last day
func (registry Registry) FetchLast() (rates []Rate, errors []error) {
	for _, provider := range registry {
		r, errs := provider.FetchLast()
		rates = append(rates, r...)
		errors = append(errors, errs...)
	}
	return
}

// FetchHistory returns exchange rates from all registered providers from history
func (registry Registry) FetchHistory() (rates []Rate, errors []error) {
	for _, provider := range registry {
		r, errs := provider.FetchHistory()
		rates = append(rates, r...)
		errors = append(errors, errs...)
	}
	return
}
<|endoftext|>


<commit_before>package csvutil

import (
	"encoding/csv"
	"os"
)

func NewReader(path string, comma rune) (*csv.Reader, *os.File, error) {
	var myCsv *csv.Reader
	var file *os.File
	file, err := os.Open(path)
	if err != nil {
		return myCsv, file, err
	}
	reader := csv.NewReader(file)
	reader.Comma = comma
	return reader, file, nil
}
<commit_msg>add strip BOM support to csvutil<commit_after>package csvutil

import (
	"encoding/csv"
	"os"
)

/*

For UTF-8 BOM, csv.Reader.Read() will return error = "line 1, column 1: bare \" in non-quoted-field"

If you encounter this close the file and call again with stripBom = true

*/

func NewReader(path string, comma rune, stripBom bool) (*csv.Reader, *os.File, error) {
	var myCsv *csv.Reader
	var file *os.File
	file, err := os.Open(path)
	if err != nil {
		return myCsv, file, err
	}
	if stripBom {
		b3 := make([]byte, 3)
		_, err := file.Read(b3)
		if err != nil {
			return myCsv, file, err
		}
	}
	reader := csv.NewReader(file)
	reader.Comma = comma
	return reader, file, nil
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package rpcdb

import (
	"sync/atomic"

	"golang.org/x/net/context"

	"github.com/ava-labs/avalanchego/database"
	"github.com/ava-labs/avalanchego/database/nodb"
	"github.com/ava-labs/avalanchego/database/rpcdb/rpcdbproto"
	"github.com/ava-labs/avalanchego/utils"
	"github.com/ava-labs/avalanchego/utils/units"
	"github.com/ava-labs/avalanchego/utils/wrappers"
)

const (
	maxBatchSize = 128 * units.KiB // 64 KiB

	// baseElementSize is an approximation of the protobuf encoding overhead per
	// element
	baseElementSize = 8 // bytes
)

var (
	_ database.Database = &DatabaseClient{}
	_ database.Batch    = &batch{}
	_ database.Iterator = &iterator{}
)

// DatabaseClient is an implementation of database that talks over RPC.
type DatabaseClient struct {
	client rpcdbproto.DatabaseClient

	batchIndex int64
}

// NewClient returns a database instance connected to a remote database instance
func NewClient(client rpcdbproto.DatabaseClient) *DatabaseClient {
	return &DatabaseClient{client: client}
}

// Has attempts to return if the database has a key with the provided value.
func (db *DatabaseClient) Has(key []byte) (bool, error) {
	resp, err := db.client.Has(context.Background(), &rpcdbproto.HasRequest{
		Key: key,
	})
	if err != nil {
		return false, err
	}
	return resp.Has, errCodeToError[resp.Err]
}

// Get attempts to return the value that was mapped to the key that was provided
func (db *DatabaseClient) Get(key []byte) ([]byte, error) {
	resp, err := db.client.Get(context.Background(), &rpcdbproto.GetRequest{
		Key: key,
	})
	if err != nil {
		return nil, err
	}
	return resp.Value, errCodeToError[resp.Err]
}

// Put attempts to set the value this key maps to
func (db *DatabaseClient) Put(key, value []byte) error {
	resp, err := db.client.Put(context.Background(), &rpcdbproto.PutRequest{
		Key:   key,
		Value: value,
	})
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

// Delete attempts to remove any mapping from the key
func (db *DatabaseClient) Delete(key []byte) error {
	resp, err := db.client.Delete(context.Background(), &rpcdbproto.DeleteRequest{
		Key: key,
	})
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

// NewBatch returns a new batch
func (db *DatabaseClient) NewBatch() database.Batch { return &batch{db: db} }

// NewIterator implements the Database interface
func (db *DatabaseClient) NewIterator() database.Iterator {
	return db.NewIteratorWithStartAndPrefix(nil, nil)
}

// NewIteratorWithStart implements the Database interface
func (db *DatabaseClient) NewIteratorWithStart(start []byte) database.Iterator {
	return db.NewIteratorWithStartAndPrefix(start, nil)
}

// NewIteratorWithPrefix implements the Database interface
func (db *DatabaseClient) NewIteratorWithPrefix(prefix []byte) database.Iterator {
	return db.NewIteratorWithStartAndPrefix(nil, prefix)
}

// NewIteratorWithStartAndPrefix returns a new empty iterator
func (db *DatabaseClient) NewIteratorWithStartAndPrefix(start, prefix []byte) database.Iterator {
	resp, err := db.client.NewIteratorWithStartAndPrefix(context.Background(), &rpcdbproto.NewIteratorWithStartAndPrefixRequest{
		Start:  start,
		Prefix: prefix,
	})
	if err != nil {
		return &nodb.Iterator{Err: err}
	}
	return &iterator{
		db: db,
		id: resp.Id,
	}
}

// Stat attempts to return the statistic of this database
func (db *DatabaseClient) Stat(property string) (string, error) {
	resp, err := db.client.Stat(context.Background(), &rpcdbproto.StatRequest{
		Property: property,
	})
	if err != nil {
		return "", err
	}
	return resp.Stat, errCodeToError[resp.Err]
}

// Compact attempts to optimize the space utilization in the provided range
func (db *DatabaseClient) Compact(start, limit []byte) error {
	resp, err := db.client.Compact(context.Background(), &rpcdbproto.CompactRequest{
		Start: start,
		Limit: limit,
	})
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

// Close attempts to close the database
func (db *DatabaseClient) Close() error {
	resp, err := db.client.Close(context.Background(), &rpcdbproto.CloseRequest{})
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

type keyValue struct {
	key    []byte
	value  []byte
	delete bool
}

type batch struct {
	db     *DatabaseClient
	writes []keyValue
	size   int
}

func (b *batch) Put(key, value []byte) error {
	b.writes = append(b.writes, keyValue{utils.CopyBytes(key), utils.CopyBytes(value), false})
	b.size += len(key) + len(value)
	return nil
}

func (b *batch) Delete(key []byte) error {
	b.writes = append(b.writes, keyValue{utils.CopyBytes(key), nil, true})
	b.size += len(key)
	return nil
}

func (b *batch) Size() int { return b.size }

func (b *batch) Write() error {
	request := &rpcdbproto.WriteBatchRequest{
		Id:        atomic.AddInt64(&b.db.batchIndex, 1),
		Continues: true,
	}
	currentSize := 0
	keySet := make(map[string]struct{}, len(b.writes))
	for i := len(b.writes) - 1; i >= 0; i-- {
		kv := b.writes[i]
		key := string(kv.key)
		if _, overwritten := keySet[key]; overwritten {
			continue
		}
		keySet[key] = struct{}{}

		sizeChange := baseElementSize + len(kv.key) + len(kv.value)
		if newSize := currentSize + sizeChange; newSize > maxBatchSize {
			resp, err := b.db.client.WriteBatch(context.Background(), request)
			if err != nil {
				return err
			}
			if err := errCodeToError[resp.Err]; err != nil {
				return err
			}
			currentSize = 0
			request.Deletes = request.Deletes[:0]
			request.Puts = request.Puts[:0]
		}
		currentSize += sizeChange

		if kv.delete {
			request.Deletes = append(request.Deletes, &rpcdbproto.DeleteRequest{
				Key: kv.key,
			})
		} else {
			request.Puts = append(request.Puts, &rpcdbproto.PutRequest{
				Key:   kv.key,
				Value: kv.value,
			})
		}
	}

	request.Continues = false
	resp, err := b.db.client.WriteBatch(context.Background(), request)
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

func (b *batch) Reset() {
	if cap(b.writes) > len(b.writes)*database.MaxExcessCapacityFactor {
		b.writes = make([]keyValue, 0, cap(b.writes)/database.CapacityReductionFactor)
	} else {
		b.writes = b.writes[:0]
	}
	b.size = 0
}

func (b *batch) Replay(w database.KeyValueWriter) error {
	for _, keyvalue := range b.writes {
		if keyvalue.delete {
			if err := w.Delete(keyvalue.key); err != nil {
				return err
			}
		} else if err := w.Put(keyvalue.key, keyvalue.value); err != nil {
			return err
		}
	}
	return nil
}

func (b *batch) Inner() database.Batch { return b }

type iterator struct {
	db *DatabaseClient
	id uint64

	data []*rpcdbproto.PutRequest
	errs wrappers.Errs
}

// Next attempts to move the iterator to the next element and returns if this
// succeeded
func (it *iterator) Next() bool {
	if len(it.data) > 1 {
		it.data = it.data[1:]
		return true
	}

	resp, err := it.db.client.IteratorNext(context.Background(), &rpcdbproto.IteratorNextRequest{
		Id: it.id,
	})
	if err != nil {
		it.errs.Add(err)
		return false
	}
	it.data = resp.Data
	return len(it.data) > 0
}

// Error returns any that occurred while iterating
func (it *iterator) Error() error {
	if it.errs.Errored() {
		return it.errs.Err
	}

	resp, err := it.db.client.IteratorError(context.Background(), &rpcdbproto.IteratorErrorRequest{
		Id: it.id,
	})
	if err != nil {
		it.errs.Add(err)
	} else {
		it.errs.Add(errCodeToError[resp.Err])
	}
	return it.errs.Err
}

// Key returns the key of the current element
func (it *iterator) Key() []byte {
	if len(it.data) == 0 {
		return nil
	}
	return it.data[0].Key
}

// Value returns the value of the current element
func (it *iterator) Value() []byte {
	if len(it.data) == 0 {
		return nil
	}
	return it.data[0].Value
}

// Release frees any resources held by the iterator
func (it *iterator) Release() {
	resp, err := it.db.client.IteratorRelease(context.Background(), &rpcdbproto.IteratorReleaseRequest{
		Id: it.id,
	})
	if err != nil {
		it.errs.Add(err)
	} else {
		it.errs.Add(errCodeToError[resp.Err])
	}
}
<commit_msg>removed old comment<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package rpcdb

import (
	"sync/atomic"

	"golang.org/x/net/context"

	"github.com/ava-labs/avalanchego/database"
	"github.com/ava-labs/avalanchego/database/nodb"
	"github.com/ava-labs/avalanchego/database/rpcdb/rpcdbproto"
	"github.com/ava-labs/avalanchego/utils"
	"github.com/ava-labs/avalanchego/utils/units"
	"github.com/ava-labs/avalanchego/utils/wrappers"
)

const (
	maxBatchSize = 128 * units.KiB

	// baseElementSize is an approximation of the protobuf encoding overhead per
	// element
	baseElementSize = 8 // bytes
)

var (
	_ database.Database = &DatabaseClient{}
	_ database.Batch    = &batch{}
	_ database.Iterator = &iterator{}
)

// DatabaseClient is an implementation of database that talks over RPC.
type DatabaseClient struct {
	client rpcdbproto.DatabaseClient

	batchIndex int64
}

// NewClient returns a database instance connected to a remote database instance
func NewClient(client rpcdbproto.DatabaseClient) *DatabaseClient {
	return &DatabaseClient{client: client}
}

// Has attempts to return if the database has a key with the provided value.
func (db *DatabaseClient) Has(key []byte) (bool, error) {
	resp, err := db.client.Has(context.Background(), &rpcdbproto.HasRequest{
		Key: key,
	})
	if err != nil {
		return false, err
	}
	return resp.Has, errCodeToError[resp.Err]
}

// Get attempts to return the value that was mapped to the key that was provided
func (db *DatabaseClient) Get(key []byte) ([]byte, error) {
	resp, err := db.client.Get(context.Background(), &rpcdbproto.GetRequest{
		Key: key,
	})
	if err != nil {
		return nil, err
	}
	return resp.Value, errCodeToError[resp.Err]
}

// Put attempts to set the value this key maps to
func (db *DatabaseClient) Put(key, value []byte) error {
	resp, err := db.client.Put(context.Background(), &rpcdbproto.PutRequest{
		Key:   key,
		Value: value,
	})
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

// Delete attempts to remove any mapping from the key
func (db *DatabaseClient) Delete(key []byte) error {
	resp, err := db.client.Delete(context.Background(), &rpcdbproto.DeleteRequest{
		Key: key,
	})
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

// NewBatch returns a new batch
func (db *DatabaseClient) NewBatch() database.Batch { return &batch{db: db} }

// NewIterator implements the Database interface
func (db *DatabaseClient) NewIterator() database.Iterator {
	return db.NewIteratorWithStartAndPrefix(nil, nil)
}

// NewIteratorWithStart implements the Database interface
func (db *DatabaseClient) NewIteratorWithStart(start []byte) database.Iterator {
	return db.NewIteratorWithStartAndPrefix(start, nil)
}

// NewIteratorWithPrefix implements the Database interface
func (db *DatabaseClient) NewIteratorWithPrefix(prefix []byte) database.Iterator {
	return db.NewIteratorWithStartAndPrefix(nil, prefix)
}

// NewIteratorWithStartAndPrefix returns a new empty iterator
func (db *DatabaseClient) NewIteratorWithStartAndPrefix(start, prefix []byte) database.Iterator {
	resp, err := db.client.NewIteratorWithStartAndPrefix(context.Background(), &rpcdbproto.NewIteratorWithStartAndPrefixRequest{
		Start:  start,
		Prefix: prefix,
	})
	if err != nil {
		return &nodb.Iterator{Err: err}
	}
	return &iterator{
		db: db,
		id: resp.Id,
	}
}

// Stat attempts to return the statistic of this database
func (db *DatabaseClient) Stat(property string) (string, error) {
	resp, err := db.client.Stat(context.Background(), &rpcdbproto.StatRequest{
		Property: property,
	})
	if err != nil {
		return "", err
	}
	return resp.Stat, errCodeToError[resp.Err]
}

// Compact attempts to optimize the space utilization in the provided range
func (db *DatabaseClient) Compact(start, limit []byte) error {
	resp, err := db.client.Compact(context.Background(), &rpcdbproto.CompactRequest{
		Start: start,
		Limit: limit,
	})
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

// Close attempts to close the database
func (db *DatabaseClient) Close() error {
	resp, err := db.client.Close(context.Background(), &rpcdbproto.CloseRequest{})
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

type keyValue struct {
	key    []byte
	value  []byte
	delete bool
}

type batch struct {
	db     *DatabaseClient
	writes []keyValue
	size   int
}

func (b *batch) Put(key, value []byte) error {
	b.writes = append(b.writes, keyValue{utils.CopyBytes(key), utils.CopyBytes(value), false})
	b.size += len(key) + len(value)
	return nil
}

func (b *batch) Delete(key []byte) error {
	b.writes = append(b.writes, keyValue{utils.CopyBytes(key), nil, true})
	b.size += len(key)
	return nil
}

func (b *batch) Size() int { return b.size }

func (b *batch) Write() error {
	request := &rpcdbproto.WriteBatchRequest{
		Id:        atomic.AddInt64(&b.db.batchIndex, 1),
		Continues: true,
	}
	currentSize := 0
	keySet := make(map[string]struct{}, len(b.writes))
	for i := len(b.writes) - 1; i >= 0; i-- {
		kv := b.writes[i]
		key := string(kv.key)
		if _, overwritten := keySet[key]; overwritten {
			continue
		}
		keySet[key] = struct{}{}

		sizeChange := baseElementSize + len(kv.key) + len(kv.value)
		if newSize := currentSize + sizeChange; newSize > maxBatchSize {
			resp, err := b.db.client.WriteBatch(context.Background(), request)
			if err != nil {
				return err
			}
			if err := errCodeToError[resp.Err]; err != nil {
				return err
			}
			currentSize = 0
			request.Deletes = request.Deletes[:0]
			request.Puts = request.Puts[:0]
		}
		currentSize += sizeChange

		if kv.delete {
			request.Deletes = append(request.Deletes, &rpcdbproto.DeleteRequest{
				Key: kv.key,
			})
		} else {
			request.Puts = append(request.Puts, &rpcdbproto.PutRequest{
				Key:   kv.key,
				Value: kv.value,
			})
		}
	}

	request.Continues = false
	resp, err := b.db.client.WriteBatch(context.Background(), request)
	if err != nil {
		return err
	}
	return errCodeToError[resp.Err]
}

func (b *batch) Reset() {
	if cap(b.writes) > len(b.writes)*database.MaxExcessCapacityFactor {
		b.writes = make([]keyValue, 0, cap(b.writes)/database.CapacityReductionFactor)
	} else {
		b.writes = b.writes[:0]
	}
	b.size = 0
}

func (b *batch) Replay(w database.KeyValueWriter) error {
	for _, keyvalue := range b.writes {
		if keyvalue.delete {
			if err := w.Delete(keyvalue.key); err != nil {
				return err
			}
		} else if err := w.Put(keyvalue.key, keyvalue.value); err != nil {
			return err
		}
	}
	return nil
}

func (b *batch) Inner() database.Batch { return b }

type iterator struct {
	db *DatabaseClient
	id uint64

	data []*rpcdbproto.PutRequest
	errs wrappers.Errs
}

// Next attempts to move the iterator to the next element and returns if this
// succeeded
func (it *iterator) Next() bool {
	if len(it.data) > 1 {
		it.data = it.data[1:]
		return true
	}

	resp, err := it.db.client.IteratorNext(context.Background(), &rpcdbproto.IteratorNextRequest{
		Id: it.id,
	})
	if err != nil {
		it.errs.Add(err)
		return false
	}
	it.data = resp.Data
	return len(it.data) > 0
}

// Error returns any that occurred while iterating
func (it *iterator) Error() error {
	if it.errs.Errored() {
		return it.errs.Err
	}

	resp, err := it.db.client.IteratorError(context.Background(), &rpcdbproto.IteratorErrorRequest{
		Id: it.id,
	})
	if err != nil {
		it.errs.Add(err)
	} else {
		it.errs.Add(errCodeToError[resp.Err])
	}
	return it.errs.Err
}

// Key returns the key of the current element
func (it *iterator) Key() []byte {
	if len(it.data) == 0 {
		return nil
	}
	return it.data[0].Key
}

// Value returns the value of the current element
func (it *iterator) Value() []byte {
	if len(it.data) == 0 {
		return nil
	}
	return it.data[0].Value
}

// Release frees any resources held by the iterator
func (it *iterator) Release() {
	resp, err := it.db.client.IteratorRelease(context.Background(), &rpcdbproto.IteratorReleaseRequest{
		Id: it.id,
	})
	if err != nil {
		it.errs.Add(err)
	} else {
		it.errs.Add(errCodeToError[resp.Err])
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cmd

import (
	"fmt"
	"strconv"
	"strings"

	charmresource "gopkg.in/juju/charm.v6-unstable/resource"

	"github.com/juju/errors"
	"github.com/juju/juju/resource"
	"github.com/juju/names"
)

type charmResourcesFormatter struct {
	resources []charmresource.Resource
}

func newCharmResourcesFormatter(resources []charmresource.Resource) *charmResourcesFormatter {
	// Note that unlike the "juju status" code, we don't worry
	// about "compatVersion".
	crf := charmResourcesFormatter{
		resources: resources,
	}
	return &crf
}

func (crf *charmResourcesFormatter) format() []FormattedCharmResource {
	if crf.resources == nil {
		return nil
	}

	var formatted []FormattedCharmResource
	for _, res := range crf.resources {
		formatted = append(formatted, FormatCharmResource(res))
	}
	return formatted
}

// FormatCharmResource converts the resource info into a FormattedCharmResource.
func FormatCharmResource(res charmresource.Resource) FormattedCharmResource {
	return FormattedCharmResource{
		Name:        res.Name,
		Type:        res.Type.String(),
		Path:        res.Path,
		Comment:       res.Comment,
		Revision:    res.Revision,
		Origin:      res.Origin.String(),
		Fingerprint: res.Fingerprint.String(), // ...the hex string.
		Size:        res.Size,
	}
}

// FormatSvcResource converts the resource info into a FormattedServiceResource.
func FormatSvcResource(res resource.Resource) FormattedSvcResource {
	used := !res.IsPlaceholder()
	return FormattedSvcResource{
		ID:               res.ID,
		ServiceID:        res.ServiceID,
		Name:             res.Name,
		Type:             res.Type.String(),
		Path:             res.Path,
		Description:      res.Description,
		Revision:         res.Revision,
		Origin:           res.Origin.String(),
		Fingerprint:      res.Fingerprint.String(),
		Size:             res.Size,
		Used:             used,
		Timestamp:        res.Timestamp,
		Username:         res.Username,
		combinedRevision: combinedRevision(res),
		combinedOrigin:   combinedOrigin(used, res),
		usedYesNo:        usedYesNo(used),
	}
}

// FormatDetailResource converts the arguments into a FormattedServiceResource.
func FormatDetailResource(tag names.UnitTag, svc, unit resource.Resource) (FormattedDetailResource, error) {
	// note that the unit resource can be a zero value here, to indicate that
	// the unit has not downloaded that resource yet.

	unitNum, err := unitNum(tag)
	if err != nil {
		return FormattedDetailResource{}, errors.Trace(err)
	}
	return FormattedDetailResource{
		UnitID:     tag.Id(),
		unitNumber: unitNum,
		Unit:       FormatSvcResource(unit),
		Expected:   FormatSvcResource(svc),
	}, nil
}

func combinedRevision(r resource.Resource) string {
	switch r.Origin {
	case charmresource.OriginStore:
		return fmt.Sprintf("%d", r.Revision)
	case charmresource.OriginUpload:
		if !r.Timestamp.IsZero() {
			return r.Timestamp.Format("2006-02-01T15:04")
		}
	}
	return "-"
}

func combinedOrigin(used bool, r resource.Resource) string {
	if r.Origin == charmresource.OriginUpload && used && r.Username != "" {
		return r.Username
	}
	if r.Origin == charmresource.OriginStore {
		return "charmstore"
	}
	return r.Origin.String()
}

func usedYesNo(used bool) string {
	if used {
		return "yes"
	}
	return "no"
}

func unitNum(unit names.UnitTag) (int, error) {
	vals := strings.SplitN(unit.Id(), "/", 2)
	if len(vals) != 2 {
		return 0, errors.Errorf("%q is not a valid unit ID", unit.Id())
	}
	num, err := strconv.Atoi(vals[1])
	if err != nil {
		return 0, errors.Annotatef(err, "%q is not a valid unit ID", unit.Id())
	}
	return num, nil
}
<commit_msg>fix merge problem<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cmd

import (
	"fmt"
	"strconv"
	"strings"

	charmresource "gopkg.in/juju/charm.v6-unstable/resource"

	"github.com/juju/errors"
	"github.com/juju/juju/resource"
	"github.com/juju/names"
)

type charmResourcesFormatter struct {
	resources []charmresource.Resource
}

func newCharmResourcesFormatter(resources []charmresource.Resource) *charmResourcesFormatter {
	// Note that unlike the "juju status" code, we don't worry
	// about "compatVersion".
	crf := charmResourcesFormatter{
		resources: resources,
	}
	return &crf
}

func (crf *charmResourcesFormatter) format() []FormattedCharmResource {
	if crf.resources == nil {
		return nil
	}

	var formatted []FormattedCharmResource
	for _, res := range crf.resources {
		formatted = append(formatted, FormatCharmResource(res))
	}
	return formatted
}

// FormatCharmResource converts the resource info into a FormattedCharmResource.
func FormatCharmResource(res charmresource.Resource) FormattedCharmResource {
	return FormattedCharmResource{
		Name:        res.Name,
		Type:        res.Type.String(),
		Path:        res.Path,
		Description: res.Description,
		Revision:    res.Revision,
		Origin:      res.Origin.String(),
		Fingerprint: res.Fingerprint.String(), // ...the hex string.
		Size:        res.Size,
	}
}

// FormatSvcResource converts the resource info into a FormattedServiceResource.
func FormatSvcResource(res resource.Resource) FormattedSvcResource {
	used := !res.IsPlaceholder()
	return FormattedSvcResource{
		ID:               res.ID,
		ServiceID:        res.ServiceID,
		Name:             res.Name,
		Type:             res.Type.String(),
		Path:             res.Path,
		Description:      res.Description,
		Revision:         res.Revision,
		Origin:           res.Origin.String(),
		Fingerprint:      res.Fingerprint.String(),
		Size:             res.Size,
		Used:             used,
		Timestamp:        res.Timestamp,
		Username:         res.Username,
		combinedRevision: combinedRevision(res),
		combinedOrigin:   combinedOrigin(used, res),
		usedYesNo:        usedYesNo(used),
	}
}

// FormatDetailResource converts the arguments into a FormattedServiceResource.
func FormatDetailResource(tag names.UnitTag, svc, unit resource.Resource) (FormattedDetailResource, error) {
	// note that the unit resource can be a zero value here, to indicate that
	// the unit has not downloaded that resource yet.

	unitNum, err := unitNum(tag)
	if err != nil {
		return FormattedDetailResource{}, errors.Trace(err)
	}
	return FormattedDetailResource{
		UnitID:     tag.Id(),
		unitNumber: unitNum,
		Unit:       FormatSvcResource(unit),
		Expected:   FormatSvcResource(svc),
	}, nil
}

func combinedRevision(r resource.Resource) string {
	switch r.Origin {
	case charmresource.OriginStore:
		return fmt.Sprintf("%d", r.Revision)
	case charmresource.OriginUpload:
		if !r.Timestamp.IsZero() {
			return r.Timestamp.Format("2006-02-01T15:04")
		}
	}
	return "-"
}

func combinedOrigin(used bool, r resource.Resource) string {
	if r.Origin == charmresource.OriginUpload && used && r.Username != "" {
		return r.Username
	}
	if r.Origin == charmresource.OriginStore {
		return "charmstore"
	}
	return r.Origin.String()
}

func usedYesNo(used bool) string {
	if used {
		return "yes"
	}
	return "no"
}

func unitNum(unit names.UnitTag) (int, error) {
	vals := strings.SplitN(unit.Id(), "/", 2)
	if len(vals) != 2 {
		return 0, errors.Errorf("%q is not a valid unit ID", unit.Id())
	}
	num, err := strconv.Atoi(vals[1])
	if err != nil {
		return 0, errors.Annotatef(err, "%q is not a valid unit ID", unit.Id())
	}
	return num, nil
}
<|endoftext|>


<commit_before>package datadog

import (
	"encoding/json"
	"fmt"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/zorkian/go-datadog-api"
	"strconv"
	"strings"
)

func thresholdSchema() *schema.Schema {
	return &schema.Schema{
		Type:     schema.TypeMap,
		Required: true,
		Elem: &schema.Resource{
			Schema: map[string]*schema.Schema{
				"ok": &schema.Schema{
					Type:     schema.TypeFloat,
					Optional: true,
				},
				"warning": &schema.Schema{
					Type:     schema.TypeFloat,
					Optional: true,
				},
				"critical": &schema.Schema{
					Type:     schema.TypeFloat,
					Required: true,
				},
			},
		},
	}
}

func getThresholds(d *schema.ResourceData) (string, datadog.ThresholdCount) {
	t := datadog.ThresholdCount{}

	var threshold string

	if r, ok := d.GetOk("thresholds.ok"); ok {
		t.Ok = json.Number(r.(string))
	}

	if r, ok := d.GetOk("thresholds.warning"); ok {
		t.Warning = json.Number(r.(string))
	}

	if r, ok := d.GetOk("thresholds.critical"); ok {
		threshold = r.(string)
		t.Critical = json.Number(r.(string))
	}

	return threshold, t
}

func resourceDatadogGenericDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*datadog.Client)

	i, err := strconv.Atoi(d.Id())
	if err != nil {
		return err
	}

	if err = client.DeleteMonitor(i); err != nil {
		return err
	}

	return nil
}

func resourceDatadogGenericExists(d *schema.ResourceData, meta interface{}) (b bool, e error) {
	// Exists - This is called to verify a resource still exists. It is called prior to Read,
	// and lowers the burden of Read to be able to assume the resource exists.
	client := meta.(*datadog.Client)

	// Workaround to handle upgrades from < 0.0.4
	if strings.Contains(d.Id(), "__") {
		return false, fmt.Errorf("Monitor ID contains __, which is pre v0.0.4 old behaviour.\n    You have the following options:\n" +
			"    * Run scripts/migration_helper.py to generate a new statefile and clean up monitors\n" +
			"        * Alternatively:\n        * Delete all your metric_check resources and recreate them\n" +
			"        * Manually remove half of the resources and hack the state file")
	}

	i, err := strconv.Atoi(d.Id())
	if err != nil {
		return false, err
	}

	if _, err = client.GetMonitor(i); err != nil {
		if strings.Contains(err.Error(), "404 Not Found") {
			return false, nil
		}
		return false, err
	}

	return true, nil
}

func resourceDatadogGenericRead(d *schema.ResourceData, meta interface{}) error {
	// TODO: Add support for read function.
	/* Read - This is called to resync the local state with the remote state.
	Terraform guarantees that an existing ID will be set. This ID should be
	used to look up the resource. Any remote data should be updated into the
	local data. No changes to the remote resource are to be made.
	*/

	return nil
}

func monitorCreator(d *schema.ResourceData, meta interface{}, m *datadog.Monitor) error {
	client := meta.(*datadog.Client)

	m, err := client.CreateMonitor(m)
	if err != nil {
		return fmt.Errorf("error updating montor: %s", err.Error())
	}

	d.SetId(strconv.Itoa(m.Id))

	return nil
}

func monitorUpdater(d *schema.ResourceData, meta interface{}, m *datadog.Monitor) error {
	client := meta.(*datadog.Client)

	i, err := strconv.Atoi(d.Id())
	if err != nil {
		return err
	}

	m.Id = i

	if err = client.UpdateMonitor(m); err != nil {
		return fmt.Errorf("error updating montor: %s", err.Error())
	}

	return nil
}
<commit_msg>Add link to script.<commit_after>package datadog

import (
	"encoding/json"
	"fmt"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/zorkian/go-datadog-api"
	"strconv"
	"strings"
)

func thresholdSchema() *schema.Schema {
	return &schema.Schema{
		Type:     schema.TypeMap,
		Required: true,
		Elem: &schema.Resource{
			Schema: map[string]*schema.Schema{
				"ok": &schema.Schema{
					Type:     schema.TypeFloat,
					Optional: true,
				},
				"warning": &schema.Schema{
					Type:     schema.TypeFloat,
					Optional: true,
				},
				"critical": &schema.Schema{
					Type:     schema.TypeFloat,
					Required: true,
				},
			},
		},
	}
}

func getThresholds(d *schema.ResourceData) (string, datadog.ThresholdCount) {
	t := datadog.ThresholdCount{}

	var threshold string

	if r, ok := d.GetOk("thresholds.ok"); ok {
		t.Ok = json.Number(r.(string))
	}

	if r, ok := d.GetOk("thresholds.warning"); ok {
		t.Warning = json.Number(r.(string))
	}

	if r, ok := d.GetOk("thresholds.critical"); ok {
		threshold = r.(string)
		t.Critical = json.Number(r.(string))
	}

	return threshold, t
}

func resourceDatadogGenericDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*datadog.Client)

	i, err := strconv.Atoi(d.Id())
	if err != nil {
		return err
	}

	if err = client.DeleteMonitor(i); err != nil {
		return err
	}

	return nil
}

func resourceDatadogGenericExists(d *schema.ResourceData, meta interface{}) (b bool, e error) {
	// Exists - This is called to verify a resource still exists. It is called prior to Read,
	// and lowers the burden of Read to be able to assume the resource exists.
	client := meta.(*datadog.Client)

	// Workaround to handle upgrades from < 0.0.4
	if strings.Contains(d.Id(), "__") {
		return false, fmt.Errorf("Monitor ID contains __, which is pre v0.0.4 old behaviour.\n    You have the following options:\n" +
			"    * Run https://github.com/ojongerius/terraform-provider-datadog/blob/master/scripts/migration_helper.py to generate a new statefile and clean up monitors\n" +
			"    * Mannualy fix this by deleting all your metric_check resources and recreate them, " +
			"or manually remove half of the resources and hack the state file.\n")
	}

	i, err := strconv.Atoi(d.Id())
	if err != nil {
		return false, err
	}

	if _, err = client.GetMonitor(i); err != nil {
		if strings.Contains(err.Error(), "404 Not Found") {
			return false, nil
		}
		return false, err
	}

	return true, nil
}

func resourceDatadogGenericRead(d *schema.ResourceData, meta interface{}) error {
	// TODO: Add support for read function.
	/* Read - This is called to resync the local state with the remote state.
	Terraform guarantees that an existing ID will be set. This ID should be
	used to look up the resource. Any remote data should be updated into the
	local data. No changes to the remote resource are to be made.
	*/

	return nil
}

func monitorCreator(d *schema.ResourceData, meta interface{}, m *datadog.Monitor) error {
	client := meta.(*datadog.Client)

	m, err := client.CreateMonitor(m)
	if err != nil {
		return fmt.Errorf("error updating montor: %s", err.Error())
	}

	d.SetId(strconv.Itoa(m.Id))

	return nil
}

func monitorUpdater(d *schema.ResourceData, meta interface{}, m *datadog.Monitor) error {
	client := meta.(*datadog.Client)

	i, err := strconv.Atoi(d.Id())
	if err != nil {
		return err
	}

	m.Id = i

	if err = client.UpdateMonitor(m); err != nil {
		return fmt.Errorf("error updating montor: %s", err.Error())
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/net/context"
)

type (
	CommandConfig struct {
		Template []string            `json:"-"`
		Options  map[string][]string `json:"options,omitempty"`
		Dryrun   bool                `json:"dryrun,omitempty"`
	}

	Job struct {
		config *CommandConfig
		// https://godoc.org/google.golang.org/genproto/googleapis/pubsub/v1#ReceivedMessage
		message      *JobMessage
		notification *ProgressNotification
		storage      Storage

		// These are set at at setupWorkspace
		workspace     string
		downloads_dir string
		uploads_dir   string

		// These are set at setupDownloadFiles
		downloadFileMap     map[string]string
		remoteDownloadFiles interface{}
		localDownloadFiles  interface{}
	}
)

func (job *Job) run(ctx context.Context) error {
	verr := job.message.Validate()
	if verr != nil {
		log.Fatalf("Invalid Message: AckId: %v, Message: %v, error: %v\n", job.message.raw.AckId, job.message.raw.Message, verr)
		err := job.withNotify(CANCELLING, job.message.Ack)()
		if err != nil {
			return err
		}
		return verr
	}

	go job.message.sendMADPeriodically()
	defer job.message.Done()

	job.notification.notify(PROCESSING, job.message.MessageId(), "info")
	err := job.setupWorkspace(ctx, func() error {
		err := job.withNotify(PREPARING, job.setupDownloadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(DOWNLOADING, job.downloadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(EXECUTING, job.execute)()
		if err != nil {
			return err
		}

		err = job.withNotify(UPLOADING, job.uploadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(ACKSENDING, job.message.Ack)()
		if err != nil {
			return err
		}

		return nil
	})
	job.notification.notify(CLEANUP, job.message.MessageId(), "info")
	return err
}

func (job *Job) withNotify(progress int, f func() error) func() error {
	msg_id := job.message.MessageId()
	return func() error {
		job.notification.notify(progress, msg_id, "info")
		err := f()
		if err != nil {
			job.notification.notify(progress+2, msg_id, "error")
			return err
		}
		job.notification.notify(progress+1, msg_id, "info")
		return nil
	}
}

func (job *Job) setupWorkspace(ctx context.Context, f func() error) error {
	dir, err := ioutil.TempDir("", "workspace")
	if err != nil {
		log.Fatal(err)
		return err
	}
	defer os.RemoveAll(dir) // clean up

	subdirs := []string{
		filepath.Join(dir, "downloads"),
		filepath.Join(dir, "uploads"),
	}
	for _, subdir := range subdirs {
		err := os.MkdirAll(subdir, 0700)
		if err != nil {
			return err
		}
	}
	job.workspace = dir
	job.downloads_dir = subdirs[0]
	job.uploads_dir = subdirs[1]
	return f()
}

func (job *Job) setupDownloadFiles() error {
	job.downloadFileMap = map[string]string{}
	job.remoteDownloadFiles = job.parseJson(job.message.Attribute("download_files"))
	objects := job.flatten(job.remoteDownloadFiles)
	remoteUrls := []string{}
	for _, obj := range objects {
		switch obj.(type) {
		case string:
			remoteUrls = append(remoteUrls, obj.(string))
		default:
			log.Printf("Invalid download file URL: %v [%T]", obj, obj)
		}
	}
	for _, remote_url := range remoteUrls {
		url, err := url.Parse(remote_url)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remote_url, err)
			return err
		}
		urlstr := fmt.Sprintf("gs://%v%v", url.Host, url.Path)
		destPath := filepath.Join(job.downloads_dir, url.Host, url.Path)
		job.downloadFileMap[urlstr] = destPath
	}
	job.localDownloadFiles = job.copyWithFileMap(job.remoteDownloadFiles)
	return nil
}

func (job *Job) copyWithFileMap(obj interface{}) interface{} {
	switch obj.(type) {
	case map[string]interface{}:
		result := map[string]interface{}{}
		for k, v := range obj.(map[string]interface{}) {
			result[k] = job.copyWithFileMap(v)
		}
		return result
	case []interface{}:
		result := []interface{}{}
		for _, v := range obj.([]interface{}) {
			result = append(result, job.copyWithFileMap(v))
		}
		return result
	case string:
		return job.downloadFileMap[obj.(string)]
	default:
		return obj
	}
}

func (job *Job) buildVariable() *Variable {
	return &Variable{
		data: map[string]interface{}{
			"workspace":             job.workspace,
			"downloads_dir":         job.downloads_dir,
			"uploads_dir":           job.uploads_dir,
			"download_files":        job.localDownloadFiles,
			"local_download_files":  job.localDownloadFiles,
			"remote_download_files": job.remoteDownloadFiles,
			"attrs":                 job.message.raw.Message.Attributes,
			"attributes":            job.message.raw.Message.Attributes,
			"data":                  job.message.raw.Message.Data,
		},
	}
}

func (job *Job) build() (*exec.Cmd, error) {
	v := job.buildVariable()

	values, err := job.extract(v, job.config.Template)
	if err != nil {
		return nil, err
	}
	if len(job.config.Options) > 0 {
		key := strings.Join(values, " ")
		t := job.config.Options[key]
		if t == nil {
			t = job.config.Options["default"]
		}
		if t != nil {
			values, err = job.extract(v, t)
			if err != nil {
				return nil, err
			}
		}
	}
	cmd := exec.Command(values[0], values[1:]...)
	return cmd, nil
}

func (job *Job) extract(v *Variable, values []string) ([]string, error) {
	result := []string{}
	for _, src := range values {
		extracted, err := v.expand(src)
		if err != nil {
			return nil, err
		}
		vals := strings.Split(extracted, v.separator)
		for _, val := range vals {
			result = append(result, val)
		}
	}
	return result, nil
}

func (job *Job) downloadFiles() error {
	for remoteURL, destPath := range job.downloadFileMap {
		url, err := url.Parse(remoteURL)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remoteURL, err)
			return err
		}

		dir := path.Dir(destPath)
		err = os.MkdirAll(dir, 0700)
		if err != nil {
			return err
		}

		err = job.storage.Download(url.Host, url.Path[1:], destPath)
		if err != nil {
			return err
		}
	}
	return nil
}

func (job *Job) execute() error {
	cmd, err := job.build()
	if err != nil {
		log.Fatalf("Command build Error template: %v msg: %v cause of %v\n", job.config.Template, job.message, err)
		return err
	}
	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &out
	log.Printf("EXECUTE running: %v\n", cmd)
	err = cmd.Run()
	if err != nil {
		log.Printf("Command Error: cmd: %v cause of %v\n%v\n", cmd, err, out.String())
		return err
	}
	return nil
}

func (job *Job) uploadFiles() error {
	localPaths, err := job.listFiles(job.uploads_dir)
	if err != nil {
		return err
	}
	for _, localPath := range localPaths {
		relPath, err := filepath.Rel(job.uploads_dir, localPath)
		if err != nil {
			log.Fatalf("Error getting relative path of %v: %v\n", localPath, err)
			return err
		}
		sep := string([]rune{os.PathSeparator})
		parts := strings.Split(relPath, sep)
		bucket := parts[0]
		object := strings.Join(parts[1:], sep)
		err = job.storage.Upload(bucket, object, localPath)
		if err != nil {
			log.Fatalf("Error uploading %v to gs://%v/%v: %v\n", localPath, bucket, object, err)
			return err
		}
	}
	return nil
}

func (job *Job) listFiles(dir string) ([]string, error) {
	result := []string{}
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			result = append(result, path)
		}
		return nil
	})
	if err != nil {
		log.Fatalf("Error listing upload files: %v\n", err)
		return nil, err
	}
	return result, nil
}

func (job *Job) parseJson(str string) interface{} {
	matched, err := regexp.MatchString(`\A\[.*\]\z|\A\{.*\}\z|`, str)
	if err != nil {
		return str
	}
	if !matched {
		return str
	}
	var dest interface{}
	err = json.Unmarshal([]byte(str), &dest)
	if err != nil {
		return str
	}
	return dest
}

func (job *Job) flatten(obj interface{}) []interface{} {
	// Support only unmarshalled object from JSON
	// See https://golang.org/pkg/encoding/json/#Unmarshal also
	switch obj.(type) {
	case []interface{}:
		res := []interface{}{}
		for _, i := range obj.([]interface{}) {
			switch i.(type) {
			case bool, float64, string, nil:
				res = append(res, i)
			default:
				for _, j := range job.flatten(i) {
					res = append(res, j)
				}
			}
		}
		return res
	case map[string]interface{}:
		values := []interface{}{}
		for _, val := range obj.(map[string]interface{}) {
			values = append(values, val)
		}
		return job.flatten(values)
	default:
		return []interface{}{obj}
	}
}
<commit_msg>:+1: Don't return error for invalid message but just ignore it<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/net/context"
)

type (
	CommandConfig struct {
		Template []string            `json:"-"`
		Options  map[string][]string `json:"options,omitempty"`
		Dryrun   bool                `json:"dryrun,omitempty"`
	}

	Job struct {
		config *CommandConfig
		// https://godoc.org/google.golang.org/genproto/googleapis/pubsub/v1#ReceivedMessage
		message      *JobMessage
		notification *ProgressNotification
		storage      Storage

		// These are set at at setupWorkspace
		workspace     string
		downloads_dir string
		uploads_dir   string

		// These are set at setupDownloadFiles
		downloadFileMap     map[string]string
		remoteDownloadFiles interface{}
		localDownloadFiles  interface{}
	}
)

func (job *Job) run(ctx context.Context) error {
	verr := job.message.Validate()
	if verr != nil {
		log.Printf("Invalid Message: MessageId: %v, Message: %v, error: %v\n", job.message.MessageId(), job.message.raw.Message, verr)
		err := job.withNotify(CANCELLING, job.message.Ack)()
		if err != nil {
			return err
		}
		return nil
	}

	go job.message.sendMADPeriodically()
	defer job.message.Done()

	job.notification.notify(PROCESSING, job.message.MessageId(), "info")
	err := job.setupWorkspace(ctx, func() error {
		err := job.withNotify(PREPARING, job.setupDownloadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(DOWNLOADING, job.downloadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(EXECUTING, job.execute)()
		if err != nil {
			return err
		}

		err = job.withNotify(UPLOADING, job.uploadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(ACKSENDING, job.message.Ack)()
		if err != nil {
			return err
		}

		return nil
	})
	job.notification.notify(CLEANUP, job.message.MessageId(), "info")
	return err
}

func (job *Job) withNotify(progress int, f func() error) func() error {
	msg_id := job.message.MessageId()
	return func() error {
		job.notification.notify(progress, msg_id, "info")
		err := f()
		if err != nil {
			job.notification.notify(progress+2, msg_id, "error")
			return err
		}
		job.notification.notify(progress+1, msg_id, "info")
		return nil
	}
}

func (job *Job) setupWorkspace(ctx context.Context, f func() error) error {
	dir, err := ioutil.TempDir("", "workspace")
	if err != nil {
		log.Fatal(err)
		return err
	}
	defer os.RemoveAll(dir) // clean up

	subdirs := []string{
		filepath.Join(dir, "downloads"),
		filepath.Join(dir, "uploads"),
	}
	for _, subdir := range subdirs {
		err := os.MkdirAll(subdir, 0700)
		if err != nil {
			return err
		}
	}
	job.workspace = dir
	job.downloads_dir = subdirs[0]
	job.uploads_dir = subdirs[1]
	return f()
}

func (job *Job) setupDownloadFiles() error {
	job.downloadFileMap = map[string]string{}
	job.remoteDownloadFiles = job.parseJson(job.message.Attribute("download_files"))
	objects := job.flatten(job.remoteDownloadFiles)
	remoteUrls := []string{}
	for _, obj := range objects {
		switch obj.(type) {
		case string:
			remoteUrls = append(remoteUrls, obj.(string))
		default:
			log.Printf("Invalid download file URL: %v [%T]", obj, obj)
		}
	}
	for _, remote_url := range remoteUrls {
		url, err := url.Parse(remote_url)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remote_url, err)
			return err
		}
		urlstr := fmt.Sprintf("gs://%v%v", url.Host, url.Path)
		destPath := filepath.Join(job.downloads_dir, url.Host, url.Path)
		job.downloadFileMap[urlstr] = destPath
	}
	job.localDownloadFiles = job.copyWithFileMap(job.remoteDownloadFiles)
	return nil
}

func (job *Job) copyWithFileMap(obj interface{}) interface{} {
	switch obj.(type) {
	case map[string]interface{}:
		result := map[string]interface{}{}
		for k, v := range obj.(map[string]interface{}) {
			result[k] = job.copyWithFileMap(v)
		}
		return result
	case []interface{}:
		result := []interface{}{}
		for _, v := range obj.([]interface{}) {
			result = append(result, job.copyWithFileMap(v))
		}
		return result
	case string:
		return job.downloadFileMap[obj.(string)]
	default:
		return obj
	}
}

func (job *Job) buildVariable() *Variable {
	return &Variable{
		data: map[string]interface{}{
			"workspace":             job.workspace,
			"downloads_dir":         job.downloads_dir,
			"uploads_dir":           job.uploads_dir,
			"download_files":        job.localDownloadFiles,
			"local_download_files":  job.localDownloadFiles,
			"remote_download_files": job.remoteDownloadFiles,
			"attrs":                 job.message.raw.Message.Attributes,
			"attributes":            job.message.raw.Message.Attributes,
			"data":                  job.message.raw.Message.Data,
		},
	}
}

func (job *Job) build() (*exec.Cmd, error) {
	v := job.buildVariable()

	values, err := job.extract(v, job.config.Template)
	if err != nil {
		return nil, err
	}
	if len(job.config.Options) > 0 {
		key := strings.Join(values, " ")
		t := job.config.Options[key]
		if t == nil {
			t = job.config.Options["default"]
		}
		if t != nil {
			values, err = job.extract(v, t)
			if err != nil {
				return nil, err
			}
		}
	}
	cmd := exec.Command(values[0], values[1:]...)
	return cmd, nil
}

func (job *Job) extract(v *Variable, values []string) ([]string, error) {
	result := []string{}
	for _, src := range values {
		extracted, err := v.expand(src)
		if err != nil {
			return nil, err
		}
		vals := strings.Split(extracted, v.separator)
		for _, val := range vals {
			result = append(result, val)
		}
	}
	return result, nil
}

func (job *Job) downloadFiles() error {
	for remoteURL, destPath := range job.downloadFileMap {
		url, err := url.Parse(remoteURL)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remoteURL, err)
			return err
		}

		dir := path.Dir(destPath)
		err = os.MkdirAll(dir, 0700)
		if err != nil {
			return err
		}

		err = job.storage.Download(url.Host, url.Path[1:], destPath)
		if err != nil {
			return err
		}
	}
	return nil
}

func (job *Job) execute() error {
	cmd, err := job.build()
	if err != nil {
		log.Fatalf("Command build Error template: %v msg: %v cause of %v\n", job.config.Template, job.message, err)
		return err
	}
	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &out
	log.Printf("EXECUTE running: %v\n", cmd)
	err = cmd.Run()
	if err != nil {
		log.Printf("Command Error: cmd: %v cause of %v\n%v\n", cmd, err, out.String())
		return err
	}
	return nil
}

func (job *Job) uploadFiles() error {
	localPaths, err := job.listFiles(job.uploads_dir)
	if err != nil {
		return err
	}
	for _, localPath := range localPaths {
		relPath, err := filepath.Rel(job.uploads_dir, localPath)
		if err != nil {
			log.Fatalf("Error getting relative path of %v: %v\n", localPath, err)
			return err
		}
		sep := string([]rune{os.PathSeparator})
		parts := strings.Split(relPath, sep)
		bucket := parts[0]
		object := strings.Join(parts[1:], sep)
		err = job.storage.Upload(bucket, object, localPath)
		if err != nil {
			log.Fatalf("Error uploading %v to gs://%v/%v: %v\n", localPath, bucket, object, err)
			return err
		}
	}
	return nil
}

func (job *Job) listFiles(dir string) ([]string, error) {
	result := []string{}
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			result = append(result, path)
		}
		return nil
	})
	if err != nil {
		log.Fatalf("Error listing upload files: %v\n", err)
		return nil, err
	}
	return result, nil
}

func (job *Job) parseJson(str string) interface{} {
	matched, err := regexp.MatchString(`\A\[.*\]\z|\A\{.*\}\z|`, str)
	if err != nil {
		return str
	}
	if !matched {
		return str
	}
	var dest interface{}
	err = json.Unmarshal([]byte(str), &dest)
	if err != nil {
		return str
	}
	return dest
}

func (job *Job) flatten(obj interface{}) []interface{} {
	// Support only unmarshalled object from JSON
	// See https://golang.org/pkg/encoding/json/#Unmarshal also
	switch obj.(type) {
	case []interface{}:
		res := []interface{}{}
		for _, i := range obj.([]interface{}) {
			switch i.(type) {
			case bool, float64, string, nil:
				res = append(res, i)
			default:
				for _, j := range job.flatten(i) {
					res = append(res, j)
				}
			}
		}
		return res
	case map[string]interface{}:
		values := []interface{}{}
		for _, val := range obj.(map[string]interface{}) {
			values = append(values, val)
		}
		return job.flatten(values)
	default:
		return []interface{}{obj}
	}
}
<|endoftext|>