<commit_before>package repository

import (
	"fmt"
	"reflect"

	log "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	repositorypkg "github.com/argoproj/argo-cd/pkg/apiclient/repository"
	appsv1 "github.com/argoproj/argo-cd/pkg/apis/application/v1alpha1"
	"github.com/argoproj/argo-cd/reposerver/apiclient"
	servercache "github.com/argoproj/argo-cd/server/cache"
	"github.com/argoproj/argo-cd/server/rbacpolicy"
	"github.com/argoproj/argo-cd/util"
	"github.com/argoproj/argo-cd/util/argo"
	"github.com/argoproj/argo-cd/util/db"
	"github.com/argoproj/argo-cd/util/rbac"
	"github.com/argoproj/argo-cd/util/settings"
)

// Server provides a Repository service
type Server struct {
	db            db.ArgoDB
	repoClientset apiclient.Clientset
	enf           *rbac.Enforcer
	cache         *servercache.Cache
	settings      *settings.SettingsManager
}

// NewServer returns a new instance of the Repository service
func NewServer(
	repoClientset apiclient.Clientset,
	db db.ArgoDB,
	enf *rbac.Enforcer,
	cache *servercache.Cache,
	settings *settings.SettingsManager,
) *Server {
	return &Server{
		db:            db,
		repoClientset: repoClientset,
		enf:           enf,
		cache:         cache,
		settings:      settings,
	}
}

// Get the connection state for a given repository URL by connecting to the
// repo and evaluate the results. Unless forceRefresh is set to true, the
// result may be retrieved out of the cache.
func (s *Server) getConnectionState(ctx context.Context, url string, forceRefresh bool) appsv1.ConnectionState {
	if !forceRefresh {
		if connectionState, err := s.cache.GetRepoConnectionState(url); err == nil {
			return connectionState
		}
	}
	now := metav1.Now()
	connectionState := appsv1.ConnectionState{
		Status:     appsv1.ConnectionStatusSuccessful,
		ModifiedAt: &now,
	}
	var err error
	repo, err := s.db.GetRepository(ctx, url)
	if err == nil {
		err = argo.TestRepo(repo)
	}
	if err != nil {
		connectionState.Status = appsv1.ConnectionStatusFailed
		connectionState.Message = fmt.Sprintf("Unable to connect to repository: %v", err)
	}
	err = s.cache.SetRepoConnectionState(url, &connectionState)
	if err != nil {
		log.Warnf("getConnectionState cache set error %s: %v", url, err)
	}
	return connectionState
}

// List returns list of repositories
// Deprecated: Use ListRepositories instead
func (s *Server) List(ctx context.Context, q *repositorypkg.RepoQuery) (*appsv1.RepositoryList, error) {
	return s.ListRepositories(ctx, q)
}

// ListRepositories returns a list of all configured repositories and the state of their connections
func (s *Server) ListRepositories(ctx context.Context, q *repositorypkg.RepoQuery) (*appsv1.RepositoryList, error) {
	repos, err := s.db.ListRepositories(ctx)
	if err != nil {
		return nil, err
	}
	items := appsv1.Repositories{}
	for _, repo := range repos {
		if s.enf.Enforce(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionGet, repo.Repo) {
			// remove secrets
			items = append(items, &appsv1.Repository{
				Repo:      repo.Repo,
				Type:      repo.Type,
				Name:      repo.Name,
				Username:  repo.Username,
				Insecure:  repo.IsInsecure(),
				EnableLFS: repo.EnableLFS,
			})
		}
	}
	err = util.RunAllAsync(len(items), func(i int) error {
		items[i].ConnectionState = s.getConnectionState(ctx, items[i].Repo, q.ForceRefresh)
		return nil
	})
	if err != nil {
		return nil, err
	}
	return &appsv1.RepositoryList{Items: items}, nil
}

// ListApps returns list of apps in the repo
func (s *Server) ListApps(ctx context.Context, q *repositorypkg.RepoAppsQuery) (*repositorypkg.RepoAppsResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionGet, q.Repo); err != nil {
		return nil, err
	}
	repo, err := s.db.GetRepository(ctx, q.Repo)
	if err != nil {
		return nil, err
	}

	// Test the repo
	conn, repoClient, err := s.repoClientset.NewRepoServerClient()
	if err != nil {
		return nil, err
	}
	defer util.Close(conn)

	apps, err := repoClient.ListApps(ctx, &apiclient.ListAppsRequest{
		Repo:     repo,
		Revision: q.Revision,
	})
	if err != nil {
		return nil, err
	}
	items := make([]*repositorypkg.AppInfo, 0)
	for app, appType := range apps.Apps {
		items = append(items, &repositorypkg.AppInfo{Path: app, Type: appType})
	}
	return &repositorypkg.RepoAppsResponse{Items: items}, nil
}

func (s *Server) GetAppDetails(ctx context.Context, q *repositorypkg.RepoAppDetailsQuery) (*apiclient.RepoAppDetailsResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionGet, q.Source.RepoURL); err != nil {
		return nil, err
	}
	repo, err := s.db.GetRepository(ctx, q.Source.RepoURL)
	if err != nil {
		return nil, err
	}
	conn, repoClient, err := s.repoClientset.NewRepoServerClient()
	if err != nil {
		return nil, err
	}
	defer util.Close(conn)
	helmRepos, err := s.db.ListHelmRepositories(ctx)
	if err != nil {
		return nil, err
	}
	buildOptions, err := s.settings.GetKustomizeBuildOptions()
	if err != nil {
		return nil, err
	}
	return repoClient.GetAppDetails(ctx, &apiclient.RepoServerAppDetailsQuery{
		Repo:   repo,
		Source: q.Source,
		Repos:  helmRepos,
		KustomizeOptions: &appsv1.KustomizeOptions{
			BuildOptions: buildOptions,
		},
	})
}

func (s *Server) GetHelmCharts(ctx context.Context, q *repositorypkg.RepoQuery) (*apiclient.HelmChartsResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionGet, q.Repo); err != nil {
		return nil, err
	}
	repo, err := s.db.GetRepository(ctx, q.Repo)
	if err != nil {
		return nil, err
	}
	conn, repoClient, err := s.repoClientset.NewRepoServerClient()
	if err != nil {
		return nil, err
	}
	defer util.Close(conn)
	return repoClient.GetHelmCharts(ctx, &apiclient.HelmChartsRequest{Repo: repo})
}

// Create creates a repository or repository credential set
// Deprecated: Use CreateRepository() instead
func (s *Server) Create(ctx context.Context, q *repositorypkg.RepoCreateRequest) (*appsv1.Repository, error) {
	return s.CreateRepository(ctx, q)
}

// CreateRepository creates a repository configuration
func (s *Server) CreateRepository(ctx context.Context, q *repositorypkg.RepoCreateRequest) (*appsv1.Repository, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionCreate, q.Repo.Repo); err != nil {
		return nil, err
	}

	var repo *appsv1.Repository
	var err error

	// check we can connect to the repo, copying any existing creds
	{
		repo := q.Repo.DeepCopy()
		if !repo.HasCredentials() {
			creds, err := s.db.GetRepositoryCredentials(ctx, repo.Repo)
			if err != nil {
				return nil, err
			}
			repo.CopyCredentialsFrom(creds)
		}
		err = argo.TestRepo(repo)
		if err != nil {
			return nil, err
		}
	}

	r := q.Repo
	r.ConnectionState = appsv1.ConnectionState{Status: appsv1.ConnectionStatusSuccessful}
	repo, err = s.db.CreateRepository(ctx, r)
	if status.Convert(err).Code() == codes.AlreadyExists {
		// act idempotent if existing spec matches new spec
		existing, getErr := s.db.GetRepository(ctx, r.Repo)
		if getErr != nil {
			return nil, status.Errorf(codes.Internal, "unable to check existing repository details: %v", getErr)
		}

		// repository ConnectionState may differ, so make consistent before testing
		existing.ConnectionState = r.ConnectionState
		if reflect.DeepEqual(existing, r) {
			repo, err = existing, nil
		} else if q.Upsert {
			return s.UpdateRepository(ctx, &repositorypkg.RepoUpdateRequest{Repo: r})
		} else {
			return nil, status.Errorf(codes.InvalidArgument, "existing repository spec is different; use upsert flag to force update")
		}
	}
	return &appsv1.Repository{Repo: repo.Repo, Type: repo.Type, Name: repo.Name}, err
}

// Update updates a repository or credential set
// Deprecated: Use UpdateRepository() instead
func (s *Server) Update(ctx context.Context, q *repositorypkg.RepoUpdateRequest) (*appsv1.Repository, error) {
	return s.UpdateRepository(ctx, q)
}

// UpdateRepository updates a repository configuration
func (s *Server) UpdateRepository(ctx context.Context, q *repositorypkg.RepoUpdateRequest) (*appsv1.Repository, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionUpdate, q.Repo.Repo); err != nil {
		return nil, err
	}
	_, err := s.db.UpdateRepository(ctx, q.Repo)
	return &appsv1.Repository{Repo: q.Repo.Repo, Type: q.Repo.Type, Name: q.Repo.Name}, err
}

// Delete removes a repository from the configuration
// Deprecated: Use DeleteRepository() instead
func (s *Server) Delete(ctx context.Context, q *repositorypkg.RepoQuery) (*repositorypkg.RepoResponse, error) {
	return s.DeleteRepository(ctx, q)
}

// DeleteRepository removes a repository from the configuration
func (s *Server) DeleteRepository(ctx context.Context, q *repositorypkg.RepoQuery) (*repositorypkg.RepoResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionDelete, q.Repo); err != nil {
		return nil, err
	}

	// invalidate cache
	if err := s.cache.SetRepoConnectionState(q.Repo, nil); err == nil {
		log.Errorf("error invalidating cache: %v", err)
	}

	err := s.db.DeleteRepository(ctx, q.Repo)
	return &repositorypkg.RepoResponse{}, err
}

// ValidateAccess checks whether access to a repository is possible with the
// given URL and credentials.
func (s *Server) ValidateAccess(ctx context.Context, q *repositorypkg.RepoAccessQuery) (*repositorypkg.RepoResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionCreate, q.Repo); err != nil {
		return nil, err
	}

	repo := &appsv1.Repository{
		Repo:              q.Repo,
		Type:              q.Type,
		Name:              q.Name,
		Username:          q.Username,
		Password:          q.Password,
		SSHPrivateKey:     q.SshPrivateKey,
		Insecure:          q.Insecure,
		TLSClientCertData: q.TlsClientCertData,
		TLSClientCertKey:  q.TlsClientCertKey,
	}

	var repoCreds *appsv1.RepoCreds
	var err error

	// If repo does not have credentials, check if there are credentials stored
	// for it and if yes, copy them
	if !repo.HasCredentials() {
		repoCreds, err = s.db.GetRepositoryCredentials(ctx, q.Repo)
		if err != nil {
			return nil, err
		}
		if repoCreds != nil {
			repo.CopyCredentialsFrom(repoCreds)
		}
	}
	err = argo.TestRepo(repo)
	if err != nil {
		return nil, err
	}
	return &repositorypkg.RepoResponse{}, nil
}
<commit_msg>Fixes panic when creating repo (#2568)<commit_after>package repository

import (
	"fmt"
	"reflect"

	log "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	repositorypkg "github.com/argoproj/argo-cd/pkg/apiclient/repository"
	appsv1 "github.com/argoproj/argo-cd/pkg/apis/application/v1alpha1"
	"github.com/argoproj/argo-cd/reposerver/apiclient"
	servercache "github.com/argoproj/argo-cd/server/cache"
	"github.com/argoproj/argo-cd/server/rbacpolicy"
	"github.com/argoproj/argo-cd/util"
	"github.com/argoproj/argo-cd/util/argo"
	"github.com/argoproj/argo-cd/util/db"
	"github.com/argoproj/argo-cd/util/rbac"
	"github.com/argoproj/argo-cd/util/settings"
)

// Server provides a Repository service
type Server struct {
	db            db.ArgoDB
	repoClientset apiclient.Clientset
	enf           *rbac.Enforcer
	cache         *servercache.Cache
	settings      *settings.SettingsManager
}

// NewServer returns a new instance of the Repository service
func NewServer(
	repoClientset apiclient.Clientset,
	db db.ArgoDB,
	enf *rbac.Enforcer,
	cache *servercache.Cache,
	settings *settings.SettingsManager,
) *Server {
	return &Server{
		db:            db,
		repoClientset: repoClientset,
		enf:           enf,
		cache:         cache,
		settings:      settings,
	}
}

// Get the connection state for a given repository URL by connecting to the
// repo and evaluate the results. Unless forceRefresh is set to true, the
// result may be retrieved out of the cache.
func (s *Server) getConnectionState(ctx context.Context, url string, forceRefresh bool) appsv1.ConnectionState {
	if !forceRefresh {
		if connectionState, err := s.cache.GetRepoConnectionState(url); err == nil {
			return connectionState
		}
	}
	now := metav1.Now()
	connectionState := appsv1.ConnectionState{
		Status:     appsv1.ConnectionStatusSuccessful,
		ModifiedAt: &now,
	}
	var err error
	repo, err := s.db.GetRepository(ctx, url)
	if err == nil {
		err = argo.TestRepo(repo)
	}
	if err != nil {
		connectionState.Status = appsv1.ConnectionStatusFailed
		connectionState.Message = fmt.Sprintf("Unable to connect to repository: %v", err)
	}
	err = s.cache.SetRepoConnectionState(url, &connectionState)
	if err != nil {
		log.Warnf("getConnectionState cache set error %s: %v", url, err)
	}
	return connectionState
}

// List returns list of repositories
// Deprecated: Use ListRepositories instead
func (s *Server) List(ctx context.Context, q *repositorypkg.RepoQuery) (*appsv1.RepositoryList, error) {
	return s.ListRepositories(ctx, q)
}

// ListRepositories returns a list of all configured repositories and the state of their connections
func (s *Server) ListRepositories(ctx context.Context, q *repositorypkg.RepoQuery) (*appsv1.RepositoryList, error) {
	repos, err := s.db.ListRepositories(ctx)
	if err != nil {
		return nil, err
	}
	items := appsv1.Repositories{}
	for _, repo := range repos {
		if s.enf.Enforce(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionGet, repo.Repo) {
			// remove secrets
			items = append(items, &appsv1.Repository{
				Repo:      repo.Repo,
				Type:      repo.Type,
				Name:      repo.Name,
				Username:  repo.Username,
				Insecure:  repo.IsInsecure(),
				EnableLFS: repo.EnableLFS,
			})
		}
	}
	err = util.RunAllAsync(len(items), func(i int) error {
		items[i].ConnectionState = s.getConnectionState(ctx, items[i].Repo, q.ForceRefresh)
		return nil
	})
	if err != nil {
		return nil, err
	}
	return &appsv1.RepositoryList{Items: items}, nil
}

// ListApps returns list of apps in the repo
func (s *Server) ListApps(ctx context.Context, q *repositorypkg.RepoAppsQuery) (*repositorypkg.RepoAppsResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionGet, q.Repo); err != nil {
		return nil, err
	}
	repo, err := s.db.GetRepository(ctx, q.Repo)
	if err != nil {
		return nil, err
	}

	// Test the repo
	conn, repoClient, err := s.repoClientset.NewRepoServerClient()
	if err != nil {
		return nil, err
	}
	defer util.Close(conn)

	apps, err := repoClient.ListApps(ctx, &apiclient.ListAppsRequest{
		Repo:     repo,
		Revision: q.Revision,
	})
	if err != nil {
		return nil, err
	}
	items := make([]*repositorypkg.AppInfo, 0)
	for app, appType := range apps.Apps {
		items = append(items, &repositorypkg.AppInfo{Path: app, Type: appType})
	}
	return &repositorypkg.RepoAppsResponse{Items: items}, nil
}

func (s *Server) GetAppDetails(ctx context.Context, q *repositorypkg.RepoAppDetailsQuery) (*apiclient.RepoAppDetailsResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionGet, q.Source.RepoURL); err != nil {
		return nil, err
	}
	repo, err := s.db.GetRepository(ctx, q.Source.RepoURL)
	if err != nil {
		return nil, err
	}
	conn, repoClient, err := s.repoClientset.NewRepoServerClient()
	if err != nil {
		return nil, err
	}
	defer util.Close(conn)
	helmRepos, err := s.db.ListHelmRepositories(ctx)
	if err != nil {
		return nil, err
	}
	buildOptions, err := s.settings.GetKustomizeBuildOptions()
	if err != nil {
		return nil, err
	}
	return repoClient.GetAppDetails(ctx, &apiclient.RepoServerAppDetailsQuery{
		Repo:   repo,
		Source: q.Source,
		Repos:  helmRepos,
		KustomizeOptions: &appsv1.KustomizeOptions{
			BuildOptions: buildOptions,
		},
	})
}

func (s *Server) GetHelmCharts(ctx context.Context, q *repositorypkg.RepoQuery) (*apiclient.HelmChartsResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionGet, q.Repo); err != nil {
		return nil, err
	}
	repo, err := s.db.GetRepository(ctx, q.Repo)
	if err != nil {
		return nil, err
	}
	conn, repoClient, err := s.repoClientset.NewRepoServerClient()
	if err != nil {
		return nil, err
	}
	defer util.Close(conn)
	return repoClient.GetHelmCharts(ctx, &apiclient.HelmChartsRequest{Repo: repo})
}

// Create creates a repository or repository credential set
// Deprecated: Use CreateRepository() instead
func (s *Server) Create(ctx context.Context, q *repositorypkg.RepoCreateRequest) (*appsv1.Repository, error) {
	return s.CreateRepository(ctx, q)
}

// CreateRepository creates a repository configuration
func (s *Server) CreateRepository(ctx context.Context, q *repositorypkg.RepoCreateRequest) (*appsv1.Repository, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionCreate, q.Repo.Repo); err != nil {
		return nil, err
	}

	var repo *appsv1.Repository
	var err error

	// check we can connect to the repo, copying any existing creds
	{
		repo := q.Repo.DeepCopy()
		if !repo.HasCredentials() {
			creds, err := s.db.GetRepositoryCredentials(ctx, repo.Repo)
			if err != nil {
				return nil, err
			}
			repo.CopyCredentialsFrom(creds)
		}
		err = argo.TestRepo(repo)
		if err != nil {
			return nil, err
		}
	}

	r := q.Repo
	r.ConnectionState = appsv1.ConnectionState{Status: appsv1.ConnectionStatusSuccessful}
	repo, err = s.db.CreateRepository(ctx, r)
	if status.Convert(err).Code() == codes.AlreadyExists {
		// act idempotent if existing spec matches new spec
		existing, getErr := s.db.GetRepository(ctx, r.Repo)
		if getErr != nil {
			return nil, status.Errorf(codes.Internal, "unable to check existing repository details: %v", getErr)
		}

		// repository ConnectionState may differ, so make consistent before testing
		existing.ConnectionState = r.ConnectionState
		if reflect.DeepEqual(existing, r) {
			repo, err = existing, nil
		} else if q.Upsert {
			return s.UpdateRepository(ctx, &repositorypkg.RepoUpdateRequest{Repo: r})
		} else {
			return nil, status.Errorf(codes.InvalidArgument, "existing repository spec is different; use upsert flag to force update")
		}
	}
	if err != nil {
		return nil, err
	}
	return &appsv1.Repository{Repo: repo.Repo, Type: repo.Type, Name: repo.Name}, nil
}

// Update updates a repository or credential set
// Deprecated: Use UpdateRepository() instead
func (s *Server) Update(ctx context.Context, q *repositorypkg.RepoUpdateRequest) (*appsv1.Repository, error) {
	return s.UpdateRepository(ctx, q)
}

// UpdateRepository updates a repository configuration
func (s *Server) UpdateRepository(ctx context.Context, q *repositorypkg.RepoUpdateRequest) (*appsv1.Repository, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionUpdate, q.Repo.Repo); err != nil {
		return nil, err
	}
	_, err := s.db.UpdateRepository(ctx, q.Repo)
	return &appsv1.Repository{Repo: q.Repo.Repo, Type: q.Repo.Type, Name: q.Repo.Name}, err
}

// Delete removes a repository from the configuration
// Deprecated: Use DeleteRepository() instead
func (s *Server) Delete(ctx context.Context, q *repositorypkg.RepoQuery) (*repositorypkg.RepoResponse, error) {
	return s.DeleteRepository(ctx, q)
}

// DeleteRepository removes a repository from the configuration
func (s *Server) DeleteRepository(ctx context.Context, q *repositorypkg.RepoQuery) (*repositorypkg.RepoResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionDelete, q.Repo); err != nil {
		return nil, err
	}

	// invalidate cache
	if err := s.cache.SetRepoConnectionState(q.Repo, nil); err == nil {
		log.Errorf("error invalidating cache: %v", err)
	}

	err := s.db.DeleteRepository(ctx, q.Repo)
	return &repositorypkg.RepoResponse{}, err
}

// ValidateAccess checks whether access to a repository is possible with the
// given URL and credentials.
func (s *Server) ValidateAccess(ctx context.Context, q *repositorypkg.RepoAccessQuery) (*repositorypkg.RepoResponse, error) {
	if err := s.enf.EnforceErr(ctx.Value("claims"), rbacpolicy.ResourceRepositories, rbacpolicy.ActionCreate, q.Repo); err != nil {
		return nil, err
	}

	repo := &appsv1.Repository{
		Repo:              q.Repo,
		Type:              q.Type,
		Name:              q.Name,
		Username:          q.Username,
		Password:          q.Password,
		SSHPrivateKey:     q.SshPrivateKey,
		Insecure:          q.Insecure,
		TLSClientCertData: q.TlsClientCertData,
		TLSClientCertKey:  q.TlsClientCertKey,
	}

	var repoCreds *appsv1.RepoCreds
	var err error

	// If repo does not have credentials, check if there are credentials stored
	// for it and if yes, copy them
	if !repo.HasCredentials() {
		repoCreds, err = s.db.GetRepositoryCredentials(ctx, q.Repo)
		if err != nil {
			return nil, err
		}
		if repoCreds != nil {
			repo.CopyCredentialsFrom(repoCreds)
		}
	}
	err = argo.TestRepo(repo)
	if err != nil {
		return nil, err
	}
	return &repositorypkg.RepoResponse{}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fatih/structs"
	"github.com/gorilla/mux"
	"github.com/malice-plugins/go-plugin-utils/database/elasticsearch"
	"github.com/malice-plugins/go-plugin-utils/utils"
	"github.com/parnurzeal/gorequest"
	"github.com/willf/bloom"
	"gopkg.in/urfave/cli.v1"
)

var (
	// Version stores the plugin's version
	Version string

	// BuildTime stores the plugin's build time
	BuildTime string

	// ErrorRate stores the bloomfilter desired error-rate
	ErrorRate string

	// HashType is the type of hash to use to build the bloomfilter
	HashType string
)

const (
	// NSRL fields
	sha1         = 0
	md5          = 1
	crc32        = 2
	fileName     = 3
	fileSize     = 4
	productCode  = 5
	opSystemCode = 6
	specialCode  = 7
)

const (
	name     = "nsrl"
	category = "intel"
)

type pluginResults struct {
	ID   string      `json:"id" gorethink:"id,omitempty"`
	Data ResultsData `json:"nsrl" gorethink:"nsrl"`
}

// Nsrl json object
type Nsrl struct {
	Results ResultsData `json:"nsrl"`
}

// ResultsData json object
type ResultsData struct {
	Found    bool   `json:"found"`
	Hash     string `json:"hash"`
	MarkDown string `json:"markdown,omitempty" structs:"markdown,omitempty"`
}

func generateMarkDownTable(n Nsrl) string {
	var tplOut bytes.Buffer

	t := template.Must(template.New("nsrl").Parse(tpl))

	err := t.Execute(&tplOut, n)
	if err != nil {
		log.Println("executing template:", err)
	}

	return tplOut.String()
}

func lineCounter(r io.Reader) (uint64, error) {
	buf := make([]byte, 32*1024)
	var count uint64
	lineSep := []byte{'\n'}

	for {
		c, err := r.Read(buf)
		count += uint64(bytes.Count(buf[:c], lineSep))

		switch {
		case err == io.EOF:
			return count, nil

		case err != nil:
			return count, err
		}
	}
}

// build bloomfilter from NSRL database
func buildFilter() {
	var err error
	// open NSRL database
	nsrlDB, err := os.Open("NSRLFile.txt")
	utils.Assert(err)
	defer nsrlDB.Close()
	// count lines in NSRL database
	lines, err := lineCounter(nsrlDB)
	utils.Assert(err)
	log.Debugf("Number of lines in NSRLFile.txt: %d", lines)
	// write line count to file LINECOUNT
	buf := new(bytes.Buffer)
	utils.Assert(binary.Write(buf, binary.LittleEndian, lines))
	utils.Assert(ioutil.WriteFile("LINECOUNT", buf.Bytes(), 0644))

	// Create new bloomfilter with size = number of lines in NSRL database
	erate, err := strconv.ParseFloat(ErrorRate, 64)
	utils.Assert(err)

	filter := bloom.NewWithEstimates(uint(lines), erate)

	// jump back to the begining of the file
	_, err = nsrlDB.Seek(0, io.SeekStart)
	utils.Assert(err)

	log.Debug("Loading NSRL database into bloomfilter")
	reader := csv.NewReader(nsrlDB)
	// strip off csv header
	_, _ = reader.Read()
	for {
		record, err := reader.Read()

		if err == io.EOF {
			break
		}
		utils.Assert(err)

		// log.Debug(record)
		filter.Add([]byte(record[sha1]))
	}

	bloomFile, err := os.Create("nsrl.bloom")
	utils.Assert(err)
	defer bloomFile.Close()

	log.Debug("Writing bloomfilter to disk")
	_, err = filter.WriteTo(bloomFile)
	utils.Assert(err)
}

// lookUp queries the NSRL bloomfilter for a hash
func lookUp(hash string, timeout int) ResultsData {

	var lines uint64
	nsrlResults := ResultsData{}

	// read line count from file LINECOUNT
	lineCount, err := ioutil.ReadFile("LINECOUNT")
	utils.Assert(err)
	buf := bytes.NewReader(lineCount)
	utils.Assert(binary.Read(buf, binary.LittleEndian, &lines))
	log.Debugf("Number of lines in NSRLFile.txt: %d", lines)

	// Create new bloomfilter with size = number of lines in NSRL database
	erate, err := strconv.ParseFloat(ErrorRate, 64)
	utils.Assert(err)

	filter := bloom.NewWithEstimates(uint(lines), erate)

	// load NSRL bloomfilter from file
	f, err := os.Open("nsrl.bloom")
	utils.Assert(err)
	_, err = filter.ReadFrom(f)
	utils.Assert(err)

	// test of existance of hash in bloomfilter
	nsrlResults.Found = filter.TestString(hash)

	return nsrlResults
}

func webService() {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/lookup/{sha1}", webLookUp)
	log.Info("web service listening on port :3993")
	log.Fatal(http.ListenAndServe(":3993", router))
}

func webLookUp(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	hash := vars["sha1"]

	hashType, _ := utils.GetHashType(hash)

	if strings.EqualFold(hashType, "sha1") {
		nsrl := Nsrl{Results: lookUp(strings.ToUpper(hash), 10)}

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		if nsrl.Results.Found {
			w.WriteHeader(http.StatusOK)
		} else {
			w.WriteHeader(http.StatusNotFound)
		}

		if err := json.NewEncoder(w).Encode(nsrl); err != nil {
			panic(err)
		}
	} else {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(w, "Please supply a proper SHA1 hash to query")
	}
}

func printStatus(resp gorequest.Response, body string, errs []error) {
	fmt.Println(body)
}

func main() {

	var elastic string

	cli.AppHelpTemplate = utils.AppHelpTemplate
	app := cli.NewApp()

	app.Name = "nsrl"
	app.Author = "blacktop"
	app.Email = "https://github.com/blacktop"
	app.Version = Version + ", BuildTime: " + BuildTime
	app.Compiled, _ = time.Parse("20060102", BuildTime)
	app.Usage = "Malice NSRL Plugin"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, V",
			Usage: "verbose output",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:    "web",
			Aliases: []string{"w"},
			Usage:   "Create a NSRL lookup web service",
			Action: func(c *cli.Context) error {
				webService()
				return nil
			},
		},
		{
			Name:    "build",
			Aliases: []string{"b"},
			Usage:   "Build bloomfilter from NSRL database",
			Action: func(c *cli.Context) error {
				if c.GlobalBool("verbose") {
					log.SetLevel(log.DebugLevel)
				}

				buildFilter()
				return nil
			},
		},
		{
			Name:      "lookup",
			Aliases:   []string{"l"},
			Usage:     "Query NSRL for hash",
			ArgsUsage: "SHA1 to query NSRL with",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:        "elasitcsearch",
					Value:       "",
					Usage:       "elasitcsearch address for Malice to store results",
					EnvVar:      "MALICE_ELASTICSEARCH",
					Destination: &elastic,
				},
				cli.BoolFlag{
					Name:   "post, p",
					Usage:  "POST results to Malice webhook",
					EnvVar: "MALICE_ENDPOINT",
				},
				cli.BoolFlag{
					Name:   "proxy, x",
					Usage:  "proxy settings for Malice webhook endpoint",
					EnvVar: "MALICE_PROXY",
				},
				cli.IntFlag{
					Name:   "timeout",
					Value:  10,
					Usage:  "malice plugin timeout (in seconds)",
					EnvVar: "MALICE_TIMEOUT",
				},
				cli.BoolFlag{
					Name:  "table, t",
					Usage: "output as Markdown table",
				},
			},
			Action: func(c *cli.Context) error {
				if c.Args().Present() {
					hash := strings.ToUpper(c.Args().First())
					hashType, _ := utils.GetHashType(hash)

					if !strings.EqualFold(hashType, "sha1") {
						log.Fatal(fmt.Errorf("please supply a valid SHA1 hash to query NSRL with"))
					}

					if c.GlobalBool("verbose") {
						log.SetLevel(log.DebugLevel)
					}

					nsrl := Nsrl{Results: lookUp(hash, c.GlobalInt("timeout"))}
					nsrl.Results.Hash = hash
					nsrl.Results.MarkDown = generateMarkDownTable(nsrl)

					// upsert into Database
					elasticsearch.InitElasticSearch(elastic)
					elasticsearch.WritePluginResultsToDatabase(elasticsearch.PluginResults{
						ID:       utils.Getopt("MALICE_SCANID", hash),
						Name:     name,
						Category: category,
						Data:     structs.Map(nsrl.Results),
					})

					if c.GlobalBool("table") {
						fmt.Println(nsrl.Results.MarkDown)
					} else {
						nsrl.Results.MarkDown = ""
						nsrlJSON, err := json.Marshal(nsrl)
						utils.Assert(err)
						if c.GlobalBool("post") {
							request := gorequest.New()
							if c.GlobalBool("proxy") {
								request = gorequest.New().Proxy(os.Getenv("MALICE_PROXY"))
							}
							request.Post(os.Getenv("MALICE_ENDPOINT")).
								Set("X-Malice-ID", utils.Getopt("MALICE_SCANID", hash)).
								Send(string(nsrlJSON)).
								End(printStatus)

							return nil
						}
						fmt.Println(string(nsrlJSON))
					}
				} else {
					log.Fatal(fmt.Errorf("please supply a SHA1 hash to query NSRL with"))
				}
				return nil
			},
		},
	}

	err := app.Run(os.Args)
	utils.Assert(err)
}
<commit_msg>fix flags<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fatih/structs"
	"github.com/gorilla/mux"
	"github.com/malice-plugins/go-plugin-utils/database/elasticsearch"
	"github.com/malice-plugins/go-plugin-utils/utils"
	"github.com/parnurzeal/gorequest"
	"github.com/willf/bloom"
	"gopkg.in/urfave/cli.v1"
)

var (
	// Version stores the plugin's version
	Version string

	// BuildTime stores the plugin's build time
	BuildTime string

	// ErrorRate stores the bloomfilter desired error-rate
	ErrorRate string

	// HashType is the type of hash to use to build the bloomfilter
	HashType string
)

const (
	// NSRL fields
	sha1         = 0
	md5          = 1
	crc32        = 2
	fileName     = 3
	fileSize     = 4
	productCode  = 5
	opSystemCode = 6
	specialCode  = 7
)

const (
	name     = "nsrl"
	category = "intel"
)

type pluginResults struct {
	ID   string      `json:"id" gorethink:"id,omitempty"`
	Data ResultsData `json:"nsrl" gorethink:"nsrl"`
}

// Nsrl json object
type Nsrl struct {
	Results ResultsData `json:"nsrl"`
}

// ResultsData json object
type ResultsData struct {
	Found    bool   `json:"found"`
	Hash     string `json:"hash"`
	MarkDown string `json:"markdown,omitempty" structs:"markdown,omitempty"`
}

func generateMarkDownTable(n Nsrl) string {
	var tplOut bytes.Buffer

	t := template.Must(template.New("nsrl").Parse(tpl))

	err := t.Execute(&tplOut, n)
	if err != nil {
		log.Println("executing template:", err)
	}

	return tplOut.String()
}

func lineCounter(r io.Reader) (uint64, error) {
	buf := make([]byte, 32*1024)
	var count uint64
	lineSep := []byte{'\n'}

	for {
		c, err := r.Read(buf)
		count += uint64(bytes.Count(buf[:c], lineSep))

		switch {
		case err == io.EOF:
			return count, nil

		case err != nil:
			return count, err
		}
	}
}

// build bloomfilter from NSRL database
func buildFilter() {
	var err error
	// open NSRL database
	nsrlDB, err := os.Open("NSRLFile.txt")
	utils.Assert(err)
	defer nsrlDB.Close()
	// count lines in NSRL database
	lines, err := lineCounter(nsrlDB)
	utils.Assert(err)
	log.Debugf("Number of lines in NSRLFile.txt: %d", lines)
	// write line count to file LINECOUNT
	buf := new(bytes.Buffer)
	utils.Assert(binary.Write(buf, binary.LittleEndian, lines))
	utils.Assert(ioutil.WriteFile("LINECOUNT", buf.Bytes(), 0644))

	// Create new bloomfilter with size = number of lines in NSRL database
	erate, err := strconv.ParseFloat(ErrorRate, 64)
	utils.Assert(err)

	filter := bloom.NewWithEstimates(uint(lines), erate)

	// jump back to the begining of the file
	_, err = nsrlDB.Seek(0, io.SeekStart)
	utils.Assert(err)

	log.Debug("Loading NSRL database into bloomfilter")
	reader := csv.NewReader(nsrlDB)
	// strip off csv header
	_, _ = reader.Read()
	for {
		record, err := reader.Read()

		if err == io.EOF {
			break
		}
		utils.Assert(err)

		// log.Debug(record)
		filter.Add([]byte(record[sha1]))
	}

	bloomFile, err := os.Create("nsrl.bloom")
	utils.Assert(err)
	defer bloomFile.Close()

	log.Debug("Writing bloomfilter to disk")
	_, err = filter.WriteTo(bloomFile)
	utils.Assert(err)
}

// lookUp queries the NSRL bloomfilter for a hash
func lookUp(hash string, timeout int) ResultsData {

	var lines uint64
	nsrlResults := ResultsData{}

	// read line count from file LINECOUNT
	lineCount, err := ioutil.ReadFile("LINECOUNT")
	utils.Assert(err)
	buf := bytes.NewReader(lineCount)
	utils.Assert(binary.Read(buf, binary.LittleEndian, &lines))
	log.Debugf("Number of lines in NSRLFile.txt: %d", lines)

	// Create new bloomfilter with size = number of lines in NSRL database
	erate, err := strconv.ParseFloat(ErrorRate, 64)
	utils.Assert(err)

	filter := bloom.NewWithEstimates(uint(lines), erate)

	// load NSRL bloomfilter from file
	f, err := os.Open("nsrl.bloom")
	utils.Assert(err)
	_, err = filter.ReadFrom(f)
	utils.Assert(err)

	// test of existance of hash in bloomfilter
	nsrlResults.Found = filter.TestString(hash)

	return nsrlResults
}

func webService() {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/lookup/{sha1}", webLookUp)
	log.Info("web service listening on port :3993")
	log.Fatal(http.ListenAndServe(":3993", router))
}

func webLookUp(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	hash := vars["sha1"]

	hashType, _ := utils.GetHashType(hash)

	if strings.EqualFold(hashType, "sha1") {
		nsrl := Nsrl{Results: lookUp(strings.ToUpper(hash), 10)}

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		if nsrl.Results.Found {
			w.WriteHeader(http.StatusOK)
		} else {
			w.WriteHeader(http.StatusNotFound)
		}

		if err := json.NewEncoder(w).Encode(nsrl); err != nil {
			panic(err)
		}
	} else {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(w, "Please supply a proper SHA1 hash to query")
	}
}

func printStatus(resp gorequest.Response, body string, errs []error) {
	fmt.Println(body)
}

func main() {

	var elastic string

	cli.AppHelpTemplate = utils.AppHelpTemplate
	app := cli.NewApp()

	app.Name = "nsrl"
	app.Author = "blacktop"
	app.Email = "https://github.com/blacktop"
	app.Version = Version + ", BuildTime: " + BuildTime
	app.Compiled, _ = time.Parse("20060102", BuildTime)
	app.Usage = "Malice NSRL Plugin"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, V",
			Usage: "verbose output",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:    "web",
			Aliases: []string{"w"},
			Usage:   "Create a NSRL lookup web service",
			Action: func(c *cli.Context) error {
				webService()
				return nil
			},
		},
		{
			Name:    "build",
			Aliases: []string{"b"},
			Usage:   "Build bloomfilter from NSRL database",
			Action: func(c *cli.Context) error {
				if c.GlobalBool("verbose") {
					log.SetLevel(log.DebugLevel)
				}

				buildFilter()
				return nil
			},
		},
		{
			Name:      "lookup",
			Aliases:   []string{"l"},
			Usage:     "Query NSRL for hash",
			ArgsUsage: "SHA1 to query NSRL with",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:        "elasitcsearch",
					Value:       "",
					Usage:       "elasitcsearch address for Malice to store results",
					EnvVar:      "MALICE_ELASTICSEARCH",
					Destination: &elastic,
				},
				cli.BoolFlag{
					Name:   "post, p",
					Usage:  "POST results to Malice webhook",
					EnvVar: "MALICE_ENDPOINT",
				},
				cli.BoolFlag{
					Name:   "proxy, x",
					Usage:  "proxy settings for Malice webhook endpoint",
					EnvVar: "MALICE_PROXY",
				},
				cli.IntFlag{
					Name:   "timeout",
					Value:  10,
					Usage:  "malice plugin timeout (in seconds)",
					EnvVar: "MALICE_TIMEOUT",
				},
				cli.BoolFlag{
					Name:  "table, t",
					Usage: "output as Markdown table",
				},
			},
			Action: func(c *cli.Context) error {
				if c.Args().Present() {
					hash := strings.ToUpper(c.Args().First())
					hashType, _ := utils.GetHashType(hash)

					if !strings.EqualFold(hashType, "sha1") {
						log.Fatal(fmt.Errorf("please supply a valid SHA1 hash to query NSRL with"))
					}

					if c.GlobalBool("verbose") {
						log.SetLevel(log.DebugLevel)
					}

					nsrl := Nsrl{Results: lookUp(hash, c.GlobalInt("timeout"))}
					nsrl.Results.Hash = hash
					nsrl.Results.MarkDown = generateMarkDownTable(nsrl)

					// upsert into Database
					elasticsearch.InitElasticSearch(elastic)
					elasticsearch.WritePluginResultsToDatabase(elasticsearch.PluginResults{
						ID:       utils.Getopt("MALICE_SCANID", hash),
						Name:     name,
						Category: category,
						Data:     structs.Map(nsrl.Results),
					})

					if c.Bool("table") {
						fmt.Println(nsrl.Results.MarkDown)
					} else {
						nsrl.Results.MarkDown = ""
						nsrlJSON, err := json.Marshal(nsrl)
						utils.Assert(err)
						if c.Bool("post") {
							request := gorequest.New()
							if c.Bool("proxy") {
								request = gorequest.New().Proxy(os.Getenv("MALICE_PROXY"))
							}
							request.Post(os.Getenv("MALICE_ENDPOINT")).
								Set("X-Malice-ID", utils.Getopt("MALICE_SCANID", hash)).
								Send(string(nsrlJSON)).
								End(printStatus)

							return nil
						}
						fmt.Println(string(nsrlJSON))
					}
				} else {
					log.Fatal(fmt.Errorf("please supply a SHA1 hash to query NSRL with"))
				}
				return nil
			},
		},
	}

	err := app.Run(os.Args)
	utils.Assert(err)
}
<|endoftext|>


<commit_before>package storage

import (
	"io/ioutil"
	"os"
	"testing"

	"github.com/endophage/gotuf/data"
	"github.com/jinzhu/gorm"
	_ "github.com/mattn/go-sqlite3"
	"github.com/stretchr/testify/assert"
)

// SampleTUF returns a sample TUFFile with the given Version (ID will have
// to be set independently)
func SampleTUF(version int) TUFFile {
	return TUFFile{
		Gun:     "testGUN",
		Role:    "root",
		Version: version,
		Data:    []byte("1"),
	}
}

func SampleUpdate(version int) MetaUpdate {
	return MetaUpdate{
		Role:    "root",
		Version: version,
		Data:    []byte("1"),
	}
}

// SetUpSQLite creates a sqlite database for testing
func SetUpSQLite(dbDir string, t *testing.T, createTable bool) (*gorm.DB, *SQLStorage) {
	dbStore, err := NewSQLStorage("sqlite3", dbDir+"test_db")
	assert.NoError(t, err)

	if createTable {
		// Create the DB tables
		err = CreateTUFTable(dbStore.DB)
		assert.NoError(t, err)

		err = CreateTimestampTable(dbStore.DB)
		assert.NoError(t, err)

		// verify that the tables are empty
		var count int
		for _, model := range [2]interface{}{&TUFFile{}, &TimestampKey{}} {
			query := dbStore.DB.Model(model).Count(&count)
			assert.NoError(t, query.Error)
			assert.Equal(t, 0, count)
		}
	}

	return &dbStore.DB, dbStore
}

// TestSQLUpdateCurrent asserts that UpdateCurrent will add a new TUF file
// if no previous version existed.
func TestSQLUpdateCurrentNew(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	// Adding a new TUF file should succeed
	err = dbStore.UpdateCurrent("testGUN", SampleUpdate(0))
	assert.NoError(t, err, "Creating a row in an empty DB failed.")

	// There should just be one row
	var rows []TUFFile
	query := gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)

	expected := SampleTUF(0)
	expected.ID = 1
	assert.Equal(t, []TUFFile{expected}, rows)
}

// TestSQLUpdateCurrentNewVersion asserts that UpdateCurrent will add a
// new (higher) version of an existing TUF file
func TestSQLUpdateCurrentNewVersion(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	// insert row
	oldVersion := SampleTUF(0)
	query := gormDB.Create(&oldVersion)
	assert.NoError(t, query.Error, "Creating a row in an empty DB failed.")

	// UpdateCurrent with a newer version should succeed
	update := SampleUpdate(2)
	err = dbStore.UpdateCurrent("testGUN", update)
	assert.NoError(t, err, "Creating a row in an empty DB failed.")

	// There should just be one row
	var rows []TUFFile
	query = gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)

	oldVersion.Model = gorm.Model{ID: 1}
	expected := SampleTUF(2)
	expected.Model = gorm.Model{ID: 2}
	assert.Equal(t, []TUFFile{oldVersion, expected}, rows)
}

// TestSQLUpdateCurrentOldVersionError asserts that an error is raised if
// trying to update to an older version of a TUF file.
func TestSQLUpdateCurrentOldVersionError(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	// insert row
	newVersion := SampleTUF(3)
	query := gormDB.Create(&newVersion)
	assert.NoError(t, query.Error, "Creating a row in an empty DB failed.")

	// UpdateCurrent should fail due to the version being lower than the
	// previous row
	err = dbStore.UpdateCurrent("testGUN", SampleUpdate(0))
	assert.Error(t, err, "Error should not be nil")
	assert.IsType(t, &ErrOldVersion{}, err,
		"Expected ErrOldVersion error type, got: %v", err)

	// There should just be one row
	var rows []TUFFile
	query = gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)

	newVersion.Model = gorm.Model{ID: 1}
	assert.Equal(t, []TUFFile{newVersion}, rows)

	dbStore.DB.Close()
}

// TestSQLUpdateMany asserts that inserting multiple updates succeeds if the
// updates do not conflict with each.
func TestSQLUpdateMany(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.UpdateMany("testGUN", []MetaUpdate{
		SampleUpdate(0),
		{
			Role:    "targets",
			Version: 1,
			Data:    []byte("2"),
		},
		SampleUpdate(2),
	})
	assert.NoError(t, err, "UpdateMany errored unexpectedly: %v", err)

	gorm1 := SampleTUF(0)
	gorm1.ID = 1
	gorm2 := TUFFile{
		Model: gorm.Model{ID: 2}, Gun: "testGUN", Role: "targets",
		Version: 1, Data: []byte("2")}
	gorm3 := SampleTUF(2)
	gorm3.ID = 3
	expected := []TUFFile{gorm1, gorm2, gorm3}

	var rows []TUFFile
	query := gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)
	assert.Equal(t, expected, rows)

	dbStore.DB.Close()
}

// TestSQLUpdateManyVersionOrder asserts that inserting updates with
// non-monotonic versions still succeeds.
func TestSQLUpdateManyVersionOrder(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.UpdateMany(
		"testGUN", []MetaUpdate{SampleUpdate(2), SampleUpdate(0)})
	assert.NoError(t, err)

	// the whole transaction should have rolled back, so there should be
	// no entries.
	gorm1 := SampleTUF(2)
	gorm1.ID = 1
	gorm2 := SampleTUF(0)
	gorm2.ID = 2

	var rows []TUFFile
	query := gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)
	assert.Equal(t, []TUFFile{gorm1, gorm2}, rows)

	dbStore.DB.Close()
}

// TestSQLUpdateManyDuplicateRollback asserts that inserting duplicate
// updates fails.
func TestSQLUpdateManyDuplicateRollback(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	update := SampleUpdate(0)
	err = dbStore.UpdateMany("testGUN", []MetaUpdate{update, update})
	assert.Error(
		t, err, "There should be an error updating the same data twice.")
	assert.IsType(t, &ErrOldVersion{}, err,
		"UpdateMany returned wrong error type")

	// the whole transaction should have rolled back, so there should be
	// no entries.
	var count int
	query := gormDB.Model(&TUFFile{}).Count(&count)
	assert.NoError(t, query.Error)
	assert.Equal(t, 0, count)

	dbStore.DB.Close()
}

func TestSQLGetCurrent(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	byt, err := dbStore.GetCurrent("testGUN", "root")
	assert.Nil(t, byt)
	assert.Error(t, err, "There should be an error Getting an empty table")
	assert.IsType(t, &ErrNotFound{}, err, "Should get a not found error")

	tuf := SampleTUF(0)
	query := gormDB.Create(&tuf)
	assert.NoError(t, query.Error, "Creating a row in an empty DB failed.")

	byt, err = dbStore.GetCurrent("testGUN", "root")
	assert.NoError(t, err, "There should not be any errors getting.")
	assert.Equal(t, []byte("1"), byt, "Returned data was incorrect")

	dbStore.DB.Close()
}

func TestSQLDelete(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	tuf := SampleTUF(0)
	query := gormDB.Create(&tuf)
	assert.NoError(t, query.Error, "Creating a row in an empty DB failed.")

	err = dbStore.Delete("testGUN")
	assert.NoError(t, err, "There should not be any errors deleting.")

	// verify deletion
	var count int
	query = gormDB.Model(&TUFFile{}).Count(&count)
	assert.NoError(t, query.Error)
	assert.Equal(t, 0, count)

	dbStore.DB.Close()
}

func TestSQLGetTimestampKeyNoKey(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	cipher, public, err := dbStore.GetTimestampKey("testGUN")
	assert.Equal(t, data.KeyAlgorithm(""), cipher)
	assert.Nil(t, public)
	assert.IsType(t, &ErrNoKey{}, err,
		"Expected ErrNoKey from GetTimestampKey")

	query := gormDB.Create(&TimestampKey{
		Gun:    "testGUN",
		Cipher: "testCipher",
		Public: []byte("1"),
	})
	assert.NoError(
		t, query.Error, "Inserting timestamp into empty DB should succeed")

	cipher, public, err = dbStore.GetTimestampKey("testGUN")
	assert.Equal(t, data.KeyAlgorithm("testCipher"), cipher,
		"Returned cipher was incorrect")
	assert.Equal(t, []byte("1"), public, "Returned pubkey was incorrect")
}

func TestSQLSetTimestampKeyExists(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.SetTimestampKey("testGUN", "testCipher", []byte("1"))
	assert.NoError(t, err, "Inserting timestamp into empty DB should succeed")

	err = dbStore.SetTimestampKey("testGUN", "testCipher", []byte("1"))
	assert.Error(t, err)
	assert.IsType(t, &ErrTimestampKeyExists{}, err,
		"Expected ErrTimestampKeyExists from SetTimestampKey")

	var rows []TimestampKey
	query := gormDB.Select("ID, Gun, Cipher, Public").Find(&rows)
	assert.NoError(t, query.Error)

	expected := TimestampKey{Gun: "testGUN", Cipher: "testCipher",
		Public: []byte("1")}
	expected.Model = gorm.Model{ID: 1}

	assert.Equal(t, []TimestampKey{expected}, rows)

	dbStore.DB.Close()
}

// TestDBCheckHealthTableMissing asserts that the health check fails if one or
// both the tables are missing.
func TestDBCheckHealthTableMissing(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	_, dbStore := SetUpSQLite(tempBaseDir, t, false)
	defer os.RemoveAll(tempBaseDir)

	// No tables, health check fails
	err = dbStore.CheckHealth()
	assert.Error(t, err, "Cannot access table:")

	// only one table existing causes health check to fail
	CreateTUFTable(dbStore.DB)
	err = dbStore.CheckHealth()
	assert.Error(t, err, "Cannot access table:")
	dbStore.DropTable(&TUFFile{})

	CreateTimestampTable(dbStore.DB)
	err = dbStore.CheckHealth()
	assert.Error(t, err, "Cannot access table:")
}

// TestDBCheckHealthDBCOnnection asserts that if the DB is not connectable, the
// health check fails.
func TestDBCheckHealthDBConnectionFail(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	_, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.Close()
	assert.NoError(t, err)

	err = dbStore.CheckHealth()
	assert.Error(t, err, "Cannot access table:")
}

// TestDBCheckHealthSuceeds asserts that if the DB is connectable and both
// tables exist, the health check succeeds.
func TestDBCheckHealthSucceeds(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	_, dbStore := SetUpSQLite(tempBaseDir, t, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.CheckHealth()
	assert.NoError(t, err)
}
<commit_msg>Move the testing.T module to be the first argument to helper functions.<commit_after>package storage

import (
	"io/ioutil"
	"os"
	"testing"

	"github.com/endophage/gotuf/data"
	"github.com/jinzhu/gorm"
	_ "github.com/mattn/go-sqlite3"
	"github.com/stretchr/testify/assert"
)

// SampleTUF returns a sample TUFFile with the given Version (ID will have
// to be set independently)
func SampleTUF(version int) TUFFile {
	return TUFFile{
		Gun:     "testGUN",
		Role:    "root",
		Version: version,
		Data:    []byte("1"),
	}
}

func SampleUpdate(version int) MetaUpdate {
	return MetaUpdate{
		Role:    "root",
		Version: version,
		Data:    []byte("1"),
	}
}

// SetUpSQLite creates a sqlite database for testing
func SetUpSQLite(t *testing.T, dbDir string, createTable bool) (*gorm.DB, *SQLStorage) {
	dbStore, err := NewSQLStorage("sqlite3", dbDir+"test_db")
	assert.NoError(t, err)

	if createTable {
		// Create the DB tables
		err = CreateTUFTable(dbStore.DB)
		assert.NoError(t, err)

		err = CreateTimestampTable(dbStore.DB)
		assert.NoError(t, err)

		// verify that the tables are empty
		var count int
		for _, model := range [2]interface{}{&TUFFile{}, &TimestampKey{}} {
			query := dbStore.DB.Model(model).Count(&count)
			assert.NoError(t, query.Error)
			assert.Equal(t, 0, count)
		}
	}

	return &dbStore.DB, dbStore
}

// TestSQLUpdateCurrent asserts that UpdateCurrent will add a new TUF file
// if no previous version existed.
func TestSQLUpdateCurrentNew(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	// Adding a new TUF file should succeed
	err = dbStore.UpdateCurrent("testGUN", SampleUpdate(0))
	assert.NoError(t, err, "Creating a row in an empty DB failed.")

	// There should just be one row
	var rows []TUFFile
	query := gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)

	expected := SampleTUF(0)
	expected.ID = 1
	assert.Equal(t, []TUFFile{expected}, rows)
}

// TestSQLUpdateCurrentNewVersion asserts that UpdateCurrent will add a
// new (higher) version of an existing TUF file
func TestSQLUpdateCurrentNewVersion(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	// insert row
	oldVersion := SampleTUF(0)
	query := gormDB.Create(&oldVersion)
	assert.NoError(t, query.Error, "Creating a row in an empty DB failed.")

	// UpdateCurrent with a newer version should succeed
	update := SampleUpdate(2)
	err = dbStore.UpdateCurrent("testGUN", update)
	assert.NoError(t, err, "Creating a row in an empty DB failed.")

	// There should just be one row
	var rows []TUFFile
	query = gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)

	oldVersion.Model = gorm.Model{ID: 1}
	expected := SampleTUF(2)
	expected.Model = gorm.Model{ID: 2}
	assert.Equal(t, []TUFFile{oldVersion, expected}, rows)
}

// TestSQLUpdateCurrentOldVersionError asserts that an error is raised if
// trying to update to an older version of a TUF file.
func TestSQLUpdateCurrentOldVersionError(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	// insert row
	newVersion := SampleTUF(3)
	query := gormDB.Create(&newVersion)
	assert.NoError(t, query.Error, "Creating a row in an empty DB failed.")

	// UpdateCurrent should fail due to the version being lower than the
	// previous row
	err = dbStore.UpdateCurrent("testGUN", SampleUpdate(0))
	assert.Error(t, err, "Error should not be nil")
	assert.IsType(t, &ErrOldVersion{}, err,
		"Expected ErrOldVersion error type, got: %v", err)

	// There should just be one row
	var rows []TUFFile
	query = gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)

	newVersion.Model = gorm.Model{ID: 1}
	assert.Equal(t, []TUFFile{newVersion}, rows)

	dbStore.DB.Close()
}

// TestSQLUpdateMany asserts that inserting multiple updates succeeds if the
// updates do not conflict with each.
func TestSQLUpdateMany(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.UpdateMany("testGUN", []MetaUpdate{
		SampleUpdate(0),
		{
			Role:    "targets",
			Version: 1,
			Data:    []byte("2"),
		},
		SampleUpdate(2),
	})
	assert.NoError(t, err, "UpdateMany errored unexpectedly: %v", err)

	gorm1 := SampleTUF(0)
	gorm1.ID = 1
	gorm2 := TUFFile{
		Model: gorm.Model{ID: 2}, Gun: "testGUN", Role: "targets",
		Version: 1, Data: []byte("2")}
	gorm3 := SampleTUF(2)
	gorm3.ID = 3
	expected := []TUFFile{gorm1, gorm2, gorm3}

	var rows []TUFFile
	query := gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)
	assert.Equal(t, expected, rows)

	dbStore.DB.Close()
}

// TestSQLUpdateManyVersionOrder asserts that inserting updates with
// non-monotonic versions still succeeds.
func TestSQLUpdateManyVersionOrder(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.UpdateMany(
		"testGUN", []MetaUpdate{SampleUpdate(2), SampleUpdate(0)})
	assert.NoError(t, err)

	// the whole transaction should have rolled back, so there should be
	// no entries.
	gorm1 := SampleTUF(2)
	gorm1.ID = 1
	gorm2 := SampleTUF(0)
	gorm2.ID = 2

	var rows []TUFFile
	query := gormDB.Select("ID, Gun, Role, Version, Data").Find(&rows)
	assert.NoError(t, query.Error)
	assert.Equal(t, []TUFFile{gorm1, gorm2}, rows)

	dbStore.DB.Close()
}

// TestSQLUpdateManyDuplicateRollback asserts that inserting duplicate
// updates fails.
func TestSQLUpdateManyDuplicateRollback(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	update := SampleUpdate(0)
	err = dbStore.UpdateMany("testGUN", []MetaUpdate{update, update})
	assert.Error(
		t, err, "There should be an error updating the same data twice.")
	assert.IsType(t, &ErrOldVersion{}, err,
		"UpdateMany returned wrong error type")

	// the whole transaction should have rolled back, so there should be
	// no entries.
	var count int
	query := gormDB.Model(&TUFFile{}).Count(&count)
	assert.NoError(t, query.Error)
	assert.Equal(t, 0, count)

	dbStore.DB.Close()
}

func TestSQLGetCurrent(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	byt, err := dbStore.GetCurrent("testGUN", "root")
	assert.Nil(t, byt)
	assert.Error(t, err, "There should be an error Getting an empty table")
	assert.IsType(t, &ErrNotFound{}, err, "Should get a not found error")

	tuf := SampleTUF(0)
	query := gormDB.Create(&tuf)
	assert.NoError(t, query.Error, "Creating a row in an empty DB failed.")

	byt, err = dbStore.GetCurrent("testGUN", "root")
	assert.NoError(t, err, "There should not be any errors getting.")
	assert.Equal(t, []byte("1"), byt, "Returned data was incorrect")

	dbStore.DB.Close()
}

func TestSQLDelete(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	tuf := SampleTUF(0)
	query := gormDB.Create(&tuf)
	assert.NoError(t, query.Error, "Creating a row in an empty DB failed.")

	err = dbStore.Delete("testGUN")
	assert.NoError(t, err, "There should not be any errors deleting.")

	// verify deletion
	var count int
	query = gormDB.Model(&TUFFile{}).Count(&count)
	assert.NoError(t, query.Error)
	assert.Equal(t, 0, count)

	dbStore.DB.Close()
}

func TestSQLGetTimestampKeyNoKey(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	cipher, public, err := dbStore.GetTimestampKey("testGUN")
	assert.Equal(t, data.KeyAlgorithm(""), cipher)
	assert.Nil(t, public)
	assert.IsType(t, &ErrNoKey{}, err,
		"Expected ErrNoKey from GetTimestampKey")

	query := gormDB.Create(&TimestampKey{
		Gun:    "testGUN",
		Cipher: "testCipher",
		Public: []byte("1"),
	})
	assert.NoError(
		t, query.Error, "Inserting timestamp into empty DB should succeed")

	cipher, public, err = dbStore.GetTimestampKey("testGUN")
	assert.Equal(t, data.KeyAlgorithm("testCipher"), cipher,
		"Returned cipher was incorrect")
	assert.Equal(t, []byte("1"), public, "Returned pubkey was incorrect")
}

func TestSQLSetTimestampKeyExists(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	gormDB, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.SetTimestampKey("testGUN", "testCipher", []byte("1"))
	assert.NoError(t, err, "Inserting timestamp into empty DB should succeed")

	err = dbStore.SetTimestampKey("testGUN", "testCipher", []byte("1"))
	assert.Error(t, err)
	assert.IsType(t, &ErrTimestampKeyExists{}, err,
		"Expected ErrTimestampKeyExists from SetTimestampKey")

	var rows []TimestampKey
	query := gormDB.Select("ID, Gun, Cipher, Public").Find(&rows)
	assert.NoError(t, query.Error)

	expected := TimestampKey{Gun: "testGUN", Cipher: "testCipher",
		Public: []byte("1")}
	expected.Model = gorm.Model{ID: 1}

	assert.Equal(t, []TimestampKey{expected}, rows)

	dbStore.DB.Close()
}

// TestDBCheckHealthTableMissing asserts that the health check fails if one or
// both the tables are missing.
func TestDBCheckHealthTableMissing(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	_, dbStore := SetUpSQLite(t, tempBaseDir, false)
	defer os.RemoveAll(tempBaseDir)

	// No tables, health check fails
	err = dbStore.CheckHealth()
	assert.Error(t, err, "Cannot access table:")

	// only one table existing causes health check to fail
	CreateTUFTable(dbStore.DB)
	err = dbStore.CheckHealth()
	assert.Error(t, err, "Cannot access table:")
	dbStore.DropTable(&TUFFile{})

	CreateTimestampTable(dbStore.DB)
	err = dbStore.CheckHealth()
	assert.Error(t, err, "Cannot access table:")
}

// TestDBCheckHealthDBCOnnection asserts that if the DB is not connectable, the
// health check fails.
func TestDBCheckHealthDBConnectionFail(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	_, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.Close()
	assert.NoError(t, err)

	err = dbStore.CheckHealth()
	assert.Error(t, err, "Cannot access table:")
}

// TestDBCheckHealthSuceeds asserts that if the DB is connectable and both
// tables exist, the health check succeeds.
func TestDBCheckHealthSucceeds(t *testing.T) {
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	_, dbStore := SetUpSQLite(t, tempBaseDir, true)
	defer os.RemoveAll(tempBaseDir)

	err = dbStore.CheckHealth()
	assert.NoError(t, err)
}
<|endoftext|>


<commit_before>package dendrite

import (
	zmq "github.com/pebbe/zmq4"
	"log"
	"sync"
	"time"
)

type controlType int

const (
	workerShutdownReq controlType = iota
	workerShutdownAllowed
	workerShutdownDenied
	workerShutdownConfirm
	workerRegisterReq
	workerRegisterAllowed
	workerRegisterDenied
)

// ZeroMQ Transport implementation
type ZMQTransport struct {
	lock              *sync.Mutex
	minHandlers       int
	maxHandlers       int
	incrHandlers      int
	activeRequests    int
	control_c         chan *workerComm
	dealer_sock       *zmq.Socket
	router_sock       *zmq.Socket
	zmq_context       *zmq.Context
	workerIdleTimeout time.Duration
}

// Creates ZeroMQ transport
// Multiplexes incoming connections which are then processed in separate go routines (workers)
// Multiplexer spawns go routines as needed, but 10 worker routines are created on start
// Every request times out after provided timeout duration
func InitZMQTransport(hostname string, timeout time.Duration) (Transport, error) {
	// initialize ZMQ Context
	context, err := zmq.NewContext()
	if err != nil {
		return nil, err
	}

	// setup router and bind() to tcp address for clients to connect to
	router_sock, err := context.NewSocket(zmq.ROUTER)
	if err != nil {
		return nil, err
	}
	err = router_sock.Bind("tcp://" + hostname)
	if err != nil {
		return nil, err
	}

	// setup dealer for
	dealer_sock, err := context.NewSocket(zmq.DEALER)
	if err != nil {
		return nil, err
	}
	err = dealer_sock.Bind("inproc://dendrite-zmqdealer")
	if err != nil {
		return nil, err
	}
	poller := zmq.NewPoller()
	poller.Add(router_sock, zmq.POLLIN)
	poller.Add(dealer_sock, zmq.POLLIN)

	transport := &ZMQTransport{
		lock:              new(sync.Mutex),
		minHandlers:       10,
		maxHandlers:       1024,
		incrHandlers:      10,
		activeRequests:    0,
		workerIdleTimeout: 10 * time.Second,
		control_c:         make(chan *workerComm),
		dealer_sock:       dealer_sock,
		router_sock:       router_sock,
		zmq_context:       context,
	}
	// proxy messages between router and dealer
	go func() {
		for {
			sockets, _ := poller.Poll(-1)
			for _, socket := range sockets {
				switch s := socket.Socket; s {
				case router_sock:
					msg, err := s.RecvBytes(0)
					if err != nil {
						log.Println("ERR: TransportListener read bytes on router failed", err)
						continue
					}
					// forward to dealer
					_, err = dealer_sock.SendBytes(msg, 0)
					if err != nil {
						log.Println("ERR: TransportListener forward to dealer failed", err)
						continue
					}
					transport.activeRequests += 1
				case dealer_sock:
					msg, err := s.RecvBytes(0)
					if err != nil {
						log.Println("ERR: TransportListener read bytes on dealer failed", err)
						continue
					}
					// forward up to router
					_, err = router_sock.SendBytes(msg, 0)
					if err != nil {
						log.Println("ERR: TransportListener forward to router failed", err)
						continue
					}
					transport.activeRequests -= 1
				}
			}
		}
	}()

	// Scheduler goroutine keeps track of running workers
	// It spawns new ones if needed, and cancels ones that are idling
	go func() {
		sched_ticker := time.NewTicker(60 * time.Second)
		workers := make(map[*workerComm]bool)
		// fire up initial set of workers
		//for i := 0; i < transport.minHandlers; i++ {
		for i := 0; i < 15; i++ {
			//go handleZMQreq(transport)
		}
		for {
			select {
			case comm := <-transport.control_c:
				// worker sent something...
				msg := <-comm.worker_out
				switch {
				case msg == workerRegisterReq:
					if len(workers) == transport.maxHandlers {
						comm.worker_in <- workerRegisterDenied
						log.Println("ERR: TransportListener - max number of workers reached")
						continue
					}
					if _, ok := workers[comm]; ok {
						// worker already registered
						continue
					}
					comm.worker_in <- workerRegisterAllowed
					workers[comm] = true
					log.Println("INFO: TransportListener - registered new worker, total:", len(workers))
				case msg == workerShutdownReq:
					//log.Println("Got shutdown req")
					if len(workers) > transport.minHandlers {
						comm.worker_in <- workerShutdownAllowed
						for _ = range comm.worker_out {
							// wait until worker closes the channel
						}
						delete(workers, comm)
					} else {
						comm.worker_in <- workerShutdownDenied
					}
				}
			case <-sched_ticker.C:
				// check if requests are piling up and start more workers if that's the case
				if transport.activeRequests > 3*len(workers) {
					for i := 0; i < transport.incrHandlers; i++ {
						go handleZMQreq(transport)
					}
				}
			}
		}
	}()
	return transport, nil
}

type workerComm struct {
	worker_in  chan controlType // worker's input channel for two way communication with scheduler
	worker_out chan controlType // worker's output channel for two way communication with scheduler
	worker_ctl chan controlType // worker's control channel for communication with scheduler
}

func handleZMQreq(transport *ZMQTransport) {
	// setup REP socket
	rep_sock, err := transport.zmq_context.NewSocket(zmq.REP)
	if err != nil {
		log.Println("ERR: TransportListener worker failed to create REP socket", err)
		return
	}
	err = rep_sock.Connect("inproc://dendrite-zmqdealer")
	if err != nil {
		log.Println("ERR: TransportListener worker failed to connect to dealer", err)
		return
	}

	// setup communication channels with scheduler
	worker_in := make(chan controlType, 1)
	worker_out := make(chan controlType, 1)
	worker_ctl := make(chan controlType, 1)
	comm := &workerComm{
		worker_in:  worker_in,
		worker_out: worker_out,
		worker_ctl: worker_ctl,
	}
	// notify scheduler that we're up
	worker_out <- workerRegisterReq
	transport.control_c <- comm
	v := <-worker_in
	if v == workerRegisterDenied {
		return
	}

	// setup socket read channel
	read_c := make(chan []byte)
	poller := zmq.NewPoller()
	poller.Add(rep_sock, zmq.POLLIN)
	cancel_c := make(chan bool, 1)
	// read from socket and emit data, or stop if canceled
	go func() {
	MAINLOOP:
		for {
			// poll for 5 seconds, but then see if we should be canceled
			sockets, _ := poller.Poll(5 * time.Second)
			for _, socket := range sockets {
				msg, err := socket.Socket.RecvBytes(0)
				if err != nil {
					log.Println("ERR: TransportListener error while reading from REP, ", err)
					continue
				}
				// emit data
				read_c <- msg
			}
			// check for cancel request
			select {
			case <-cancel_c:
				break MAINLOOP
			default:
				break
			}
		}
	}()

	// read from socket and process request -- OR
	// shutdown if scheduler wants me to -- OR
	// request shutdown from scheduler if idling and exit if allowed
	ticker := time.NewTicker(transport.workerIdleTimeout)
	for {
		select {
		case data := <-read_c:
			// handle data
			log.Println("Got data", data)
			ticker.Stop()
			ticker = time.NewTicker(transport.workerIdleTimeout)

		case controlMsg := <-comm.worker_ctl:
			// got control message, probably shutdown instruction
			log.Println("Got control msg", controlMsg)
		case <-ticker.C:
			// we're idling, lets request shutdown
			comm.worker_out <- workerShutdownReq
			transport.control_c <- comm
			v := <-comm.worker_in
			if v == workerShutdownAllowed {
				log.Println("Shuting down now....")
				close(comm.worker_out)
				cancel_c <- true
				close(cancel_c)
				return
			}
		}
	}
}
<commit_msg>fixes initial worker pool<commit_after>package dendrite

import (
	zmq "github.com/pebbe/zmq4"
	"log"
	"sync"
	"time"
)

type controlType int

const (
	workerShutdownReq controlType = iota
	workerShutdownAllowed
	workerShutdownDenied
	workerShutdownConfirm
	workerRegisterReq
	workerRegisterAllowed
	workerRegisterDenied
)

// ZeroMQ Transport implementation
type ZMQTransport struct {
	lock              *sync.Mutex
	minHandlers       int
	maxHandlers       int
	incrHandlers      int
	activeRequests    int
	control_c         chan *workerComm
	dealer_sock       *zmq.Socket
	router_sock       *zmq.Socket
	zmq_context       *zmq.Context
	workerIdleTimeout time.Duration
}

// Creates ZeroMQ transport
// Multiplexes incoming connections which are then processed in separate go routines (workers)
// Multiplexer spawns go routines as needed, but 10 worker routines are created on start
// Every request times out after provided timeout duration
func InitZMQTransport(hostname string, timeout time.Duration) (Transport, error) {
	// initialize ZMQ Context
	context, err := zmq.NewContext()
	if err != nil {
		return nil, err
	}

	// setup router and bind() to tcp address for clients to connect to
	router_sock, err := context.NewSocket(zmq.ROUTER)
	if err != nil {
		return nil, err
	}
	err = router_sock.Bind("tcp://" + hostname)
	if err != nil {
		return nil, err
	}

	// setup dealer for
	dealer_sock, err := context.NewSocket(zmq.DEALER)
	if err != nil {
		return nil, err
	}
	err = dealer_sock.Bind("inproc://dendrite-zmqdealer")
	if err != nil {
		return nil, err
	}
	poller := zmq.NewPoller()
	poller.Add(router_sock, zmq.POLLIN)
	poller.Add(dealer_sock, zmq.POLLIN)

	transport := &ZMQTransport{
		lock:              new(sync.Mutex),
		minHandlers:       10,
		maxHandlers:       1024,
		incrHandlers:      10,
		activeRequests:    0,
		workerIdleTimeout: 10 * time.Second,
		control_c:         make(chan *workerComm),
		dealer_sock:       dealer_sock,
		router_sock:       router_sock,
		zmq_context:       context,
	}
	// proxy messages between router and dealer
	go func() {
		for {
			sockets, _ := poller.Poll(-1)
			for _, socket := range sockets {
				switch s := socket.Socket; s {
				case router_sock:
					msg, err := s.RecvBytes(0)
					if err != nil {
						log.Println("ERR: TransportListener read bytes on router failed", err)
						continue
					}
					// forward to dealer
					_, err = dealer_sock.SendBytes(msg, 0)
					if err != nil {
						log.Println("ERR: TransportListener forward to dealer failed", err)
						continue
					}
					transport.activeRequests += 1
				case dealer_sock:
					msg, err := s.RecvBytes(0)
					if err != nil {
						log.Println("ERR: TransportListener read bytes on dealer failed", err)
						continue
					}
					// forward up to router
					_, err = router_sock.SendBytes(msg, 0)
					if err != nil {
						log.Println("ERR: TransportListener forward to router failed", err)
						continue
					}
					transport.activeRequests -= 1
				}
			}
		}
	}()

	// Scheduler goroutine keeps track of running workers
	// It spawns new ones if needed, and cancels ones that are idling
	go func() {
		sched_ticker := time.NewTicker(60 * time.Second)
		workers := make(map[*workerComm]bool)
		// fire up initial set of workers
		for i := 0; i < transport.minHandlers; i++ {
			go handleZMQreq(transport)
		}
		for {
			select {
			case comm := <-transport.control_c:
				// worker sent something...
				msg := <-comm.worker_out
				switch {
				case msg == workerRegisterReq:
					if len(workers) == transport.maxHandlers {
						comm.worker_in <- workerRegisterDenied
						log.Println("ERR: TransportListener - max number of workers reached")
						continue
					}
					if _, ok := workers[comm]; ok {
						// worker already registered
						continue
					}
					comm.worker_in <- workerRegisterAllowed
					workers[comm] = true
					log.Println("INFO: TransportListener - registered new worker, total:", len(workers))
				case msg == workerShutdownReq:
					//log.Println("Got shutdown req")
					if len(workers) > transport.minHandlers {
						comm.worker_in <- workerShutdownAllowed
						for _ = range comm.worker_out {
							// wait until worker closes the channel
						}
						delete(workers, comm)
					} else {
						comm.worker_in <- workerShutdownDenied
					}
				}
			case <-sched_ticker.C:
				// check if requests are piling up and start more workers if that's the case
				if transport.activeRequests > 3*len(workers) {
					for i := 0; i < transport.incrHandlers; i++ {
						go handleZMQreq(transport)
					}
				}
			}
		}
	}()
	return transport, nil
}

type workerComm struct {
	worker_in  chan controlType // worker's input channel for two way communication with scheduler
	worker_out chan controlType // worker's output channel for two way communication with scheduler
	worker_ctl chan controlType // worker's control channel for communication with scheduler
}

func handleZMQreq(transport *ZMQTransport) {
	// setup REP socket
	rep_sock, err := transport.zmq_context.NewSocket(zmq.REP)
	if err != nil {
		log.Println("ERR: TransportListener worker failed to create REP socket", err)
		return
	}
	err = rep_sock.Connect("inproc://dendrite-zmqdealer")
	if err != nil {
		log.Println("ERR: TransportListener worker failed to connect to dealer", err)
		return
	}

	// setup communication channels with scheduler
	worker_in := make(chan controlType, 1)
	worker_out := make(chan controlType, 1)
	worker_ctl := make(chan controlType, 1)
	comm := &workerComm{
		worker_in:  worker_in,
		worker_out: worker_out,
		worker_ctl: worker_ctl,
	}
	// notify scheduler that we're up
	worker_out <- workerRegisterReq
	transport.control_c <- comm
	v := <-worker_in
	if v == workerRegisterDenied {
		return
	}

	// setup socket read channel
	read_c := make(chan []byte)
	poller := zmq.NewPoller()
	poller.Add(rep_sock, zmq.POLLIN)
	cancel_c := make(chan bool, 1)
	// read from socket and emit data, or stop if canceled
	go func() {
	MAINLOOP:
		for {
			// poll for 5 seconds, but then see if we should be canceled
			sockets, _ := poller.Poll(5 * time.Second)
			for _, socket := range sockets {
				msg, err := socket.Socket.RecvBytes(0)
				if err != nil {
					log.Println("ERR: TransportListener error while reading from REP, ", err)
					continue
				}
				// emit data
				read_c <- msg
			}
			// check for cancel request
			select {
			case <-cancel_c:
				break MAINLOOP
			default:
				break
			}
		}
	}()

	// read from socket and process request -- OR
	// shutdown if scheduler wants me to -- OR
	// request shutdown from scheduler if idling and exit if allowed
	ticker := time.NewTicker(transport.workerIdleTimeout)
	for {
		select {
		case data := <-read_c:
			// handle data
			log.Println("Got data", data)
			ticker.Stop()
			ticker = time.NewTicker(transport.workerIdleTimeout)

		case controlMsg := <-comm.worker_ctl:
			// got control message, probably shutdown instruction
			log.Println("Got control msg", controlMsg)
		case <-ticker.C:
			// we're idling, lets request shutdown
			comm.worker_out <- workerShutdownReq
			transport.control_c <- comm
			v := <-comm.worker_in
			if v == workerShutdownAllowed {
				log.Println("Shuting down now....")
				close(comm.worker_out)
				cancel_c <- true
				close(cancel_c)
				return
			}
		}
	}
}
<|endoftext|>


<commit_before>package sqlite3

import "fmt"

// PageKind describes what kind of page is.
type PageKind byte

/*
An index B-Tree internal node
An index B-Tree leaf node
A table B-Tree internal node
A table B-Tree leaf node
An overflow page
A freelist page
A pointer map page
The locking page
*/

const (
	intKeyKind   PageKind = 0x01
	zeroDataKind          = 0x02
	leafDataKind          = 0x04
	leafKind              = 0x08

	BTreeInteriorIndexKind = intKeyKind
	BTreeInteriorTableKind = leafDataKind | intKeyKind
	BTreeLeafIndexKind     = zeroDataKind | leafKind
	BTreeLeafTableKind     = leafDataKind | intKeyKind | leafKind

	pkLockByte
	pkFreelistTrunk
	pkFreelistLeaf
	pkPayloadOverflow
	pkPointerMap
)

func (pk PageKind) String() string {
	switch pk {
	case BTreeInteriorIndexKind:
		return "BTreeInteriorIndex"
	case BTreeInteriorTableKind:
		return "BTreeInteriorTable"
	case BTreeLeafIndexKind:
		return "BTreeLeafIndex"
	case BTreeLeafTableKind:
		return "BTreeLeafTable"
	}
	panic("unreachable")
}

func newPage(i int, pb pageBuffer) (Page, error) {
	pk := PageKind(pb.buf[0])
	switch pk {
	case BTreeInteriorIndexKind:
		panic("not implemented")
	case BTreeInteriorTableKind:
		panic("not implemented")
	case BTreeLeafIndexKind:
		panic("not implemented")
	case BTreeLeafTableKind:
		return newBtreeLeafTable(i, pb)

	}

	panic(fmt.Errorf("invalid PageKind value (%d)", pk))
}

// Page is a page on disk
type Page interface {
	ID() int        // page number. starts at 1. id=0 means no page.
	Kind() PageKind // the kind of page
	Size() int      // the size of the page in bytes
}

// pageBuffer
type pageBuffer struct {
	pos int
	buf []byte
}

func (p *pageBuffer) Seek(offset int64, whence int) (ret int64, err error) {
	switch whence {
	case 0:
		offset := int(offset)
		if offset > len(p.buf) {
			return 0, fmt.Errorf("sqlite: offset too big (%d)", offset)
		}
		p.pos = offset
	case 1:
		offset := int(offset)
		pos := p.pos + offset
		if pos > len(p.buf) {
			return 0, fmt.Errorf("sqlite: offset too big (%d)", offset)
		}
		p.pos = pos
	case 2:
		offset := int(offset)
		pos := len(p.buf) - offset
		if pos < 0 {
			return 0, fmt.Errorf("sqlite: offset too big (%d)", offset)
		}
		p.pos = pos
	}
	return int64(p.pos), nil
}

func (p *pageBuffer) Pos() int {
	return p.pos
}

func (p *pageBuffer) Bytes() []byte {
	return p.buf[p.pos:]
}

func (p *pageBuffer) Decode(ptr interface{}) error {
	n, err := unmarshal(p.buf[p.pos:], ptr)
	if err != nil {
		return err
	}
	p.pos += int(n)
	return err
}

func (p *pageBuffer) Read(data []byte) (int, error) {
	n := copy(data, p.buf[p.pos:p.pos+len(data)])
	if n != len(data) {
		return n, fmt.Errorf("error. read too few bytes: %d. want %d", n, len(data))
	}
	p.pos += n
	return n, nil
}

func (p *pageBuffer) Uvarint() (uint64, int) {
	v, n := uvarint(p.Bytes())
	if n <= 0 {
		return v, n
	}
	p.pos += int(n)
	return v, n
}
<commit_msg>page: gofmt<commit_after>package sqlite3

import "fmt"

// PageKind describes what kind of page is.
type PageKind byte

/*
An index B-Tree internal node
An index B-Tree leaf node
A table B-Tree internal node
A table B-Tree leaf node
An overflow page
A freelist page
A pointer map page
The locking page
*/

const (
	intKeyKind   PageKind = 0x01
	zeroDataKind          = 0x02
	leafDataKind          = 0x04
	leafKind              = 0x08

	BTreeInteriorIndexKind = intKeyKind
	BTreeInteriorTableKind = leafDataKind | intKeyKind
	BTreeLeafIndexKind     = zeroDataKind | leafKind
	BTreeLeafTableKind     = leafDataKind | intKeyKind | leafKind

	pkLockByte
	pkFreelistTrunk
	pkFreelistLeaf
	pkPayloadOverflow
	pkPointerMap
)

func (pk PageKind) String() string {
	switch pk {
	case BTreeInteriorIndexKind:
		return "BTreeInteriorIndex"
	case BTreeInteriorTableKind:
		return "BTreeInteriorTable"
	case BTreeLeafIndexKind:
		return "BTreeLeafIndex"
	case BTreeLeafTableKind:
		return "BTreeLeafTable"
	}
	panic("unreachable")
}

/*
func newPage(i int, pb pageBuffer) (Page, error) {
	pk := PageKind(pb.buf[0])
	switch pk {
	case BTreeInteriorIndexKind:
		panic("not implemented")
	case BTreeInteriorTableKind:
		panic("not implemented")
	case BTreeLeafIndexKind:
		panic("not implemented")
	case BTreeLeafTableKind:
		return newBtreeLeafTable(i, pb)

	}

	panic(fmt.Errorf("invalid PageKind value (%d)", pk))
}
*/

// Page is a page on disk
type Page interface {
	ID() int        // page number. starts at 1. id=0 means no page.
	Kind() PageKind // the kind of page
	Size() int      // the size of the page in bytes
}

// pageBuffer
type pageBuffer struct {
	pos int
	buf []byte
}

func (p *pageBuffer) Seek(offset int64, whence int) (ret int64, err error) {
	switch whence {
	case 0:
		offset := int(offset)
		if offset > len(p.buf) {
			return 0, fmt.Errorf("sqlite: offset too big (%d)", offset)
		}
		p.pos = offset
	case 1:
		offset := int(offset)
		pos := p.pos + offset
		if pos > len(p.buf) {
			return 0, fmt.Errorf("sqlite: offset too big (%d)", offset)
		}
		p.pos = pos
	case 2:
		offset := int(offset)
		pos := len(p.buf) - offset
		if pos < 0 {
			return 0, fmt.Errorf("sqlite: offset too big (%d)", offset)
		}
		p.pos = pos
	}
	return int64(p.pos), nil
}

func (p *pageBuffer) Pos() int {
	return p.pos
}

func (p *pageBuffer) Bytes() []byte {
	return p.buf[p.pos:]
}

func (p *pageBuffer) Decode(ptr interface{}) error {
	n, err := unmarshal(p.buf[p.pos:], ptr)
	if err != nil {
		return err
	}
	p.pos += int(n)
	return err
}

func (p *pageBuffer) Read(data []byte) (int, error) {
	n := copy(data, p.buf[p.pos:p.pos+len(data)])
	if n != len(data) {
		return n, fmt.Errorf("error. read too few bytes: %d. want %d", n, len(data))
	}
	p.pos += n
	return n, nil
}

func (p *pageBuffer) Uvarint() (uint64, int) {
	v, n := uvarint(p.Bytes())
	if n <= 0 {
		return v, n
	}
	p.pos += int(n)
	return v, n
}
<|endoftext|>


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

package migrationtarget

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

	"github.com/juju/errors"
	"github.com/juju/httprequest"
	"github.com/juju/version"
	"gopkg.in/juju/charm.v6-unstable"
	"gopkg.in/juju/names.v2"

	"github.com/juju/juju/api/base"
	"github.com/juju/juju/apiserver/params"
	coremigration "github.com/juju/juju/core/migration"
	"github.com/juju/juju/tools"
	jujuversion "github.com/juju/juju/version"
)

type httpClient interface {
	HTTPClient() (*httprequest.Client, error)
}

// NewClient returns a new Client based on an existing API connection.
func NewClient(caller base.APICaller) *Client {
	return &Client{
		caller:            base.NewFacadeCaller(caller, "MigrationTarget"),
		httpClientFactory: caller,
	}
}

// Client is the client-side API for the MigrationTarget facade. It is
// used by the migrationmaster worker when talking to the target
// controller during a migration.
type Client struct {
	caller            base.FacadeCaller
	httpClientFactory httpClient
}

func (c *Client) Prechecks(model coremigration.ModelInfo) error {
	args := params.MigrationModelInfo{
		UUID:         model.UUID,
		Name:         model.Name,
		OwnerTag:     model.Owner.String(),
		AgentVersion: model.AgentVersion,
	}
	return c.caller.FacadeCall("Prechecks", args, nil)
}

// Import takes a serialized model and imports it into the target
// controller.
func (c *Client) Import(bytes []byte) error {
	serialized := params.SerializedModel{Bytes: bytes}
	return c.caller.FacadeCall("Import", serialized, nil)
}

// Abort removes all data relating to a previously imported model.
func (c *Client) Abort(modelUUID string) error {
	args := params.ModelArgs{ModelTag: names.NewModelTag(modelUUID).String()}
	return c.caller.FacadeCall("Abort", args, nil)
}

// Activate marks a migrated model as being ready to use.
func (c *Client) Activate(modelUUID string) error {
	args := params.ModelArgs{ModelTag: names.NewModelTag(modelUUID).String()}
	return c.caller.FacadeCall("Activate", args, nil)
}

// UploadCharm sends the content to the API server using an HTTP post in order
// to add the charm binary to the model specified.
func (c *Client) UploadCharm(modelUUID string, curl *charm.URL, content io.ReadSeeker) (*charm.URL, error) {
	args := url.Values{}
	args.Add("schema", curl.Schema)
	args.Add("user", curl.User)
	args.Add("series", curl.Series)
	args.Add("revision", strconv.Itoa(curl.Revision))
	apiURI := url.URL{Path: "/migrate/charms", RawQuery: args.Encode()}

	contentType := "application/zip"
	var resp params.CharmsResponse
	if err := c.httpPost(modelUUID, content, apiURI.String(), contentType, &resp); err != nil {
		return nil, errors.Trace(err)
	}

	curl, err := charm.ParseURL(resp.CharmURL)
	if err != nil {
		return nil, errors.Annotatef(err, "bad charm URL in response")
	}
	return curl, nil
}

// UploadTools uploads tools at the specified location to the API server over HTTPS
// for the specified model.
func (c *Client) UploadTools(modelUUID string, r io.ReadSeeker, vers version.Binary, additionalSeries ...string) (tools.List, error) {
	endpoint := fmt.Sprintf("/migrate/tools?binaryVersion=%s&series=%s", vers, strings.Join(additionalSeries, ","))
	contentType := "application/x-tar-gz"
	var resp params.ToolsResult
	if err := c.httpPost(modelUUID, r, endpoint, contentType, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return resp.ToolsList, nil
}

func (c *Client) httpPost(modelUUID string, content io.ReadSeeker, endpoint, contentType string, response interface{}) error {
	req, err := http.NewRequest("POST", endpoint, nil)
	if err != nil {
		return errors.Annotate(err, "cannot create upload request")
	}
	req.Header.Set("Content-Type", contentType)
	req.Header.Set(params.MigrationModelHTTPHeader, modelUUID)

	// The returned httpClient sets the base url to /model/<uuid> if it can.
	httpClient, err := c.httpClientFactory.HTTPClient()
	if err != nil {
		return errors.Trace(err)
	}

	if err := httpClient.Do(req, content, response); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// OpenLogTransferStream connects to the migration logtransfer
// endpoint on the target controller and returns a stream that JSON
// logs records can be fed into. The objects written should be params.LogRecords.
func (c *Client) OpenLogTransferStream(modelUUID string) (base.Stream, error) {
	attrs := url.Values{}
	attrs.Set("jujuclientversion", jujuversion.Current.String())
	headers := http.Header{}
	headers.Set(params.MigrationModelHTTPHeader, modelUUID)
	caller := c.caller.RawAPICaller()
	stream, err := caller.ConnectControllerStream("/migrate/logtransfer", attrs, headers)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return stream, nil
}

// LatestLogTime asks the target controller for the time of the latest
// log record it has seen. This can be used to make the log transfer
// restartable.
func (c *Client) LatestLogTime(modelUUID string) (time.Time, error) {
	var result time.Time
	args := params.ModelArgs{names.NewModelTag(modelUUID).String()}
	err := c.caller.FacadeCall("LatestLogTime", args, &result)
	if err != nil {
		return time.Time{}, errors.Trace(err)
	}
	return result, nil
}
<commit_msg>api/migrationtarget: Added UploadResource<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package migrationtarget

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

	"github.com/juju/errors"
	"github.com/juju/httprequest"
	"github.com/juju/version"
	"gopkg.in/juju/charm.v6-unstable"
	"gopkg.in/juju/charm.v6-unstable/resource"
	"gopkg.in/juju/names.v2"

	"github.com/juju/juju/api/base"
	"github.com/juju/juju/apiserver/params"
	coremigration "github.com/juju/juju/core/migration"
	"github.com/juju/juju/tools"
	jujuversion "github.com/juju/juju/version"
)

type httpClient interface {
	HTTPClient() (*httprequest.Client, error)
}

// NewClient returns a new Client based on an existing API connection.
func NewClient(caller base.APICaller) *Client {
	return &Client{
		caller:            base.NewFacadeCaller(caller, "MigrationTarget"),
		httpClientFactory: caller,
	}
}

// Client is the client-side API for the MigrationTarget facade. It is
// used by the migrationmaster worker when talking to the target
// controller during a migration.
type Client struct {
	caller            base.FacadeCaller
	httpClientFactory httpClient
}

func (c *Client) Prechecks(model coremigration.ModelInfo) error {
	args := params.MigrationModelInfo{
		UUID:         model.UUID,
		Name:         model.Name,
		OwnerTag:     model.Owner.String(),
		AgentVersion: model.AgentVersion,
	}
	return c.caller.FacadeCall("Prechecks", args, nil)
}

// Import takes a serialized model and imports it into the target
// controller.
func (c *Client) Import(bytes []byte) error {
	serialized := params.SerializedModel{Bytes: bytes}
	return c.caller.FacadeCall("Import", serialized, nil)
}

// Abort removes all data relating to a previously imported model.
func (c *Client) Abort(modelUUID string) error {
	args := params.ModelArgs{ModelTag: names.NewModelTag(modelUUID).String()}
	return c.caller.FacadeCall("Abort", args, nil)
}

// Activate marks a migrated model as being ready to use.
func (c *Client) Activate(modelUUID string) error {
	args := params.ModelArgs{ModelTag: names.NewModelTag(modelUUID).String()}
	return c.caller.FacadeCall("Activate", args, nil)
}

// UploadCharm sends the content to the API server using an HTTP post in order
// to add the charm binary to the model specified.
func (c *Client) UploadCharm(modelUUID string, curl *charm.URL, content io.ReadSeeker) (*charm.URL, error) {
	args := url.Values{}
	args.Add("schema", curl.Schema)
	args.Add("user", curl.User)
	args.Add("series", curl.Series)
	args.Add("revision", strconv.Itoa(curl.Revision))
	apiURI := url.URL{Path: "/migrate/charms", RawQuery: args.Encode()}

	contentType := "application/zip"
	var resp params.CharmsResponse
	if err := c.httpPost(modelUUID, content, apiURI.String(), contentType, &resp); err != nil {
		return nil, errors.Trace(err)
	}

	curl, err := charm.ParseURL(resp.CharmURL)
	if err != nil {
		return nil, errors.Annotatef(err, "bad charm URL in response")
	}
	return curl, nil
}

// UploadTools uploads tools at the specified location to the API server over HTTPS
// for the specified model.
func (c *Client) UploadTools(modelUUID string, r io.ReadSeeker, vers version.Binary, additionalSeries ...string) (tools.List, error) {
	endpoint := fmt.Sprintf("/migrate/tools?binaryVersion=%s&series=%s", vers, strings.Join(additionalSeries, ","))
	contentType := "application/x-tar-gz"
	var resp params.ToolsResult
	if err := c.httpPost(modelUUID, r, endpoint, contentType, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return resp.ToolsList, nil
}

// XXX needs tests
// XXX
func (c *Client) UploadResource(modelUUID string, res resource.Resource, r io.ReadSeeker) error {
	args := url.Values{}
	args.Add("name", res.Name)
	args.Add("type", res.Type.String())
	args.Add("path", res.Path)
	args.Add("description", res.Description)
	args.Add("origin", res.Origin.String())
	args.Add("revision", fmt.Sprintf("%d", res.Revision))
	args.Add("size", fmt.Sprintf("%d", res.Size))
	args.Add("fingerprint", res.Fingerprint.Hex())
	uri := "/migrate/resources?" + args.Encode()
	contentType := "application/octet-stream"
	err := c.httpPost(modelUUID, r, uri, contentType, nil)
	return errors.Trace(err)
}

func (c *Client) httpPost(modelUUID string, content io.ReadSeeker, endpoint, contentType string, response interface{}) error {
	req, err := http.NewRequest("POST", endpoint, nil)
	if err != nil {
		return errors.Annotate(err, "cannot create upload request")
	}
	req.Header.Set("Content-Type", contentType)
	req.Header.Set(params.MigrationModelHTTPHeader, modelUUID)

	// The returned httpClient sets the base url to /model/<uuid> if it can.
	httpClient, err := c.httpClientFactory.HTTPClient()
	if err != nil {
		return errors.Trace(err)
	}

	if err := httpClient.Do(req, content, response); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// OpenLogTransferStream connects to the migration logtransfer
// endpoint on the target controller and returns a stream that JSON
// logs records can be fed into. The objects written should be params.LogRecords.
func (c *Client) OpenLogTransferStream(modelUUID string) (base.Stream, error) {
	attrs := url.Values{}
	attrs.Set("jujuclientversion", jujuversion.Current.String())
	headers := http.Header{}
	headers.Set(params.MigrationModelHTTPHeader, modelUUID)
	caller := c.caller.RawAPICaller()
	stream, err := caller.ConnectControllerStream("/migrate/logtransfer", attrs, headers)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return stream, nil
}

// LatestLogTime asks the target controller for the time of the latest
// log record it has seen. This can be used to make the log transfer
// restartable.
func (c *Client) LatestLogTime(modelUUID string) (time.Time, error) {
	var result time.Time
	args := params.ModelArgs{names.NewModelTag(modelUUID).String()}
	err := c.caller.FacadeCall("LatestLogTime", args, &result)
	if err != nil {
		return time.Time{}, errors.Trace(err)
	}
	return result, nil
}
<|endoftext|>