<commit_before>package google

import (
	"fmt"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/serviceusage/v1"
	"log"
	"strings"
	"time"
)

var ignoredProjectServices = []string{"dataproc-control.googleapis.com", "source.googleapis.com", "stackdriverprovisioning.googleapis.com"}

// These services can only be enabled as a side-effect of enabling other services,
// so don't bother storing them in the config or using them for diffing.
var ignoredProjectServicesSet = golangSetFromStringSlice(ignoredProjectServices)

// Service Renames
// we expect when a service is renamed:
// - both service names will continue to be able to be set
// - setting one will effectively enable the other as a dependent
// - GET will return whichever service name is requested
// - LIST responses will not contain the old service name
// renames may be reverted, though, so we should canonicalise both ways until
// the old service is fully removed from the provider
//
// We handle service renames in the provider by pretending that we've read both
// the old and new service names from the API if we see either, and only setting
// the one(s) that existed in prior state in config (if any). If neither exists,
// we'll set the new service name in state.
// Additionally, in case of service rename rollbacks or unexpected early
// removals of services, if we fail to create or delete a service that's been
// renamed we'll retry using an alternate name.
// We try creation by the user-specified value followed by the other value.
// We try deletion by the old value followed by the new value.

// map from old -> new names of services that have been renamed
// these should be removed during major provider versions. comment here with
// "DEPRECATED FOR {{version}} next to entries slated for removal in {{version}}
// upon removal, we should disallow the old name from being used even if it's
// not gone from the underlying API yet
var renamedServices = map[string]string{
	"bigquery-json.googleapis.com": "bigquery.googleapis.com", // DEPRECATED FOR 3.0.0
}

// renamedServices in reverse (new -> old)
var renamedServicesByNewServiceNames = reverseStringMap(renamedServices)

// renamedServices expressed as both old -> new and new -> old
var renamedServicesByOldAndNewServiceNames = mergeStringMaps(renamedServices, renamedServicesByNewServiceNames)

func resourceGoogleProjectService() *schema.Resource {
	return &schema.Resource{
		Create: resourceGoogleProjectServiceCreate,
		Read:   resourceGoogleProjectServiceRead,
		Delete: resourceGoogleProjectServiceDelete,
		Update: resourceGoogleProjectServiceUpdate,

		Importer: &schema.ResourceImporter{
			State: resourceGoogleProjectServiceImport,
		},

		Timeouts: &schema.ResourceTimeout{
			Create: schema.DefaultTimeout(20 * time.Minute),
			Update: schema.DefaultTimeout(20 * time.Minute),
			Read:   schema.DefaultTimeout(10 * time.Minute),
			Delete: schema.DefaultTimeout(20 * time.Minute),
		},

		Schema: map[string]*schema.Schema{
			"service": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: StringNotInSlice(ignoredProjectServices, false),
			},
			"project": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"disable_dependent_services": {
				Type:     schema.TypeBool,
				Optional: true,
			},

			"disable_on_destroy": {
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
		},
	}
}

func resourceGoogleProjectServiceImport(d *schema.ResourceData, m interface{}) ([]*schema.ResourceData, error) {
	parts := strings.Split(d.Id(), "/")
	if len(parts) != 2 {
		return nil, fmt.Errorf("Invalid google_project_service id format for import, expecting `{project}/{service}`, found %s", d.Id())
	}
	d.Set("project", parts[0])
	d.Set("service", parts[1])
	return []*schema.ResourceData{d}, nil
}

func resourceGoogleProjectServiceCreate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	srv := d.Get("service").(string)
	err = BatchRequestEnableServices(map[string]struct{}{srv: {}}, project, d, config)
	if err != nil {
		return err
	}

	id, err := replaceVars(d, config, "{{project}}/{{service}}")
	if err != nil {
		return fmt.Errorf("unable to construct ID: %s", err)
	}
	d.SetId(id)
	return resourceGoogleProjectServiceRead(d, meta)
}

func resourceGoogleProjectServiceRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}
	srv := d.Get("service").(string)

	enabled, err := isServiceEnabled(project, srv, config)
	if err != nil {
		return handleNotFoundError(err, d, fmt.Sprintf("Project Service %s", d.Id()))
	}
	if enabled {
		d.Set("project", project)
		d.Set("service", srv)
		return nil
	}

	// The service is was not found in enabled services - remove it from state
	log.Printf("[DEBUG] service %s not in enabled services for project %s, removing from state", srv, project)
	d.SetId("")
	return nil
}

func resourceGoogleProjectServiceDelete(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	if disable := d.Get("disable_on_destroy"); !(disable.(bool)) {
		log.Printf("[WARN] Project service %q disable_on_destroy is false, skip disabling service", d.Id())
		d.SetId("")
		return nil
	}

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	service := d.Get("service").(string)
	disableDependencies := d.Get("disable_dependent_services").(bool)
	if err = disableServiceUsageProjectService(service, project, d, config, disableDependencies); err != nil {
		return handleNotFoundError(err, d, fmt.Sprintf("Project Service %s", d.Id()))
	}

	d.SetId("")
	return nil
}

func resourceGoogleProjectServiceUpdate(d *schema.ResourceData, meta interface{}) error {
	// This update method is no-op because the only updatable fields
	// are state/config-only, i.e. they aren't sent in requests to the API.
	return nil
}

// Retrieve enablement state for a given project's service
func isServiceEnabled(project, serviceName string, config *Config) (bool, error) {
	// Verify project for services still exists
	p, err := config.clientResourceManager.Projects.Get(project).Do()
	if err != nil {
		return false, err
	}
	if p.LifecycleState == "DELETE_REQUESTED" {
		// Construct a 404 error for handleNotFoundError
		return false, &googleapi.Error{
			Code:    404,
			Message: "Project deletion was requested",
		}
	}

	resourceName := fmt.Sprintf("projects/%s/services/%s", project, serviceName)
	var srv *serviceusage.GoogleApiServiceusageV1Service
	err = retryTime(func() error {
		var currErr error
		srv, currErr = config.clientServiceUsage.Services.Get(resourceName).Do()
		return currErr
	}, 10)
	if err != nil {
		return false, errwrap.Wrapf(fmt.Sprintf("Failed to list enabled services for project %s: {{err}}", project), err)
	}
	return srv.State == "ENABLED", nil
}

// Disables a project service.
func disableServiceUsageProjectService(service, project string, d *schema.ResourceData, config *Config, disableDependentServices bool) error {
	err := retryTimeDuration(func() error {
		name := fmt.Sprintf("projects/%s/services/%s", project, service)
		sop, err := config.clientServiceUsage.Services.Disable(name, &serviceusage.DisableServiceRequest{
			DisableDependentServices: disableDependentServices,
		}).Do()
		if err != nil {
			return err
		}
		// Wait for the operation to complete
		waitErr := serviceUsageOperationWait(config, sop, "api to disable")
		if waitErr != nil {
			return waitErr
		}
		return nil
	}, d.Timeout(schema.TimeoutDelete))
	if err != nil {
		return fmt.Errorf("Error disabling service %q for project %q: %v", service, project, err)
	}
	return nil
}
<commit_msg>Update service renames comment (#4623)<commit_after>package google

import (
	"fmt"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/serviceusage/v1"
	"log"
	"strings"
	"time"
)

var ignoredProjectServices = []string{"dataproc-control.googleapis.com", "source.googleapis.com", "stackdriverprovisioning.googleapis.com"}

// These services can only be enabled as a side-effect of enabling other services,
// so don't bother storing them in the config or using them for diffing.
var ignoredProjectServicesSet = golangSetFromStringSlice(ignoredProjectServices)

// Service Renames
// we expect when a service is renamed:
// - both service names will continue to be able to be set
// - setting one will effectively enable the other as a dependent
// - GET will return whichever service name is requested
// - LIST responses will not contain the old service name
// renames may be reverted, though, so we should canonicalise both ways until
// the old service is fully removed from the provider
//
// We handle service renames in the provider by pretending that we've read both
// the old and new service names from the API if we see either, and only setting
// the one(s) that existed in prior state in config (if any). If neither exists,
// we'll set the old service name in state.
// Additionally, in case of service rename rollbacks or unexpected early
// removals of services, if we fail to create or delete a service that's been
// renamed we'll retry using an alternate name.
// We try creation by the user-specified value followed by the other value.
// We try deletion by the old value followed by the new value.

// map from old -> new names of services that have been renamed
// these should be removed during major provider versions. comment here with
// "DEPRECATED FOR {{version}} next to entries slated for removal in {{version}}
// upon removal, we should disallow the old name from being used even if it's
// not gone from the underlying API yet
var renamedServices = map[string]string{
	"bigquery-json.googleapis.com": "bigquery.googleapis.com", // DEPRECATED FOR 3.0.0
}

// renamedServices in reverse (new -> old)
var renamedServicesByNewServiceNames = reverseStringMap(renamedServices)

// renamedServices expressed as both old -> new and new -> old
var renamedServicesByOldAndNewServiceNames = mergeStringMaps(renamedServices, renamedServicesByNewServiceNames)

func resourceGoogleProjectService() *schema.Resource {
	return &schema.Resource{
		Create: resourceGoogleProjectServiceCreate,
		Read:   resourceGoogleProjectServiceRead,
		Delete: resourceGoogleProjectServiceDelete,
		Update: resourceGoogleProjectServiceUpdate,

		Importer: &schema.ResourceImporter{
			State: resourceGoogleProjectServiceImport,
		},

		Timeouts: &schema.ResourceTimeout{
			Create: schema.DefaultTimeout(20 * time.Minute),
			Update: schema.DefaultTimeout(20 * time.Minute),
			Read:   schema.DefaultTimeout(10 * time.Minute),
			Delete: schema.DefaultTimeout(20 * time.Minute),
		},

		Schema: map[string]*schema.Schema{
			"service": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: StringNotInSlice(ignoredProjectServices, false),
			},
			"project": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"disable_dependent_services": {
				Type:     schema.TypeBool,
				Optional: true,
			},

			"disable_on_destroy": {
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
		},
	}
}

func resourceGoogleProjectServiceImport(d *schema.ResourceData, m interface{}) ([]*schema.ResourceData, error) {
	parts := strings.Split(d.Id(), "/")
	if len(parts) != 2 {
		return nil, fmt.Errorf("Invalid google_project_service id format for import, expecting `{project}/{service}`, found %s", d.Id())
	}
	d.Set("project", parts[0])
	d.Set("service", parts[1])
	return []*schema.ResourceData{d}, nil
}

func resourceGoogleProjectServiceCreate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	srv := d.Get("service").(string)
	err = BatchRequestEnableServices(map[string]struct{}{srv: {}}, project, d, config)
	if err != nil {
		return err
	}

	id, err := replaceVars(d, config, "{{project}}/{{service}}")
	if err != nil {
		return fmt.Errorf("unable to construct ID: %s", err)
	}
	d.SetId(id)
	return resourceGoogleProjectServiceRead(d, meta)
}

func resourceGoogleProjectServiceRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}
	srv := d.Get("service").(string)

	enabled, err := isServiceEnabled(project, srv, config)
	if err != nil {
		return handleNotFoundError(err, d, fmt.Sprintf("Project Service %s", d.Id()))
	}
	if enabled {
		d.Set("project", project)
		d.Set("service", srv)
		return nil
	}

	// The service is was not found in enabled services - remove it from state
	log.Printf("[DEBUG] service %s not in enabled services for project %s, removing from state", srv, project)
	d.SetId("")
	return nil
}

func resourceGoogleProjectServiceDelete(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	if disable := d.Get("disable_on_destroy"); !(disable.(bool)) {
		log.Printf("[WARN] Project service %q disable_on_destroy is false, skip disabling service", d.Id())
		d.SetId("")
		return nil
	}

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	service := d.Get("service").(string)
	disableDependencies := d.Get("disable_dependent_services").(bool)
	if err = disableServiceUsageProjectService(service, project, d, config, disableDependencies); err != nil {
		return handleNotFoundError(err, d, fmt.Sprintf("Project Service %s", d.Id()))
	}

	d.SetId("")
	return nil
}

func resourceGoogleProjectServiceUpdate(d *schema.ResourceData, meta interface{}) error {
	// This update method is no-op because the only updatable fields
	// are state/config-only, i.e. they aren't sent in requests to the API.
	return nil
}

// Retrieve enablement state for a given project's service
func isServiceEnabled(project, serviceName string, config *Config) (bool, error) {
	// Verify project for services still exists
	p, err := config.clientResourceManager.Projects.Get(project).Do()
	if err != nil {
		return false, err
	}
	if p.LifecycleState == "DELETE_REQUESTED" {
		// Construct a 404 error for handleNotFoundError
		return false, &googleapi.Error{
			Code:    404,
			Message: "Project deletion was requested",
		}
	}

	resourceName := fmt.Sprintf("projects/%s/services/%s", project, serviceName)
	var srv *serviceusage.GoogleApiServiceusageV1Service
	err = retryTime(func() error {
		var currErr error
		srv, currErr = config.clientServiceUsage.Services.Get(resourceName).Do()
		return currErr
	}, 10)
	if err != nil {
		return false, errwrap.Wrapf(fmt.Sprintf("Failed to list enabled services for project %s: {{err}}", project), err)
	}
	return srv.State == "ENABLED", nil
}

// Disables a project service.
func disableServiceUsageProjectService(service, project string, d *schema.ResourceData, config *Config, disableDependentServices bool) error {
	err := retryTimeDuration(func() error {
		name := fmt.Sprintf("projects/%s/services/%s", project, service)
		sop, err := config.clientServiceUsage.Services.Disable(name, &serviceusage.DisableServiceRequest{
			DisableDependentServices: disableDependentServices,
		}).Do()
		if err != nil {
			return err
		}
		// Wait for the operation to complete
		waitErr := serviceUsageOperationWait(config, sop, "api to disable")
		if waitErr != nil {
			return waitErr
		}
		return nil
	}, d.Timeout(schema.TimeoutDelete))
	if err != nil {
		return fmt.Errorf("Error disabling service %q for project %q: %v", service, project, err)
	}
	return nil
}
<|endoftext|>


<commit_before>package storage

import (
	"fmt"
	"github.com/pkg/errors"
	"github.com/viant/toolbox"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"path"
	"time"
)

var fileMode os.FileMode = 0644

//Service represents abstract way to accessing local or remote storage
type fileStorageService struct{}

func openFileFromUrl(URL string) (*os.File, error) {
	parsedUrl, err := url.Parse(URL)
	if err != nil {
		return nil, err
	}
	if parsedUrl.Scheme != "file" {
		return nil, fmt.Errorf("Invalid schema, expected file but had: %v", parsedUrl.Scheme)
	}
	return os.Open(parsedUrl.Path)
}

//List returns a list of object for supplied url
func (s *fileStorageService) List(URL string) ([]Object, error) {
	file, err := openFileFromUrl(URL)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	stat, err := file.Stat()
	if err == nil {
		if ! stat.IsDir() {
			modTime := stat.ModTime()
			return []Object{
				newFileObject(URL, StorageObjectContentType, stat, &modTime, stat.Size()),
			}, nil
		}
	}

	files, err := file.Readdir(0)
	if err != nil {
		return nil, err
	}
	var result = make([]Object, 0)
	for _, fileInfo := range files {
		fileUrl := URL + "/" + fileInfo.Name()
		objectType := StorageObjectContentType
		if fileInfo.IsDir() {
			objectType = StorageObjectFolderType
		}
		modTime := fileInfo.ModTime()
		result = append(result, newFileObject(fileUrl, objectType, fileInfo, &modTime, fileInfo.Size()))
	}
	return result, nil
}

//Exists returns true if resource exists
func (s *fileStorageService) Exists(URL string) (bool, error) {
	parsedUrl, err := url.Parse(URL)
	if err != nil {
		return false, err
	}
	if parsedUrl.Scheme != "file" {
		return false, fmt.Errorf("Invalid schema, expected file but had: %v", parsedUrl.Scheme)
	}
	return toolbox.FileExists(parsedUrl.Path), nil
}

//Object returns a Object for supplied url
func (s *fileStorageService) StorageObject(URL string) (Object, error) {
	file, err := openFileFromUrl(URL)
	if err != nil {
		return nil, err
	}
	defer file.Close()
	fileInfo, err := os.Stat(file.Name())
	if err != nil {
		return nil, err
	}
	objectType := 0
	switch mode := fileInfo.Mode(); {

	case mode.IsDir():
		// do directory stuff
		objectType = StorageObjectFolderType
	case mode.IsRegular():
		objectType = StorageObjectContentType
	}
	modTime := fileInfo.ModTime()
	return newFileObject(URL, objectType, &fileInfo, &modTime, fileInfo.Size()), nil
}

//Download returns reader for downloaded storage object
func (s *fileStorageService) Download(object Object) (io.Reader, error) {
	reader, _, err := toolbox.OpenReaderFromURL(object.URL())
	return reader, err
}

//Upload uploads provided reader content for supplied url.
func (s *fileStorageService) Upload(URL string, reader io.Reader) error {
	parsedUrl, err := url.Parse(URL)
	if err != nil {
		return err
	}
	if parsedUrl.Scheme != "file" {
		return fmt.Errorf("Invalid schema, expected file but had: %v", parsedUrl.Scheme)
	}

	parentDir, _ := path.Split(parsedUrl.Path)
	err = toolbox.CreateDirIfNotExist(parentDir)
	if err != nil {
		return err
	}
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(parsedUrl.Path, data, fileMode)
}

func (s *fileStorageService) Register(schema string, service Service) error {
	return errors.New("unsupported")
}

//Delete removes passed in storage object
func (s *fileStorageService) Delete(object Object) error {
	fileName, err := toolbox.FileFromURL(object.URL())
	if err != nil {
		return err
	}
	return os.Remove(fileName)
}

type fileStorageObject struct {
	*AbstractObject
}

func (o *fileStorageObject) Unwrap(target interface{}) error {
	if fileInfo, casted := target.(**os.FileInfo); casted {

		source, ok := o.Source.(*os.FileInfo)
		if !ok {
			return fmt.Errorf("Failed to cast %T into %T", o.Source, target)
		}
		*fileInfo = source
		return nil
	}

	return fmt.Errorf("unsuported target %T", target)
}

func newFileObject(url string, objectType int, source interface{}, lastModified *time.Time, size int64) Object {
	abstract := NewAbstractStorageObject(url, source, objectType, lastModified, size)
	result := &fileStorageObject{
		AbstractObject: abstract,
	}
	result.AbstractObject.Object = result
	return result
}
<commit_msg>patch open file leackage<commit_after>package storage

import (
	"fmt"
	"github.com/pkg/errors"
	"github.com/viant/toolbox"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"path"
	"time"
	"bytes"
)

var fileMode os.FileMode = 0644

//Service represents abstract way to accessing local or remote storage
type fileStorageService struct{}

func openFileFromUrl(URL string) (*os.File, error) {
	parsedUrl, err := url.Parse(URL)
	if err != nil {
		return nil, err
	}
	if parsedUrl.Scheme != "file"  {
		return nil, fmt.Errorf("Invalid schema, expected file but had: %v", parsedUrl.Scheme)
	}
	return os.Open(parsedUrl.Path)
}

//List returns a list of object for supplied url
func (s *fileStorageService) List(URL string) ([]Object, error) {
	file, err := openFileFromUrl(URL)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	stat, err := file.Stat()
	if err == nil {
		if ! stat.IsDir() {
			modTime := stat.ModTime()
			return []Object{
				newFileObject(URL, StorageObjectContentType, stat, &modTime, stat.Size()),
			}, nil
		}
	}

	files, err := file.Readdir(0)
	if err != nil {
		return nil, err
	}
	var result = make([]Object, 0)
	for _, fileInfo := range files {
		fileUrl := URL + "/" + fileInfo.Name()
		objectType := StorageObjectContentType
		if fileInfo.IsDir() {
			objectType = StorageObjectFolderType
		}
		modTime := fileInfo.ModTime()
		result = append(result, newFileObject(fileUrl, objectType, fileInfo, &modTime, fileInfo.Size()))
	}
	return result, nil
}

//Exists returns true if resource exists
func (s *fileStorageService) Exists(URL string) (bool, error) {
	parsedUrl, err := url.Parse(URL)
	if err != nil {
		return false, err
	}
	if parsedUrl.Scheme != "file" {
		return false, fmt.Errorf("Invalid schema, expected file but had: %v", parsedUrl.Scheme)
	}
	return toolbox.FileExists(parsedUrl.Path), nil
}

//Object returns a Object for supplied url
func (s *fileStorageService) StorageObject(URL string) (Object, error) {
	file, err := openFileFromUrl(URL)
	if err != nil {
		return nil, err
	}
	defer file.Close()
	fileInfo, err := os.Stat(file.Name())
	if err != nil {
		return nil, err
	}
	objectType := 0
	switch mode := fileInfo.Mode(); {

	case mode.IsDir():
		// do directory stuff
		objectType = StorageObjectFolderType
	case mode.IsRegular():
		objectType = StorageObjectContentType
	}
	modTime := fileInfo.ModTime()
	return newFileObject(URL, objectType, &fileInfo, &modTime, fileInfo.Size()), nil
}

//Download returns reader for downloaded storage object
func (s *fileStorageService) Download(object Object) (io.Reader, error) {
	reader, _, err := toolbox.OpenReaderFromURL(object.URL())
	if err != nil {
		defer reader.Close()
	}
	data, err :=  ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}
	return bytes.NewReader(data), err
}

//Upload uploads provided reader content for supplied url.
func (s *fileStorageService) Upload(URL string, reader io.Reader) error {
	parsedUrl, err := url.Parse(URL)
	if err != nil {
		return err
	}
	if parsedUrl.Scheme != "file" {
		return fmt.Errorf("Invalid schema, expected file but had: %v", parsedUrl.Scheme)
	}

	parentDir, _ := path.Split(parsedUrl.Path)
	err = toolbox.CreateDirIfNotExist(parentDir)
	if err != nil {
		return err
	}
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(parsedUrl.Path, data, fileMode)
}

func (s *fileStorageService) Register(schema string, service Service) error {
	return errors.New("unsupported")
}

//Delete removes passed in storage object
func (s *fileStorageService) Delete(object Object) error {
	fileName, err := toolbox.FileFromURL(object.URL())
	if err != nil {
		return err
	}
	return os.Remove(fileName)
}

type fileStorageObject struct {
	*AbstractObject
}

func (o *fileStorageObject) Unwrap(target interface{}) error {
	if fileInfo, casted := target.(**os.FileInfo); casted {

		source, ok := o.Source.(*os.FileInfo)
		if !ok {
			return fmt.Errorf("Failed to cast %T into %T", o.Source, target)
		}
		*fileInfo = source
		return nil
	}

	return fmt.Errorf("unsuported target %T", target)
}

func newFileObject(url string, objectType int, source interface{}, lastModified *time.Time, size int64) Object {
	abstract := NewAbstractStorageObject(url, source, objectType, lastModified, size)
	result := &fileStorageObject{
		AbstractObject: abstract,
	}
	result.AbstractObject.Object = result
	return result
}
<|endoftext|>


<commit_before>package main

import (
    "os"
    "os/signal"
    "syscall"
    "log/syslog"
)

var g_logger, _ = syslog.NewLogger(syslog.LOG_NOTICE | syslog.LOG_CRON, 0)
var g_err_logger, _ = syslog.NewLogger(syslog.LOG_ERR | syslog.LOG_CRON, 0)

func stopServerOnSignal(server *IpcServer, jm *JobManager) {
    // Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM, os.Kill)

	// Block until a signal is received.
	<-c
	g_logger.Printf("Got signal.\n")
	server.Stop()
	jm.Cancel()
	jm.Wait()
	os.Exit(0)
}

func main() {
    var err error

    // run them
	manager, err := NewJobManager()
	if err != nil {
        g_err_logger.Printf("Error: %v\n", err)
        os.Exit(1)
    }
	cmdChan, err := manager.Launch()
	if err != nil {
        g_err_logger.Printf("Error: %v\n", err)
        os.Exit(1)
    }

    // make IPC server
    ipcServer := NewIpcServer(cmdChan)
    go stopServerOnSignal(ipcServer, manager)
    err = ipcServer.Launch()
    if err != nil {
        g_err_logger.Printf("Error: %v\n", err)
        os.Exit(1)
    }
    defer ipcServer.Stop()

    manager.Wait()
    g_logger.Printf("Job manager stopped.\n")
}
<commit_msg>Issue #16: Check g_err_logger for nil.<commit_after>package main

import (
    "os"
    "os/signal"
    "syscall"
    "log/syslog"
)

var g_logger, _ = syslog.NewLogger(syslog.LOG_NOTICE | syslog.LOG_CRON, 0)
var g_err_logger, _ = syslog.NewLogger(syslog.LOG_ERR | syslog.LOG_CRON, 0)

func stopServerOnSignal(server *IpcServer, jm *JobManager) {
    // Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM, os.Kill)

	// Block until a signal is received.
	<-c
	g_logger.Printf("Got signal.\n")
	server.Stop()
	jm.Cancel()
	jm.Wait()
	os.Exit(0)
}

func main() {
    var err error

    // run them
	manager, err := NewJobManager()
	if err != nil {
        if g_err_logger != nil {
            g_err_logger.Printf("Error: %v\n", err)
        }
        os.Exit(1)
    }
	cmdChan, err := manager.Launch()
	if err != nil {
        if g_err_logger != nil {
            g_err_logger.Printf("Error: %v\n", err)
        }
        os.Exit(1)
    }

    // make IPC server
    ipcServer := NewIpcServer(cmdChan)
    go stopServerOnSignal(ipcServer, manager)
    err = ipcServer.Launch()
    if err != nil {
        if g_err_logger != nil {
            g_err_logger.Printf("Error: %v\n", err)
        }
        os.Exit(1)
    }
    defer ipcServer.Stop()

    manager.Wait()
    g_logger.Printf("Job manager stopped.\n")
}
<|endoftext|>


<commit_before>/* Provides a low-level Go interface to the systemd journal C API.

All public methods map closely to the sd-journal API functions. See the
sd-journal.h documentation[1] for information about each function.

[1] http://www.freedesktop.org/software/systemd/man/sd-journal.html
*/
package journal

/*
#cgo pkg-config: libsystemd-journal
#include <systemd/sd-journal.h>
#include <stdlib.h>
#include <syslog.h>

int go_sd_journal_print(int priority, char* s) {
	int r;
	r = sd_journal_print(priority, "%s", s);
	return r;
}
*/
import "C"
import (
	"fmt"
	"unsafe"
)

// Journal entry field strings which correspond to:
// http://www.freedesktop.org/software/systemd/man/systemd.journal-fields.html
const (
	SD_JOURNAL_FIELD_SYSTEMD_UNIT = "_SYSTEMD_UNIT"
)

// Journal event constants
const (
	SD_JOURNAL_NOP        = int(C.SD_JOURNAL_NOP)
	SD_JOURNAL_APPEND     = int(C.SD_JOURNAL_APPEND)
	SD_JOURNAL_INVALIDATE = int(C.SD_JOURNAL_INVALIDATE)
)

// A Journal is a Go wrapper of an sd_journal structure.
type Journal struct {
	cjournal *C.sd_journal
}

// A Match is a convenience wrapper to describe filters supplied to AddMatch.
type Match struct {
	Field string
	Value string
}

// String returns a string representation of a Match suitable for use with AddMatch.
func (m *Match) String() string {
	return m.Field + "=" + m.Value
}

func NewJournal() (*Journal, error) {
	j := &Journal{}
	err := C.sd_journal_open(&j.cjournal, C.SD_JOURNAL_LOCAL_ONLY)

	if err < 0 {
		return nil, fmt.Errorf("failed to open journal: %s", err)
	}

	return j, nil
}

func (j *Journal) Close() error {
	C.sd_journal_close(j.cjournal)
	return nil
}

func (j *Journal) AddMatch(match string) error {
	m := C.CString(match)
	defer C.free(unsafe.Pointer(m))

	C.sd_journal_add_match(j.cjournal, unsafe.Pointer(m), C.size_t(len(match)))
	return nil
}

func (j *Journal) Next() (int, error) {
	r := C.sd_journal_next(j.cjournal)

	if r < 0 {
		return int(r), fmt.Errorf("failed to iterate journal: %d", r)
	}

	return int(r), nil
}

func (j *Journal) Previous() (uint64, error) {
	r := C.sd_journal_previous(j.cjournal)

	if r < 0 {
		return uint64(r), fmt.Errorf("failed to iterate journal: %d", r)
	}

	return uint64(r), nil
}

func (j *Journal) PreviousSkip(skip uint64) (uint64, error) {
	r := C.sd_journal_previous_skip(j.cjournal, C.uint64_t(skip))

	if r < 0 {
		return uint64(r), fmt.Errorf("failed to iterate journal: %d", r)
	}

	return uint64(r), nil
}

func (j *Journal) GetData(field string) (string, error) {
	f := C.CString(field)
	defer C.free(unsafe.Pointer(f))

	var d unsafe.Pointer
	var l C.size_t

	err := C.sd_journal_get_data(j.cjournal, f, &d, &l)

	if err < 0 {
		return "", fmt.Errorf("failed to read message: %d", err)
	} else {
		msg := C.GoStringN((*C.char)(d), C.int(l))
		return msg, nil
	}
}

func (j *Journal) GetRealtimeUsec() (uint64, error) {
	var usec C.uint64_t

	r := C.sd_journal_get_realtime_usec(j.cjournal, &usec)

	if r < 0 {
		return 0, fmt.Errorf("error getting timestamp for entry: %d", r)
	}

	return uint64(usec), nil
}

func (j *Journal) Print(priority int, message string) error {
	m := C.CString(message)
	defer C.free(unsafe.Pointer(m))

	err := C.go_sd_journal_print(C.LOG_INFO, m)

	if err == 0 {
		return nil
	} else {
		return fmt.Errorf("failed to print message: %s", err)
	}
}

func (j *Journal) SeekTail() error {
	err := C.sd_journal_seek_tail(j.cjournal)

	if err != 0 {
		return fmt.Errorf("failed to seek to tail of journal: %s", err)
	} else {
		return nil
	}
}

func (j *Journal) SeekRealtimeUsec(usec uint64) error {
	err := C.sd_journal_seek_realtime_usec(j.cjournal, C.uint64_t(usec))

	if err != 0 {
		return fmt.Errorf("failed to seek to %d: %d", usec, int(err))
	}

	return nil
}

func (j *Journal) Wait(timeout uint64) int {
	r := C.sd_journal_wait(j.cjournal, C.uint64_t(timeout))

	return int(r)
}
<commit_msg>journal: clean up error handling<commit_after>/* Provides a low-level Go interface to the systemd journal C API.

All public methods map closely to the sd-journal API functions. See the
sd-journal.h documentation[1] for information about each function.

[1] http://www.freedesktop.org/software/systemd/man/sd-journal.html
*/
package journal

/*
#cgo pkg-config: libsystemd-journal
#include <systemd/sd-journal.h>
#include <stdlib.h>
#include <syslog.h>

int go_sd_journal_print(int priority, char* s) {
	int r;
	r = sd_journal_print(priority, "%s", s);
	return r;
}
*/
import "C"
import (
	"fmt"
	"unsafe"
)

// Journal entry field strings which correspond to:
// http://www.freedesktop.org/software/systemd/man/systemd.journal-fields.html
const (
	SD_JOURNAL_FIELD_SYSTEMD_UNIT = "_SYSTEMD_UNIT"
)

// Journal event constants
const (
	SD_JOURNAL_NOP        = int(C.SD_JOURNAL_NOP)
	SD_JOURNAL_APPEND     = int(C.SD_JOURNAL_APPEND)
	SD_JOURNAL_INVALIDATE = int(C.SD_JOURNAL_INVALIDATE)
)

// A Journal is a Go wrapper of an sd_journal structure.
type Journal struct {
	cjournal *C.sd_journal
}

// A Match is a convenience wrapper to describe filters supplied to AddMatch.
type Match struct {
	Field string
	Value string
}

// String returns a string representation of a Match suitable for use with AddMatch.
func (m *Match) String() string {
	return m.Field + "=" + m.Value
}

func NewJournal() (*Journal, error) {
	j := &Journal{}
	err := C.sd_journal_open(&j.cjournal, C.SD_JOURNAL_LOCAL_ONLY)

	if err < 0 {
		return nil, fmt.Errorf("failed to open journal: %s", err)
	}

	return j, nil
}

func (j *Journal) Close() error {
	C.sd_journal_close(j.cjournal)
	return nil
}

func (j *Journal) AddMatch(match string) error {
	m := C.CString(match)
	defer C.free(unsafe.Pointer(m))

	C.sd_journal_add_match(j.cjournal, unsafe.Pointer(m), C.size_t(len(match)))
	return nil
}

func (j *Journal) Next() (int, error) {
	r := C.sd_journal_next(j.cjournal)

	if r < 0 {
		return int(r), fmt.Errorf("failed to iterate journal: %d", r)
	}

	return int(r), nil
}

func (j *Journal) Previous() (uint64, error) {
	r := C.sd_journal_previous(j.cjournal)

	if r < 0 {
		return uint64(r), fmt.Errorf("failed to iterate journal: %d", r)
	}

	return uint64(r), nil
}

func (j *Journal) PreviousSkip(skip uint64) (uint64, error) {
	r := C.sd_journal_previous_skip(j.cjournal, C.uint64_t(skip))

	if r < 0 {
		return uint64(r), fmt.Errorf("failed to iterate journal: %d", r)
	}

	return uint64(r), nil
}

func (j *Journal) GetData(field string) (string, error) {
	f := C.CString(field)
	defer C.free(unsafe.Pointer(f))

	var d unsafe.Pointer
	var l C.size_t

	err := C.sd_journal_get_data(j.cjournal, f, &d, &l)

	if err < 0 {
		return "", fmt.Errorf("failed to read message: %d", err)
	}

	msg := C.GoStringN((*C.char)(d), C.int(l))
	return msg, nil
}

func (j *Journal) GetRealtimeUsec() (uint64, error) {
	var usec C.uint64_t

	r := C.sd_journal_get_realtime_usec(j.cjournal, &usec)

	if r < 0 {
		return 0, fmt.Errorf("error getting timestamp for entry: %d", r)
	}

	return uint64(usec), nil
}

func (j *Journal) Print(priority int, message string) error {
	m := C.CString(message)
	defer C.free(unsafe.Pointer(m))

	err := C.go_sd_journal_print(C.LOG_INFO, m)

	if err != 0 {
		return fmt.Errorf("failed to print message: %s", err)
	}

	return nil
}

func (j *Journal) SeekTail() error {
	err := C.sd_journal_seek_tail(j.cjournal)

	if err != 0 {
		return fmt.Errorf("failed to seek to tail of journal: %s", err)
	}

	return nil
}

func (j *Journal) SeekRealtimeUsec(usec uint64) error {
	err := C.sd_journal_seek_realtime_usec(j.cjournal, C.uint64_t(usec))

	if err != 0 {
		return fmt.Errorf("failed to seek to %d: %d", usec, int(err))
	}

	return nil
}

func (j *Journal) Wait(timeout uint64) int {
	r := C.sd_journal_wait(j.cjournal, C.uint64_t(timeout))

	return int(r)
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package license

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/admpub/license_gen/lib"
	"github.com/admpub/log"
	"github.com/admpub/restful/rest"
	"github.com/shirou/gopsutil/cpu"
	"github.com/webx-top/com"
	"github.com/webx-top/echo"
)

var (
	trackerURL         = `http://www.webx.top/product/script/nging/tracker.js`
	productURL         = `http://www.webx.top/product/detail/nging`
	licenseURL         = `http://www.webx.top/product/license/nging`
	versionURL         = `http://www.webx.top/product/version/nging`
	licenseFileName    = `license.key`
	licenseFile        = filepath.Join(echo.Wd(), licenseFileName)
	licenseExists      bool
	licenseError       = lib.UnlicensedVersion
	licenseData        *lib.LicenseData
	licenseVersion     string
	machineID          string
	domain             string
	emptyLicense       = lib.LicenseData{}
	ErrLicenseNotFound = errors.New(`License does not exist`)
	SkipLicenseCheck   = true
)

type ServerURL struct {
	Tracker         string //用于统计分析的js地址
	Product         string //该产品的详情介绍页面网址
	License         string //许可证验证和许可证下载API网址
	Version         string //该产品最新版本信息API网址
	LicenseFileName string //许可证文件名称
}

func (s *ServerURL) Apply() {
	trackerURL = s.Tracker
	productURL = s.Product
	licenseURL = s.License
	versionURL = s.Version
	if len(s.LicenseFileName) > 0 {
		licenseFileName = s.LicenseFileName
		licenseFile = filepath.Join(echo.Wd(), licenseFileName)
	}
}

func SetServerURL(s *ServerURL) {
	if s != nil {
		s.Apply()
	}
}

func SetVersion(version string) {
	licenseVersion = version
}

func Version() string {
	return licenseVersion
}

func ProductURL() string {
	return productURL
}

func TrackerURL() string {
	return trackerURL
}

func FilePath() string {
	return licenseFile
}

func FileName() string {
	return licenseFileName
}

func Exists() bool {
	return licenseExists
}

func Error() error {
	return licenseError
}

func License() lib.LicenseData {
	if licenseData == nil {
		return emptyLicense
	}
	return *licenseData
}

// Check 检查权限
func Check(machineID, domain string) error {
	if SkipLicenseCheck {
		return nil
	}
	licenseError = validateFromOfficial(machineID, domain)
	if licenseError != ErrConnectionNotFound {
		return licenseError
	}
	//当官方服务器不可用时才验证本地许可证
	licenseError = Validate()
	return licenseError
}

func Ok(domain string) bool {
	if SkipLicenseCheck {
		return true
	}
	switch licenseError {
	case nil:
		if licenseData == nil {
			licenseError = lib.UnlicensedVersion
			return false
		}
		if !licenseData.Info.Expiration.IsZero() && time.Now().After(licenseData.Info.Expiration) {
			licenseError = lib.ExpiredLicense
			return false
		}
		return true
	default:
		machineID, _ := MachineID()
		err := Check(machineID, domain)
		if err == nil {
			licenseError = nil
			return true
		}
		log.Warn(err)
	}
	return false
}

// Validation 定义验证器
type Validation struct {
	NowVersions []string
}

// Validate 参数验证器
func (v *Validation) Validate(data *lib.LicenseData) error {
	if err := data.CheckExpiration(); err != nil {
		return err
	}
	if err := data.CheckVersion(v.NowVersions...); err != nil {
		return err
	}
	mid, err := MachineID()
	if err != nil {
		return err
	}
	if data.Info.MachineID != mid {
		return lib.InvalidMachineID
	}
	return nil
}

// Validate 验证授权
func Validate() error {
	licenseExists = com.FileExists(FilePath())
	if !licenseExists {
		licenseError = ErrLicenseNotFound
		return licenseError
	}
	b, err := ioutil.ReadFile(FilePath())
	if err != nil {
		return err
	}
	validator := &Validation{
		NowVersions: []string{licenseVersion},
	}
	licenseData, err = lib.CheckLicenseStringAndReturning(string(b), PublicKey(), validator)
	return err
}

// Save 保存授权文件
func Save(b []byte) error {
	return ioutil.WriteFile(licenseFile, b, os.ModePerm)
}

// Generate 生成演示版证书
func Generate(privBytes []byte, pemSaveDirs ...string) error {
	var err error
	if privBytes == nil {
		var pubBytes []byte
		pubBytes, privBytes, err = lib.GenerateCertificateData(2048)
		if err != nil {
			return err
		}
		publicKey = string(pubBytes)
		var pemSaveDir string
		if len(pemSaveDirs) > 0 {
			pemSaveDir = pemSaveDirs[0]
		} else {
			pemSaveDir = filepath.Join(echo.Wd(), `data`)
		}
		if len(pemSaveDir) > 0 {
			err = ioutil.WriteFile(filepath.Join(pemSaveDir, `nging.pem.pub`), pubBytes, os.ModePerm)
			if err != nil {
				return err
			}
			err = ioutil.WriteFile(filepath.Join(pemSaveDir, `nging.pem`), privBytes, os.ModePerm)
			if err != nil {
				return err
			}
		}
	}
	info := &lib.LicenseInfo{
		Name:       `demo`,
		LicenseID:  `0`,
		Version:    licenseVersion,
		Expiration: time.Now().Add(30 * 24 * time.Hour),
	}
	info.MachineID, err = MachineID()
	if err != nil {
		return err
	}
	licBytes, err := lib.GenerateLicense(info, string(privBytes))
	if err != nil {
		return err
	}
	return Save(licBytes)
}

// MachineID 生成当前机器的机器码
func MachineID() (string, error) {
	if len(machineID) > 0 {
		return machineID, nil
	}
	addrs, err := lib.MACAddresses(false)
	if err != nil {
		return ``, err
	}
	if len(addrs) < 1 {
		return ``, lib.ErrorMachineID
	}
	cpuInfo, err := cpu.Info()
	if err != nil {
		return ``, err
	}
	var cpuID string
	if len(cpuInfo) > 0 {
		cpuID = cpuInfo[0].PhysicalID
		if len(cpuID) == 0 {
			cpuID = com.Md5(com.Dump(cpuInfo, false))
		}
	}
	machineID = com.MakePassword(lib.Hash(addrs[0])+`#`+cpuID, `coscms`, 3, 8, 19)
	return machineID, err
}

// FullLicenseURL 包含完整参数的授权网址
func FullLicenseURL(machineID, domain string) string {
	return licenseURL + `?os=` + runtime.GOOS + `&arch=` + runtime.GOARCH + `&version=` + licenseVersion + `&machineID=` + machineID + `&domain=` + domain + `&time=` + time.Now().Format(`20060102-150405`)
}

// Download 从官方服务器重新下载许可证
func Download(machineID, domain string) error {
	operation := `获取授权证书失败：%v`
	response := rest.Get(FullLicenseURL(machineID, domain) + `&pipe=download`)
	if response.StatusCode != http.StatusOK {
		return fmt.Errorf(operation, response.Status)
	}
	officialResp := &OfficialResp{}
	err := response.FillUp(officialResp)
	if err != nil {
		return fmt.Errorf(operation, err)
	}
	if officialResp.Code != 1 {
		return fmt.Errorf(`%v`, officialResp.Info)
	}
	if officialResp.Data == nil {
		return errors.New(`下载证书失败：官方数据返回异常`)
	}
	if com.FileExists(licenseFile) {
		err = os.Rename(licenseFile, licenseFile+`.`+time.Now().Format(`20060102150405`))
		if err != nil {
			return err
		}
	}
	licenseData = &officialResp.Data.LicenseData
	b, err := com.JSONEncode(licenseData, `  `)
	if err != nil {
		b = []byte(err.Error())
	}
	err = ioutil.WriteFile(licenseFile, b, os.ModePerm)
	if err != nil {
		return fmt.Errorf(`保存授权证书失败：%v`, err)
	}
	return err
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package license

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/admpub/license_gen/lib"
	"github.com/admpub/log"
	"github.com/admpub/restful/rest"
	"github.com/shirou/gopsutil/cpu"
	"github.com/webx-top/com"
	"github.com/webx-top/echo"
)

var (
	trackerURL         = `//www.webx.top/product/script/nging/tracker.js`
	productURL         = `http://www.webx.top/product/detail/nging`
	licenseURL         = `http://www.webx.top/product/license/nging`
	versionURL         = `http://www.webx.top/product/version/nging`
	licenseFileName    = `license.key`
	licenseFile        = filepath.Join(echo.Wd(), licenseFileName)
	licenseExists      bool
	licenseError       = lib.UnlicensedVersion
	licenseData        *lib.LicenseData
	licenseVersion     string
	machineID          string
	domain             string
	emptyLicense       = lib.LicenseData{}
	ErrLicenseNotFound = errors.New(`License does not exist`)
	SkipLicenseCheck   = true
)

type ServerURL struct {
	Tracker         string //用于统计分析的js地址
	Product         string //该产品的详情介绍页面网址
	License         string //许可证验证和许可证下载API网址
	Version         string //该产品最新版本信息API网址
	LicenseFileName string //许可证文件名称
}

func (s *ServerURL) Apply() {
	trackerURL = s.Tracker
	productURL = s.Product
	licenseURL = s.License
	versionURL = s.Version
	if len(s.LicenseFileName) > 0 {
		licenseFileName = s.LicenseFileName
		licenseFile = filepath.Join(echo.Wd(), licenseFileName)
	}
}

func SetServerURL(s *ServerURL) {
	if s != nil {
		s.Apply()
	}
}

func SetVersion(version string) {
	licenseVersion = version
}

func Version() string {
	return licenseVersion
}

func ProductURL() string {
	return productURL
}

func TrackerURL() string {
	return trackerURL
}

func FilePath() string {
	return licenseFile
}

func FileName() string {
	return licenseFileName
}

func Exists() bool {
	return licenseExists
}

func Error() error {
	return licenseError
}

func License() lib.LicenseData {
	if licenseData == nil {
		return emptyLicense
	}
	return *licenseData
}

// Check 检查权限
func Check(machineID, domain string) error {
	if SkipLicenseCheck {
		return nil
	}
	licenseError = validateFromOfficial(machineID, domain)
	if licenseError != ErrConnectionNotFound {
		return licenseError
	}
	//当官方服务器不可用时才验证本地许可证
	licenseError = Validate()
	return licenseError
}

func Ok(domain string) bool {
	if SkipLicenseCheck {
		return true
	}
	switch licenseError {
	case nil:
		if licenseData == nil {
			licenseError = lib.UnlicensedVersion
			return false
		}
		if !licenseData.Info.Expiration.IsZero() && time.Now().After(licenseData.Info.Expiration) {
			licenseError = lib.ExpiredLicense
			return false
		}
		return true
	default:
		machineID, _ := MachineID()
		err := Check(machineID, domain)
		if err == nil {
			licenseError = nil
			return true
		}
		log.Warn(err)
	}
	return false
}

// Validation 定义验证器
type Validation struct {
	NowVersions []string
}

// Validate 参数验证器
func (v *Validation) Validate(data *lib.LicenseData) error {
	if err := data.CheckExpiration(); err != nil {
		return err
	}
	if err := data.CheckVersion(v.NowVersions...); err != nil {
		return err
	}
	mid, err := MachineID()
	if err != nil {
		return err
	}
	if data.Info.MachineID != mid {
		return lib.InvalidMachineID
	}
	return nil
}

// Validate 验证授权
func Validate() error {
	licenseExists = com.FileExists(FilePath())
	if !licenseExists {
		licenseError = ErrLicenseNotFound
		return licenseError
	}
	b, err := ioutil.ReadFile(FilePath())
	if err != nil {
		return err
	}
	validator := &Validation{
		NowVersions: []string{licenseVersion},
	}
	licenseData, err = lib.CheckLicenseStringAndReturning(string(b), PublicKey(), validator)
	return err
}

// Save 保存授权文件
func Save(b []byte) error {
	return ioutil.WriteFile(licenseFile, b, os.ModePerm)
}

// Generate 生成演示版证书
func Generate(privBytes []byte, pemSaveDirs ...string) error {
	var err error
	if privBytes == nil {
		var pubBytes []byte
		pubBytes, privBytes, err = lib.GenerateCertificateData(2048)
		if err != nil {
			return err
		}
		publicKey = string(pubBytes)
		var pemSaveDir string
		if len(pemSaveDirs) > 0 {
			pemSaveDir = pemSaveDirs[0]
		} else {
			pemSaveDir = filepath.Join(echo.Wd(), `data`)
		}
		if len(pemSaveDir) > 0 {
			err = ioutil.WriteFile(filepath.Join(pemSaveDir, `nging.pem.pub`), pubBytes, os.ModePerm)
			if err != nil {
				return err
			}
			err = ioutil.WriteFile(filepath.Join(pemSaveDir, `nging.pem`), privBytes, os.ModePerm)
			if err != nil {
				return err
			}
		}
	}
	info := &lib.LicenseInfo{
		Name:       `demo`,
		LicenseID:  `0`,
		Version:    licenseVersion,
		Expiration: time.Now().Add(30 * 24 * time.Hour),
	}
	info.MachineID, err = MachineID()
	if err != nil {
		return err
	}
	licBytes, err := lib.GenerateLicense(info, string(privBytes))
	if err != nil {
		return err
	}
	return Save(licBytes)
}

// MachineID 生成当前机器的机器码
func MachineID() (string, error) {
	if len(machineID) > 0 {
		return machineID, nil
	}
	addrs, err := lib.MACAddresses(false)
	if err != nil {
		return ``, err
	}
	if len(addrs) < 1 {
		return ``, lib.ErrorMachineID
	}
	cpuInfo, err := cpu.Info()
	if err != nil {
		return ``, err
	}
	var cpuID string
	if len(cpuInfo) > 0 {
		cpuID = cpuInfo[0].PhysicalID
		if len(cpuID) == 0 {
			cpuID = com.Md5(com.Dump(cpuInfo, false))
		}
	}
	machineID = com.MakePassword(lib.Hash(addrs[0])+`#`+cpuID, `coscms`, 3, 8, 19)
	return machineID, err
}

// FullLicenseURL 包含完整参数的授权网址
func FullLicenseURL(machineID, domain string) string {
	return licenseURL + `?os=` + runtime.GOOS + `&arch=` + runtime.GOARCH + `&version=` + licenseVersion + `&machineID=` + machineID + `&domain=` + domain + `&time=` + time.Now().Format(`20060102-150405`)
}

// Download 从官方服务器重新下载许可证
func Download(machineID, domain string) error {
	operation := `获取授权证书失败：%v`
	response := rest.Get(FullLicenseURL(machineID, domain) + `&pipe=download`)
	if response.StatusCode != http.StatusOK {
		return fmt.Errorf(operation, response.Status)
	}
	officialResp := &OfficialResp{}
	err := response.FillUp(officialResp)
	if err != nil {
		return fmt.Errorf(operation, err)
	}
	if officialResp.Code != 1 {
		return fmt.Errorf(`%v`, officialResp.Info)
	}
	if officialResp.Data == nil {
		return errors.New(`下载证书失败：官方数据返回异常`)
	}
	if com.FileExists(licenseFile) {
		err = os.Rename(licenseFile, licenseFile+`.`+time.Now().Format(`20060102150405`))
		if err != nil {
			return err
		}
	}
	licenseData = &officialResp.Data.LicenseData
	b, err := com.JSONEncode(licenseData, `  `)
	if err != nil {
		b = []byte(err.Error())
	}
	err = ioutil.WriteFile(licenseFile, b, os.ModePerm)
	if err != nil {
		return fmt.Errorf(`保存授权证书失败：%v`, err)
	}
	return err
}
<|endoftext|>


<commit_before>// +build !windows

/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package envelope transforms values for storage at rest using a Envelope provider
package envelope

import (
	"context"
	"encoding/base64"
	"fmt"
	"net"
	"reflect"
	"sync"
	"testing"
	"time"

	"google.golang.org/grpc"

	"k8s.io/apimachinery/pkg/util/uuid"
	kmsapi "k8s.io/apiserver/pkg/storage/value/encrypt/envelope/v1beta1"
)

// TestKMSPluginLateStart tests the scenario where kms-plugin pod/container starts after kube-apiserver pod/container.
// Since the Dial to kms-plugin is non-blocking we expect the construction of gRPC service to succeed even when
// kms-plugin is not yet up - dialing happens in the background.
func TestKMSPluginLateStart(t *testing.T) {
	t.Parallel()
	callTimeout := 3 * time.Second
	endpoint := getSocketName()

	service, err := NewGRPCService(endpoint, callTimeout)
	if err != nil {
		t.Fatalf("failed to create envelope service, error: %v", err)
	}
	defer destroyService(service)

	time.Sleep(callTimeout / 2)
	f, err := startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.server.Stop()

	data := []byte("test data")
	_, err = service.Encrypt(data)
	if err != nil {
		t.Fatalf("failed when execute encrypt, error: %v", err)
	}
}

// TestTimeout tests behaviour of the kube-apiserver based on the supplied timeout and delayed start of kms-plugin.
func TestTimeouts(t *testing.T) {
	t.Parallel()
	var testCases = []struct {
		desc               string
		callTimeout        time.Duration
		pluginDelay        time.Duration
		kubeAPIServerDelay time.Duration
		wantErr            string
	}{
		{
			desc:        "timeout zero - expect failure when call from kube-apiserver arrives before plugin starts",
			callTimeout: 0 * time.Second,
			pluginDelay: 3 * time.Second,
			wantErr:     "rpc error: code = DeadlineExceeded desc = context deadline exceeded",
		},
		{
			desc:               "timeout zero but kms-plugin already up - still failure - zero timeout is an invalid value",
			callTimeout:        0 * time.Second,
			pluginDelay:        0 * time.Second,
			kubeAPIServerDelay: 2 * time.Second,
			wantErr:            "rpc error: code = DeadlineExceeded desc = context deadline exceeded",
		},
		{
			desc:        "timeout greater than kms-plugin delay - expect success",
			callTimeout: 6 * time.Second,
			pluginDelay: 3 * time.Second,
		},
		{
			desc:        "timeout less than kms-plugin delay - expect failure",
			callTimeout: 3 * time.Second,
			pluginDelay: 6 * time.Second,
			wantErr:     "rpc error: code = DeadlineExceeded desc = context deadline exceeded",
		},
	}

	for _, tt := range testCases {
		tt := tt
		t.Run(tt.desc, func(t *testing.T) {
			t.Parallel()
			var (
				service         Service
				err             error
				data            = []byte("test data")
				kubeAPIServerWG sync.WaitGroup
				kmsPluginWG     sync.WaitGroup
				testCompletedWG sync.WaitGroup
				socketName      = getSocketName()
			)

			testCompletedWG.Add(1)
			defer testCompletedWG.Done()

			kubeAPIServerWG.Add(1)
			go func() {
				// Simulating late start of kube-apiserver - plugin is up before kube-apiserver, if requested by the testcase.
				time.Sleep(tt.kubeAPIServerDelay)

				service, err = NewGRPCService(socketName, tt.callTimeout)
				if err != nil {
					t.Fatalf("failed to create envelope service, error: %v", err)
				}
				defer destroyService(service)
				kubeAPIServerWG.Done()
				// Keeping kube-apiserver up to process requests.
				testCompletedWG.Wait()
			}()

			kmsPluginWG.Add(1)
			go func() {
				// Simulating delayed start of kms-plugin, kube-apiserver is up before the plugin, if requested by the testcase.
				time.Sleep(tt.pluginDelay)

				f, err := startFakeKMSProvider(kmsapiVersion, socketName)
				if err != nil {
					t.Fatalf("failed to start test KMS provider server, error: %v", err)
				}
				defer f.server.Stop()
				kmsPluginWG.Done()
				// Keeping plugin up to process requests.
				testCompletedWG.Wait()
			}()

			kubeAPIServerWG.Wait()
			_, err = service.Encrypt(data)

			if err == nil && tt.wantErr != "" {
				t.Fatalf("got nil, want %s", tt.wantErr)
			}

			if err != nil && tt.wantErr == "" {
				t.Fatalf("got %q, want nil", err.Error())
			}

			// Collecting kms-plugin - allowing plugin to clean-up.
			kmsPluginWG.Wait()
		})
	}
}

// TestIntermittentConnectionLoss tests the scenario where the connection with kms-plugin is intermittently lost.
func TestIntermittentConnectionLoss(t *testing.T) {
	t.Parallel()
	var (
		wg1      sync.WaitGroup
		wg2      sync.WaitGroup
		timeout  = 30 * time.Second
		blackOut = 1 * time.Second
		data     = []byte("test data")
		endpoint = getSocketName()
	)
	// Start KMS Plugin
	f, err := startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}

	//  connect to kms plugin
	service, err := NewGRPCService(endpoint, timeout)
	if err != nil {
		t.Fatalf("failed to create envelope service, error: %v", err)
	}
	defer destroyService(service)

	_, err = service.Encrypt(data)
	if err != nil {
		t.Fatalf("failed when execute encrypt, error: %v", err)
	}
	t.Log("Connected to KMSPlugin")

	// Stop KMS Plugin - simulating connection loss
	t.Log("KMS Plugin is stopping")
	f.server.Stop()
	time.Sleep(2 * time.Second)

	wg1.Add(1)
	wg2.Add(1)
	go func() {
		defer wg2.Done()
		// Call service to encrypt data.
		t.Log("Sending encrypt request")
		wg1.Done()
		_, err := service.Encrypt(data)
		if err != nil {
			t.Fatalf("failed when executing encrypt, error: %v", err)
		}
	}()

	wg1.Wait()
	time.Sleep(blackOut)
	// Start KMS Plugin
	f, err = startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.server.Stop()
	t.Log("Restarted KMS Plugin")

	wg2.Wait()
}

func TestUnsupportedVersion(t *testing.T) {
	t.Parallel()
	ver := "invalid"
	data := []byte("test data")
	wantErr := fmt.Errorf(versionErrorf, ver, kmsapiVersion)
	endpoint := getSocketName()

	f, err := startFakeKMSProvider(ver, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %ver", err)
	}
	defer f.server.Stop()

	s, err := NewGRPCService(endpoint, 1*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer destroyService(s)

	// Encrypt
	_, err = s.Encrypt(data)
	if err == nil || err.Error() != wantErr.Error() {
		t.Errorf("got err: %ver, want: %ver", err, wantErr)
	}

	destroyService(s)

	s, err = NewGRPCService(endpoint, 1*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer destroyService(s)

	// Decrypt
	_, err = s.Decrypt(data)
	if err == nil || err.Error() != wantErr.Error() {
		t.Errorf("got err: %ver, want: %ver", err, wantErr)
	}
}

// Normal encryption and decryption operation.
func TestGRPCService(t *testing.T) {
	t.Parallel()
	// Start a test gRPC server.
	endpoint := getSocketName()
	f, err := startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.server.Stop()

	// Create the gRPC client service.
	service, err := NewGRPCService(endpoint, 1*time.Second)
	if err != nil {
		t.Fatalf("failed to create envelope service, error: %v", err)
	}
	defer destroyService(service)

	// Call service to encrypt data.
	data := []byte("test data")
	cipher, err := service.Encrypt(data)
	if err != nil {
		t.Fatalf("failed when execute encrypt, error: %v", err)
	}

	// Call service to decrypt data.
	result, err := service.Decrypt(cipher)
	if err != nil {
		t.Fatalf("failed when execute decrypt, error: %v", err)
	}

	if !reflect.DeepEqual(data, result) {
		t.Errorf("expect: %v, but: %v", data, result)
	}
}

// Normal encryption and decryption operation by multiple go-routines.
func TestGRPCServiceConcurrentAccess(t *testing.T) {
	t.Parallel()
	// Start a test gRPC server.
	endpoint := getSocketName()
	f, err := startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.server.Stop()

	// Create the gRPC client service.
	service, err := NewGRPCService(endpoint, 15*time.Second)
	if err != nil {
		t.Fatalf("failed to create envelope service, error: %v", err)
	}
	defer destroyService(service)

	var wg sync.WaitGroup
	n := 100
	wg.Add(n)
	for i := 0; i < n; i++ {
		go func() {
			defer wg.Done()
			// Call service to encrypt data.
			data := []byte("test data")
			cipher, err := service.Encrypt(data)
			if err != nil {
				t.Errorf("failed when execute encrypt, error: %v", err)
			}

			// Call service to decrypt data.
			result, err := service.Decrypt(cipher)
			if err != nil {
				t.Errorf("failed when execute decrypt, error: %v", err)
			}

			if !reflect.DeepEqual(data, result) {
				t.Errorf("expect: %v, but: %v", data, result)
			}
		}()
	}

	wg.Wait()
}

func destroyService(service Service) {
	if service != nil {
		s := service.(*gRPCService)
		s.connection.Close()
	}
}

func getSocketName() string {
	return fmt.Sprintf("unix:///@%s.sock", uuid.NewUUID())
}

// Test all those invalid configuration for KMS provider.
func TestInvalidConfiguration(t *testing.T) {
	t.Parallel()
	// Start a test gRPC server.
	f, err := startFakeKMSProvider(kmsapiVersion, getSocketName())
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.server.Stop()

	invalidConfigs := []struct {
		name       string
		apiVersion string
		endpoint   string
	}{
		{"emptyConfiguration", kmsapiVersion, ""},
		{"invalidScheme", kmsapiVersion, "tcp://localhost:6060"},
	}

	for _, testCase := range invalidConfigs {
		t.Run(testCase.name, func(t *testing.T) {
			f.apiVersion = testCase.apiVersion
			_, err := NewGRPCService(testCase.endpoint, 1*time.Second)
			if err == nil {
				t.Fatalf("should fail to create envelope service for %s.", testCase.name)
			}
		})
	}
}

// Start the gRPC server that listens on unix socket.
func startFakeKMSProvider(version, endpoint string) (*fakeKMSPlugin, error) {
	sockFile, err := parseEndpoint(endpoint)
	if err != nil {
		return nil, fmt.Errorf("failed to parse endpoint:%q, error %v", endpoint, err)
	}
	listener, err := net.Listen(unixProtocol, sockFile)
	if err != nil {
		return nil, fmt.Errorf("failed to listen on the unix socket, error: %v", err)
	}

	s := grpc.NewServer()
	f := &fakeKMSPlugin{apiVersion: version, server: s}
	kmsapi.RegisterKeyManagementServiceServer(s, f)
	go s.Serve(listener)
	return f, nil
}

// Fake gRPC sever for remote KMS provider.
// Use base64 to simulate encrypt and decrypt.
type fakeKMSPlugin struct {
	apiVersion string
	server     *grpc.Server
}

func (s *fakeKMSPlugin) Version(ctx context.Context, request *kmsapi.VersionRequest) (*kmsapi.VersionResponse, error) {
	return &kmsapi.VersionResponse{Version: s.apiVersion, RuntimeName: "testKMS", RuntimeVersion: "0.0.1"}, nil
}

func (s *fakeKMSPlugin) Decrypt(ctx context.Context, request *kmsapi.DecryptRequest) (*kmsapi.DecryptResponse, error) {
	buf := make([]byte, base64.StdEncoding.DecodedLen(len(request.Cipher)))
	n, err := base64.StdEncoding.Decode(buf, request.Cipher)
	if err != nil {
		return nil, err
	}

	return &kmsapi.DecryptResponse{Plain: buf[:n]}, nil
}

func (s *fakeKMSPlugin) Encrypt(ctx context.Context, request *kmsapi.EncryptRequest) (*kmsapi.EncryptResponse, error) {
	buf := make([]byte, base64.StdEncoding.EncodedLen(len(request.Plain)))
	base64.StdEncoding.Encode(buf, request.Plain)
	return &kmsapi.EncryptResponse{Cipher: buf}, nil
}
<commit_msg>Fix spurious .sock files running envelope unit tests<commit_after>// +build !windows

/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package envelope transforms values for storage at rest using a Envelope provider
package envelope

import (
	"context"
	"encoding/base64"
	"fmt"
	"net"
	"os"
	"reflect"
	"runtime"
	"strings"
	"sync"
	"testing"
	"time"

	"google.golang.org/grpc"

	"k8s.io/apimachinery/pkg/util/uuid"
	kmsapi "k8s.io/apiserver/pkg/storage/value/encrypt/envelope/v1beta1"
)

// TestKMSPluginLateStart tests the scenario where kms-plugin pod/container starts after kube-apiserver pod/container.
// Since the Dial to kms-plugin is non-blocking we expect the construction of gRPC service to succeed even when
// kms-plugin is not yet up - dialing happens in the background.
func TestKMSPluginLateStart(t *testing.T) {
	t.Parallel()
	callTimeout := 3 * time.Second
	endpoint := getSocketName()

	service, err := NewGRPCService(endpoint, callTimeout)
	if err != nil {
		t.Fatalf("failed to create envelope service, error: %v", err)
	}
	defer destroyService(service)

	time.Sleep(callTimeout / 2)
	f, err := startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.Stop()

	data := []byte("test data")
	_, err = service.Encrypt(data)
	if err != nil {
		t.Fatalf("failed when execute encrypt, error: %v", err)
	}
}

// TestTimeout tests behaviour of the kube-apiserver based on the supplied timeout and delayed start of kms-plugin.
func TestTimeouts(t *testing.T) {
	t.Parallel()
	var testCases = []struct {
		desc               string
		callTimeout        time.Duration
		pluginDelay        time.Duration
		kubeAPIServerDelay time.Duration
		wantErr            string
	}{
		{
			desc:        "timeout zero - expect failure when call from kube-apiserver arrives before plugin starts",
			callTimeout: 0 * time.Second,
			pluginDelay: 3 * time.Second,
			wantErr:     "rpc error: code = DeadlineExceeded desc = context deadline exceeded",
		},
		{
			desc:               "timeout zero but kms-plugin already up - still failure - zero timeout is an invalid value",
			callTimeout:        0 * time.Second,
			pluginDelay:        0 * time.Second,
			kubeAPIServerDelay: 2 * time.Second,
			wantErr:            "rpc error: code = DeadlineExceeded desc = context deadline exceeded",
		},
		{
			desc:        "timeout greater than kms-plugin delay - expect success",
			callTimeout: 6 * time.Second,
			pluginDelay: 3 * time.Second,
		},
		{
			desc:        "timeout less than kms-plugin delay - expect failure",
			callTimeout: 3 * time.Second,
			pluginDelay: 6 * time.Second,
			wantErr:     "rpc error: code = DeadlineExceeded desc = context deadline exceeded",
		},
	}

	for _, tt := range testCases {
		tt := tt
		t.Run(tt.desc, func(t *testing.T) {
			t.Parallel()
			var (
				service         Service
				err             error
				data            = []byte("test data")
				kubeAPIServerWG sync.WaitGroup
				kmsPluginWG     sync.WaitGroup
				testCompletedWG sync.WaitGroup
				socketName      = getSocketName()
			)

			testCompletedWG.Add(1)
			defer testCompletedWG.Done()

			kubeAPIServerWG.Add(1)
			go func() {
				// Simulating late start of kube-apiserver - plugin is up before kube-apiserver, if requested by the testcase.
				time.Sleep(tt.kubeAPIServerDelay)

				service, err = NewGRPCService(socketName, tt.callTimeout)
				if err != nil {
					t.Fatalf("failed to create envelope service, error: %v", err)
				}
				defer destroyService(service)
				kubeAPIServerWG.Done()
				// Keeping kube-apiserver up to process requests.
				testCompletedWG.Wait()
			}()

			kmsPluginWG.Add(1)
			go func() {
				// Simulating delayed start of kms-plugin, kube-apiserver is up before the plugin, if requested by the testcase.
				time.Sleep(tt.pluginDelay)

				f, err := startFakeKMSProvider(kmsapiVersion, socketName)
				if err != nil {
					t.Fatalf("failed to start test KMS provider server, error: %v", err)
				}
				defer f.Stop()
				kmsPluginWG.Done()
				// Keeping plugin up to process requests.
				testCompletedWG.Wait()
			}()

			kubeAPIServerWG.Wait()
			_, err = service.Encrypt(data)

			if err == nil && tt.wantErr != "" {
				t.Fatalf("got nil, want %s", tt.wantErr)
			}

			if err != nil && tt.wantErr == "" {
				t.Fatalf("got %q, want nil", err.Error())
			}

			// Collecting kms-plugin - allowing plugin to clean-up.
			kmsPluginWG.Wait()
		})
	}
}

// TestIntermittentConnectionLoss tests the scenario where the connection with kms-plugin is intermittently lost.
func TestIntermittentConnectionLoss(t *testing.T) {
	t.Parallel()
	var (
		wg1      sync.WaitGroup
		wg2      sync.WaitGroup
		timeout  = 30 * time.Second
		blackOut = 1 * time.Second
		data     = []byte("test data")
		endpoint = getSocketName()
	)
	// Start KMS Plugin
	f, err := startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}

	//  connect to kms plugin
	service, err := NewGRPCService(endpoint, timeout)
	if err != nil {
		t.Fatalf("failed to create envelope service, error: %v", err)
	}
	defer destroyService(service)

	_, err = service.Encrypt(data)
	if err != nil {
		t.Fatalf("failed when execute encrypt, error: %v", err)
	}
	t.Log("Connected to KMSPlugin")

	// Stop KMS Plugin - simulating connection loss
	t.Log("KMS Plugin is stopping")
	f.Stop()
	time.Sleep(2 * time.Second)

	wg1.Add(1)
	wg2.Add(1)
	go func() {
		defer wg2.Done()
		// Call service to encrypt data.
		t.Log("Sending encrypt request")
		wg1.Done()
		_, err := service.Encrypt(data)
		if err != nil {
			t.Fatalf("failed when executing encrypt, error: %v", err)
		}
	}()

	wg1.Wait()
	time.Sleep(blackOut)
	// Start KMS Plugin
	f, err = startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.Stop()
	t.Log("Restarted KMS Plugin")

	wg2.Wait()
}

func TestUnsupportedVersion(t *testing.T) {
	t.Parallel()
	ver := "invalid"
	data := []byte("test data")
	wantErr := fmt.Errorf(versionErrorf, ver, kmsapiVersion)
	endpoint := getSocketName()

	f, err := startFakeKMSProvider(ver, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %ver", err)
	}
	defer f.Stop()

	s, err := NewGRPCService(endpoint, 1*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer destroyService(s)

	// Encrypt
	_, err = s.Encrypt(data)
	if err == nil || err.Error() != wantErr.Error() {
		t.Errorf("got err: %ver, want: %ver", err, wantErr)
	}

	destroyService(s)

	s, err = NewGRPCService(endpoint, 1*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer destroyService(s)

	// Decrypt
	_, err = s.Decrypt(data)
	if err == nil || err.Error() != wantErr.Error() {
		t.Errorf("got err: %ver, want: %ver", err, wantErr)
	}
}

// Normal encryption and decryption operation.
func TestGRPCService(t *testing.T) {
	t.Parallel()
	// Start a test gRPC server.
	endpoint := getSocketName()
	f, err := startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.Stop()

	// Create the gRPC client service.
	service, err := NewGRPCService(endpoint, 1*time.Second)
	if err != nil {
		t.Fatalf("failed to create envelope service, error: %v", err)
	}
	defer destroyService(service)

	// Call service to encrypt data.
	data := []byte("test data")
	cipher, err := service.Encrypt(data)
	if err != nil {
		t.Fatalf("failed when execute encrypt, error: %v", err)
	}

	// Call service to decrypt data.
	result, err := service.Decrypt(cipher)
	if err != nil {
		t.Fatalf("failed when execute decrypt, error: %v", err)
	}

	if !reflect.DeepEqual(data, result) {
		t.Errorf("expect: %v, but: %v", data, result)
	}
}

// Normal encryption and decryption operation by multiple go-routines.
func TestGRPCServiceConcurrentAccess(t *testing.T) {
	t.Parallel()
	// Start a test gRPC server.
	endpoint := getSocketName()
	f, err := startFakeKMSProvider(kmsapiVersion, endpoint)
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.Stop()

	// Create the gRPC client service.
	service, err := NewGRPCService(endpoint, 15*time.Second)
	if err != nil {
		t.Fatalf("failed to create envelope service, error: %v", err)
	}
	defer destroyService(service)

	var wg sync.WaitGroup
	n := 100
	wg.Add(n)
	for i := 0; i < n; i++ {
		go func() {
			defer wg.Done()
			// Call service to encrypt data.
			data := []byte("test data")
			cipher, err := service.Encrypt(data)
			if err != nil {
				t.Errorf("failed when execute encrypt, error: %v", err)
			}

			// Call service to decrypt data.
			result, err := service.Decrypt(cipher)
			if err != nil {
				t.Errorf("failed when execute decrypt, error: %v", err)
			}

			if !reflect.DeepEqual(data, result) {
				t.Errorf("expect: %v, but: %v", data, result)
			}
		}()
	}

	wg.Wait()
}

func destroyService(service Service) {
	if service != nil {
		s := service.(*gRPCService)
		s.connection.Close()
	}
}

func getSocketName() string {
	return fmt.Sprintf("unix:///@%s.sock", uuid.NewUUID())
}

// Test all those invalid configuration for KMS provider.
func TestInvalidConfiguration(t *testing.T) {
	t.Parallel()
	// Start a test gRPC server.
	f, err := startFakeKMSProvider(kmsapiVersion, getSocketName())
	if err != nil {
		t.Fatalf("failed to start test KMS provider server, error: %v", err)
	}
	defer f.Stop()

	invalidConfigs := []struct {
		name       string
		apiVersion string
		endpoint   string
	}{
		{"emptyConfiguration", kmsapiVersion, ""},
		{"invalidScheme", kmsapiVersion, "tcp://localhost:6060"},
	}

	for _, testCase := range invalidConfigs {
		t.Run(testCase.name, func(t *testing.T) {
			f.apiVersion = testCase.apiVersion
			_, err := NewGRPCService(testCase.endpoint, 1*time.Second)
			if err == nil {
				t.Fatalf("should fail to create envelope service for %s.", testCase.name)
			}
		})
	}
}

// Start the gRPC server that listens on unix socket.
func startFakeKMSProvider(version, endpoint string) (*fakeKMSPlugin, error) {
	sockFile, err := parseEndpoint(endpoint)
	if err != nil {
		return nil, fmt.Errorf("failed to parse endpoint:%q, error %v", endpoint, err)
	}
	listener, err := net.Listen(unixProtocol, sockFile)
	if err != nil {
		return nil, fmt.Errorf("failed to listen on the unix socket, error: %v", err)
	}

	s := grpc.NewServer()
	f := &fakeKMSPlugin{apiVersion: version, server: s, sockFile: sockFile}
	kmsapi.RegisterKeyManagementServiceServer(s, f)
	go s.Serve(listener)
	return f, nil
}

// Fake gRPC sever for remote KMS provider.
// Use base64 to simulate encrypt and decrypt.
type fakeKMSPlugin struct {
	apiVersion string
	server     *grpc.Server
	sockFile   string
}

func (s *fakeKMSPlugin) Stop() {
	// Stop the server
	s.server.Stop()
	// If this isn't a Linux abstract namespace socket, or if we're on a non-linux platform, clean up the socket file
	if !strings.HasPrefix(s.sockFile, "@") || runtime.GOOS != "linux" {
		os.Remove(s.sockFile)
	}
}

func (s *fakeKMSPlugin) Version(ctx context.Context, request *kmsapi.VersionRequest) (*kmsapi.VersionResponse, error) {
	return &kmsapi.VersionResponse{Version: s.apiVersion, RuntimeName: "testKMS", RuntimeVersion: "0.0.1"}, nil
}

func (s *fakeKMSPlugin) Decrypt(ctx context.Context, request *kmsapi.DecryptRequest) (*kmsapi.DecryptResponse, error) {
	buf := make([]byte, base64.StdEncoding.DecodedLen(len(request.Cipher)))
	n, err := base64.StdEncoding.Decode(buf, request.Cipher)
	if err != nil {
		return nil, err
	}

	return &kmsapi.DecryptResponse{Plain: buf[:n]}, nil
}

func (s *fakeKMSPlugin) Encrypt(ctx context.Context, request *kmsapi.EncryptRequest) (*kmsapi.EncryptResponse, error) {
	buf := make([]byte, base64.StdEncoding.EncodedLen(len(request.Plain)))
	base64.StdEncoding.Encode(buf, request.Plain)
	return &kmsapi.EncryptResponse{Cipher: buf}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"image"
	"image/color"
	"math"
)

type Canvas struct {
	image.RGBA
}

func NewCanvas(r image.Rectangle) *Canvas {
	canvas := new(Canvas)
	canvas.RGBA = *image.NewRGBA(r)
	return canvas
}

func (c Canvas) Clone() *Canvas {
	clone := NewCanvas(c.Bounds())
	copy(clone.Pix, c.Pix)
	return clone
}

func (c Canvas) DrawGradient() {
	size := c.Bounds().Size()
	for x := 0; x < size.X; x++ {
		for y := 0; y < size.Y; y++ {
			color := color.RGBA{
				uint8(255 * x / size.X),
				uint8(255 * y / size.Y),
				55,
				255}
			c.Set(x, y, color)
		}
	}
}

func (c Canvas) DrawLine(color color.RGBA, from Vector, to Vector) {
	delta := to.Sub(from)
	length := delta.Length()
	x_step, y_step := delta.X/length, delta.Y/length
	limit := int(length + 0.5)
	for i := 0; i < limit; i++ {
		x := from.X + float64(i)*x_step
		y := from.Y + float64(i)*y_step
		c.Set(int(x), int(y), color)
	}
}

func (c Canvas) DrawSpiral(color color.RGBA, from Vector) {
	dir := Vector{0, 2}
	last := from
	for i := 0; i < 10000; i++ {
		next := last.Add(dir)
		c.DrawLine(color, last, next)
		dir.Rotate(0.03)
		dir.Scale(0.999)
		last = next
	}
}

func (c Canvas) Blur(radius int, weight WeightFunction) {
	clone := c.Clone()
	size := c.Bounds().Size()
	for x := 0; x < size.X; x++ {
		for y := 0; y < size.Y; y++ {
			color := c.BlurPixel(x, y, radius, weight)
			clone.Set(x, y, color)
		}
	}
	copy(c.Pix, clone.Pix)
}

func (c Canvas) BlurPixel(x int, y int, radius int, weight WeightFunction) color.Color {
	weightSum := float64(0)
	size := c.Bounds().Size()
	outR, outG, outB := float64(0), float64(0), float64(0)
	for i := x - radius; i < x+radius+1; i++ {
		if i < 0 || i > size.X {
			continue
		}
		for j := y - radius; j < y+radius+1; j++ {
			if j < 0 || j > size.Y {
				continue
			}
			weight := weight.Weight(i-x, j-y)
			r, g, b, _ := c.At(i, j).RGBA()
			outR += float64(r) * weight
			outG += float64(g) * weight
			outB += float64(b) * weight
			weightSum += weight
		}
	}
	// Need to divide by 0xFF as the RGBA() function returns color values as uint32
	// and we need uint8
	return color.RGBA{
		uint8(outR / (weightSum * 0xFF)),
		uint8(outG / (weightSum * 0xFF)),
		uint8(outB / (weightSum * 0xFF)),
		255}
}

// Blur weighting functions
type WeightFunction interface {
	Weight(x int, y int) float64
}

type WeightFunctionBox struct{}

func (w WeightFunctionBox) Weight(x int, y int) float64 { return 1.0 }

type WeightFunctionDist struct{}

func (w WeightFunctionDist) Weight(x int, y int) float64 {
	d := math.Hypot(float64(x), float64(y))
	return 1 / (1 + d)
}

type WeightFunctionMotion struct {
}

func (w WeightFunctionMotion) Weight(x int, y int) float64 {
	if y != 0 {
		return 0
	}
	if x < 0 {
		return 0
	}
	return 1 / (1 + float64(x))
}
<commit_msg>Add double blur filter<commit_after>package main

import (
	"image"
	"image/color"
	"math"
)

type Canvas struct {
	image.RGBA
}

func NewCanvas(r image.Rectangle) *Canvas {
	canvas := new(Canvas)
	canvas.RGBA = *image.NewRGBA(r)
	return canvas
}

func (c Canvas) Clone() *Canvas {
	clone := NewCanvas(c.Bounds())
	copy(clone.Pix, c.Pix)
	return clone
}

func (c Canvas) DrawGradient() {
	size := c.Bounds().Size()
	for x := 0; x < size.X; x++ {
		for y := 0; y < size.Y; y++ {
			color := color.RGBA{
				uint8(255 * x / size.X),
				uint8(255 * y / size.Y),
				55,
				255}
			c.Set(x, y, color)
		}
	}
}

func (c Canvas) DrawLine(color color.RGBA, from Vector, to Vector) {
	delta := to.Sub(from)
	length := delta.Length()
	x_step, y_step := delta.X/length, delta.Y/length
	limit := int(length + 0.5)
	for i := 0; i < limit; i++ {
		x := from.X + float64(i)*x_step
		y := from.Y + float64(i)*y_step
		c.Set(int(x), int(y), color)
	}
}

func (c Canvas) DrawSpiral(color color.RGBA, from Vector) {
	dir := Vector{0, 2}
	last := from
	for i := 0; i < 10000; i++ {
		next := last.Add(dir)
		c.DrawLine(color, last, next)
		dir.Rotate(0.03)
		dir.Scale(0.999)
		last = next
	}
}

func (c Canvas) Blur(radius int, weight WeightFunction) {
	clone := c.Clone()
	size := c.Bounds().Size()
	for x := 0; x < size.X; x++ {
		for y := 0; y < size.Y; y++ {
			color := c.BlurPixel(x, y, radius, weight)
			clone.Set(x, y, color)
		}
	}
	copy(c.Pix, clone.Pix)
}

func (c Canvas) BlurPixel(x int, y int, radius int, weight WeightFunction) color.Color {
	weightSum := float64(0)
	size := c.Bounds().Size()
	outR, outG, outB := float64(0), float64(0), float64(0)
	for i := x - radius; i < x+radius+1; i++ {
		if i < 0 || i > size.X {
			continue
		}
		for j := y - radius; j < y+radius+1; j++ {
			if j < 0 || j > size.Y {
				continue
			}
			weight := weight.Weight(i-x, j-y)
			r, g, b, _ := c.At(i, j).RGBA()
			outR += float64(r) * weight
			outG += float64(g) * weight
			outB += float64(b) * weight
			weightSum += weight
		}
	}
	// Need to divide by 0xFF as the RGBA() function returns color values as uint32
	// and we need uint8
	return color.RGBA{
		uint8(outR / (weightSum * 0xFF)),
		uint8(outG / (weightSum * 0xFF)),
		uint8(outB / (weightSum * 0xFF)),
		255}
}

// Blur weighting functions
type WeightFunction interface {
	Weight(x int, y int) float64
}

type WeightFunctionBox struct{}

func (w WeightFunctionBox) Weight(x int, y int) float64 { return 1.0 }

type WeightFunctionDist struct{}

func (w WeightFunctionDist) Weight(x int, y int) float64 {
	d := math.Hypot(float64(x), float64(y))
	return 1 / (1 + d)
}

type WeightFunctionMotion struct {
}

func (w WeightFunctionMotion) Weight(x int, y int) float64 {
	if y != 0 {
		return 0
	}
	if x < 0 {
		return 0
	}
	return 0.3 + 0.7/math.Sqrt(1+float64(x))
}

type WeightFunctionDouble struct {
	split int
}

func (w WeightFunctionDouble) Weight(x int, y int) float64 {
	if y == 0 && (x == w.split || x == -w.split) {
		return 1.0
	} else {
		return 0
	}
}
<|endoftext|>


<commit_before>package survey

import (
	"errors"
	"fmt"
	tm "github.com/buger/goterm"
	"strings"

	"github.com/alecaivazis/survey/format"
)

// Choice is a prompt that presents a
type Choice struct {
	Message string
	Choices []string
	Default string
}

func (prompt *Choice) Prompt() (string, error) {
	// ask the question
	fmt.Println(format.Ask(prompt.Message, ""))

	// get the current location of the cursor
	loc, err := CursorLocation()
	// if something went wrong
	if err != nil {
		// bubble up
		return "", err
	}

	// the height of the prompt's output
	height := len(prompt.Choices)

	// the starting point of the list depends on wether or not we
	// are at the bottom of the current terminal session
	var initialLocation int
	// if the options would fit cleanly
	if loc.col+height < tm.Height() {
		// start at the current location
		initialLocation = loc.col
		// otherwise we will be placed at the bottom of the terminal after this print
	} else {
		// the we have to start printing so that we just fit
		initialLocation = tm.Height() - height
	}

	// start off with the first option selected
	sel := 0
	// if there is a default
	if prompt.Default != "" {
		// find the choice
		for i, opt := range prompt.Choices {
			// if the option correponds to the default
			if opt == prompt.Default {
				// we found our initial value
				sel = i
				// stop looking
				break
			}
		}
	}

	// print the options to start
	refreshOptions(prompt.Choices, sel, initialLocation)

	for {
		// wait for an input from the user
		ascii, keycode, err := getChar()
		// if there is an error
		if err != nil {
			// bubble up
			return "", err
		}

		// if the user sends SIGTERM (ascii 3) or presses esc (ascii 27)
		if ascii == 3 || ascii == 27 {
			// hard close
			return "", errors.New("Goodbye.")
		}

		// if the user pressed the up arrow (keycode 38) and we can decrement sel
		if keycode == 38 && sel > 0 {
			// decrement the selected index
			sel--
		}
		// if the user pressed the down arrow (keycode 40)  and we can decrement sel
		if keycode == 40 && sel < len(prompt.Choices)-1 {
			// decrement the selected index
			sel++
		}

		// if the user presses enter (ascii 13)
		if ascii == 13 {
			// we're done with the rendering loop (the current value is good)
			break
		}

		// print the options
		refreshOptions(prompt.Choices, sel, initialLocation)
	}

	// return the selected choice
	return prompt.Choices[sel], nil
}

// Cleanup removes the choices section, and renders the ask like a normal question.
func (prompt *Choice) Cleanup(val string) error {

	// the height of the prompt's output
	height := len(prompt.Choices)

	// get the current location of the cursor
	loc, err := CursorLocation()
	// if something went wrong
	if err != nil {
		// yell loudly
		return err
	}

	var initLoc int
	// if the options would fit cleanly
	if loc.col+height <= tm.Height() {
		// start at the current location
		initLoc = loc.col - height - 1
		// otherwise we will be placed at the bottom of the terminal after this print
	} else {

		// the we have to start printing so that we just fit
		initLoc = loc.col - height - 2
	}

	// find the index of the selected choice

	// start where we were told
	tm.MoveCursor(initLoc, 1)
	tm.Print(format.Response(prompt.Message, val), "\x1b[0K")
	// for each choice
	for range prompt.Choices {
		// add an empty line
		tm.Print(format.AnsiClearLine)
		// print the output
		tm.Flush()
	}
	// add an empty line
	tm.Print(format.AnsiClearLine)
	// print the output
	tm.Flush()
	tm.MoveCursor(initLoc, 1)
	tm.Flush()

	// nothing went wrong
	return nil
}

func refreshOptions(opts []string, sel int, initLoc int) {
	// we need to render the options
	tm.Print(formatChoiceOptions(opts, sel))
	tm.Flush()
	// make sure we overwrite the first line next time we print
	tm.MoveCursor(initLoc, 1)
}

func formatChoiceOptions(opts []string, selected int) string {
	// a string to acc
	acc := []string{}
	// format each option
	for i, opt := range opts {
		// by default, the option is not selected
		isSel := false
		// if this option is at the same index as the selected value
		if i == selected {
			// then the option should show a selection indicator
			isSel = true
		}

		// add the formatted option
		acc = append(acc, format.ChoiceOption(opt, isSel))
	}

	// show each option on its own line
	return strings.Join(acc, "\n")
}
<commit_msg>added comment to choice.Prompt<commit_after>package survey

import (
	"errors"
	"fmt"
	tm "github.com/buger/goterm"
	"strings"

	"github.com/alecaivazis/survey/format"
)

// Choice is a prompt that presents a list of various options to the user
// for them to select using the arrow keys and enter.
type Choice struct {
	Message string
	Choices []string
	Default string
}

// Prompt shows the list, and listens for input from the user using /dev/tty.
func (prompt *Choice) Prompt() (string, error) {
	// ask the question
	fmt.Println(format.Ask(prompt.Message, ""))

	// get the current location of the cursor
	loc, err := CursorLocation()
	// if something went wrong
	if err != nil {
		// bubble up
		return "", err
	}

	// the height of the prompt's output
	height := len(prompt.Choices)

	// the starting point of the list depends on wether or not we
	// are at the bottom of the current terminal session
	var initialLocation int
	// if the options would fit cleanly
	if loc.col+height < tm.Height() {
		// start at the current location
		initialLocation = loc.col
		// otherwise we will be placed at the bottom of the terminal after this print
	} else {
		// the we have to start printing so that we just fit
		initialLocation = tm.Height() - height
	}

	// start off with the first option selected
	sel := 0
	// if there is a default
	if prompt.Default != "" {
		// find the choice
		for i, opt := range prompt.Choices {
			// if the option correponds to the default
			if opt == prompt.Default {
				// we found our initial value
				sel = i
				// stop looking
				break
			}
		}
	}

	// print the options to start
	refreshOptions(prompt.Choices, sel, initialLocation)

	for {
		// wait for an input from the user
		ascii, keycode, err := getChar()
		// if there is an error
		if err != nil {
			// bubble up
			return "", err
		}

		// if the user sends SIGTERM (ascii 3) or presses esc (ascii 27)
		if ascii == 3 || ascii == 27 {
			// hard close
			return "", errors.New("Goodbye.")
		}

		// if the user pressed the up arrow (keycode 38) and we can decrement sel
		if keycode == 38 && sel > 0 {
			// decrement the selected index
			sel--
		}
		// if the user pressed the down arrow (keycode 40)  and we can decrement sel
		if keycode == 40 && sel < len(prompt.Choices)-1 {
			// decrement the selected index
			sel++
		}

		// if the user presses enter (ascii 13)
		if ascii == 13 {
			// we're done with the rendering loop (the current value is good)
			break
		}

		// print the options
		refreshOptions(prompt.Choices, sel, initialLocation)
	}

	// return the selected choice
	return prompt.Choices[sel], nil
}

// Cleanup removes the choices section, and renders the ask like a normal question.
func (prompt *Choice) Cleanup(val string) error {

	// the height of the prompt's output
	height := len(prompt.Choices)

	// get the current location of the cursor
	loc, err := CursorLocation()
	// if something went wrong
	if err != nil {
		// yell loudly
		return err
	}

	var initLoc int
	// if the options would fit cleanly
	if loc.col+height <= tm.Height() {
		// start at the current location
		initLoc = loc.col - height - 1
		// otherwise we will be placed at the bottom of the terminal after this print
	} else {

		// the we have to start printing so that we just fit
		initLoc = loc.col - height - 2
	}

	// find the index of the selected choice

	// start where we were told
	tm.MoveCursor(initLoc, 1)
	tm.Print(format.Response(prompt.Message, val), "\x1b[0K")
	// for each choice
	for range prompt.Choices {
		// add an empty line
		tm.Print(format.AnsiClearLine)
		// print the output
		tm.Flush()
	}
	// add an empty line
	tm.Print(format.AnsiClearLine)
	// print the output
	tm.Flush()
	tm.MoveCursor(initLoc, 1)
	tm.Flush()

	// nothing went wrong
	return nil
}

func refreshOptions(opts []string, sel int, initLoc int) {
	// we need to render the options
	tm.Print(formatChoiceOptions(opts, sel))
	tm.Flush()
	// make sure we overwrite the first line next time we print
	tm.MoveCursor(initLoc, 1)
}

func formatChoiceOptions(opts []string, selected int) string {
	// a string to acc
	acc := []string{}
	// format each option
	for i, opt := range opts {
		// by default, the option is not selected
		isSel := false
		// if this option is at the same index as the selected value
		if i == selected {
			// then the option should show a selection indicator
			isSel = true
		}

		// add the formatted option
		acc = append(acc, format.ChoiceOption(opt, isSel))
	}

	// show each option on its own line
	return strings.Join(acc, "\n")
}
<|endoftext|>


<commit_before>package putio

import (
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
)

const (
	defaultUserAgent = "go-putio"
	defaultMediaType = "application/json"
	defaultBaseURL   = "https://api.put.io"
	defaultUploadURL = "https://upload.put.io"
)

// Client manages communication with Put.io v2 API.
type Client struct {
	// HTTP client used to communicate with Put.io API
	client *http.Client

	// Base URL for API requests
	BaseURL *url.URL

	// base url for upload requests
	uploadURL *url.URL

	// User agent for client
	UserAgent string

	// Override host header for API requests
	Host string

	// ExtraHeaders are passed to the API server on every request.
	ExtraHeaders http.Header

	// Services used for communicating with the API
	Account   *AccountService
	Files     *FilesService
	Transfers *TransfersService
	Zips      *ZipsService
	Friends   *FriendsService
	Events    *EventsService
}

// NewClient returns a new Put.io API client, using the htttpClient, which must
// be a new Oauth2 enabled http.Client. If httpClient is not defined, default
// HTTP client is used.
func NewClient(httpClient *http.Client) *Client {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}

	baseURL, _ := url.Parse(defaultBaseURL)
	uploadURL, _ := url.Parse(defaultUploadURL)
	c := &Client{
		client:       httpClient,
		BaseURL:      baseURL,
		uploadURL:    uploadURL,
		UserAgent:    defaultUserAgent,
		ExtraHeaders: make(http.Header),
	}

	c.Account = &AccountService{client: c}
	c.Files = &FilesService{client: c}
	c.Transfers = &TransfersService{client: c}
	c.Zips = &ZipsService{client: c}
	c.Friends = &FriendsService{client: c}
	c.Events = &EventsService{client: c}

	return c
}

func (c *Client) ValidateToken(ctx context.Context) (userID *int64, err error) {
	req, err := c.NewRequest(ctx, "GET", "/v2/oauth2/validate", nil)
	if err != nil {
		return
	}
	var r struct {
		UserID *int64 `json:"user_id"`
	}
	resp, err := c.Do(req, &r)
	defer resp.Body.Close()
	return r.UserID, err
}

// NewRequest creates an API request. A relative URL can be provided via
// relURL, which will be resolved to the BaseURL of the Client.
func (c *Client) NewRequest(ctx context.Context, method, relURL string, body io.Reader) (*http.Request, error) {
	rel, err := url.Parse(relURL)
	if err != nil {
		return nil, err
	}

	var u *url.URL
	// XXX: workaroud for upload endpoint. upload method has a different base url,
	// so we've a special case for testing purposes.
	if relURL == "/v2/files/upload" {
		u = c.uploadURL.ResolveReference(rel)
	} else {
		u = c.BaseURL.ResolveReference(rel)
	}

	req, err := http.NewRequest(method, u.String(), body)
	if err != nil {
		return nil, err
	}

	req = req.WithContext(ctx)
	req.Header.Set("Accept", defaultMediaType)
	req.Header.Set("User-Agent", c.UserAgent)

	if c.Host != "" {
		req.Host = c.Host
	}

	// merge headers with extra headers
	for header, values := range c.ExtraHeaders {
		for _, value := range values {
			req.Header.Add(header, value)
		}
	}

	return req, nil
}

// Do sends an API request and returns the API response. The API response is
// JSON decoded and stored in the value pointed to by v, or returned as an
// error if an API error has occurred. Response body is closed at all cases except
// v is nil. If v is nil, response body is not closed and the body can be used
// for streaming.
func (c *Client) Do(r *http.Request, v interface{}) (*http.Response, error) {
	resp, err := c.client.Do(r)
	if err != nil {
		return nil, err
	}

	err = checkResponse(resp)
	if err != nil {
		// close the body at all times if there is an http error
		resp.Body.Close()
		return resp, err
	}

	if v == nil {
		return resp, nil
	}

	// close the body for all cases from here
	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(v)
	if err != nil {
		return resp, err
	}

	return resp, nil
}

// checkResponse is the entrypoint to reading the API response. If the response
// status code is not in success range, it will try to return a structured
// error.
func checkResponse(r *http.Response) error {
	if r.StatusCode >= 200 && r.StatusCode <= 399 {
		return nil
	}

	// server possibly returns json and more details
	er := &ErrorResponse{Response: r}

	er.Body, er.ParseError = ioutil.ReadAll(r.Body)
	if er.ParseError != nil {
		return er
	}
	if r.Header.Get("content-type") == "application/json" {
		er.ParseError = json.Unmarshal(er.Body, er)
		if er.ParseError != nil {
			return er
		}
	}
	return er
}
<commit_msg>check error before closing body<commit_after>package putio

import (
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
)

const (
	defaultUserAgent = "go-putio"
	defaultMediaType = "application/json"
	defaultBaseURL   = "https://api.put.io"
	defaultUploadURL = "https://upload.put.io"
)

// Client manages communication with Put.io v2 API.
type Client struct {
	// HTTP client used to communicate with Put.io API
	client *http.Client

	// Base URL for API requests
	BaseURL *url.URL

	// base url for upload requests
	uploadURL *url.URL

	// User agent for client
	UserAgent string

	// Override host header for API requests
	Host string

	// ExtraHeaders are passed to the API server on every request.
	ExtraHeaders http.Header

	// Services used for communicating with the API
	Account   *AccountService
	Files     *FilesService
	Transfers *TransfersService
	Zips      *ZipsService
	Friends   *FriendsService
	Events    *EventsService
}

// NewClient returns a new Put.io API client, using the htttpClient, which must
// be a new Oauth2 enabled http.Client. If httpClient is not defined, default
// HTTP client is used.
func NewClient(httpClient *http.Client) *Client {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}

	baseURL, _ := url.Parse(defaultBaseURL)
	uploadURL, _ := url.Parse(defaultUploadURL)
	c := &Client{
		client:       httpClient,
		BaseURL:      baseURL,
		uploadURL:    uploadURL,
		UserAgent:    defaultUserAgent,
		ExtraHeaders: make(http.Header),
	}

	c.Account = &AccountService{client: c}
	c.Files = &FilesService{client: c}
	c.Transfers = &TransfersService{client: c}
	c.Zips = &ZipsService{client: c}
	c.Friends = &FriendsService{client: c}
	c.Events = &EventsService{client: c}

	return c
}

func (c *Client) ValidateToken(ctx context.Context) (userID *int64, err error) {
	req, err := c.NewRequest(ctx, "GET", "/v2/oauth2/validate", nil)
	if err != nil {
		return
	}
	var r struct {
		UserID *int64 `json:"user_id"`
	}
	_, err = c.Do(req, &r)
	if err != nil {
		return nil, err
	}
	return r.UserID, nil
}

// NewRequest creates an API request. A relative URL can be provided via
// relURL, which will be resolved to the BaseURL of the Client.
func (c *Client) NewRequest(ctx context.Context, method, relURL string, body io.Reader) (*http.Request, error) {
	rel, err := url.Parse(relURL)
	if err != nil {
		return nil, err
	}

	var u *url.URL
	// XXX: workaroud for upload endpoint. upload method has a different base url,
	// so we've a special case for testing purposes.
	if relURL == "/v2/files/upload" {
		u = c.uploadURL.ResolveReference(rel)
	} else {
		u = c.BaseURL.ResolveReference(rel)
	}

	req, err := http.NewRequest(method, u.String(), body)
	if err != nil {
		return nil, err
	}

	req = req.WithContext(ctx)
	req.Header.Set("Accept", defaultMediaType)
	req.Header.Set("User-Agent", c.UserAgent)

	if c.Host != "" {
		req.Host = c.Host
	}

	// merge headers with extra headers
	for header, values := range c.ExtraHeaders {
		for _, value := range values {
			req.Header.Add(header, value)
		}
	}

	return req, nil
}

// Do sends an API request and returns the API response. The API response is
// JSON decoded and stored in the value pointed to by v, or returned as an
// error if an API error has occurred. Response body is closed at all cases except
// v is nil. If v is nil, response body is not closed and the body can be used
// for streaming.
func (c *Client) Do(r *http.Request, v interface{}) (*http.Response, error) {
	resp, err := c.client.Do(r)
	if err != nil {
		return nil, err
	}

	err = checkResponse(resp)
	if err != nil {
		// close the body at all times if there is an http error
		resp.Body.Close()
		return resp, err
	}

	if v == nil {
		return resp, nil
	}

	// close the body for all cases from here
	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(v)
	if err != nil {
		return resp, err
	}

	return resp, nil
}

// checkResponse is the entrypoint to reading the API response. If the response
// status code is not in success range, it will try to return a structured
// error.
func checkResponse(r *http.Response) error {
	if r.StatusCode >= 200 && r.StatusCode <= 399 {
		return nil
	}

	// server possibly returns json and more details
	er := &ErrorResponse{Response: r}

	er.Body, er.ParseError = ioutil.ReadAll(r.Body)
	if er.ParseError != nil {
		return er
	}
	if r.Header.Get("content-type") == "application/json" {
		er.ParseError = json.Unmarshal(er.Body, er)
		if er.ParseError != nil {
			return er
		}
	}
	return er
}
<|endoftext|>