<commit_before>package tftest

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/hashicorp/terraform-exec/tfexec"
	tfjson "github.com/hashicorp/terraform-json"
)

// WorkingDir represents a distinct working directory that can be used for
// running tests. Each test should construct its own WorkingDir by calling
// NewWorkingDir or RequireNewWorkingDir on its package's singleton
// tftest.Helper.
type WorkingDir struct {
	h *Helper

	// baseDir is the root of the working directory tree
	baseDir string

	// baseArgs is arguments that should be appended to all commands
	baseArgs []string

	// configDir contains the singular config file generated for each test
	configDir string

	// tf is the instance of tfexec.Terraform used for running Terraform commands
	tf *tfexec.Terraform

	// terraformExec is a path to a terraform binary, inherited from Helper
	terraformExec string

	env map[string]string
}

// Close deletes the directories and files created to represent the receiving
// working directory. After this method is called, the working directory object
// is invalid and may no longer be used.
func (wd *WorkingDir) Close() error {
	return os.RemoveAll(wd.baseDir)
}

// Setenv sets an environment variable on the WorkingDir.
func (wd *WorkingDir) Setenv(envVar, val string) {
	if wd.env == nil {
		wd.env = map[string]string{}
	}
	wd.env[envVar] = val
}

// Unsetenv removes an environment variable from the WorkingDir.
func (wd *WorkingDir) Unsetenv(envVar string) {
	delete(wd.env, envVar)
}

// GetHelper returns the Helper set on the WorkingDir.
func (wd *WorkingDir) GetHelper() *Helper {
	return wd.h
}

func (wd *WorkingDir) relativeConfigDir() (string, error) {
	relPath, err := filepath.Rel(wd.baseDir, wd.configDir)
	if err != nil {
		return "", fmt.Errorf("Error determining relative path of configuration directory: %w", err)
	}
	return relPath, nil
}

// SetConfig sets a new configuration for the working directory.
//
// This must be called at least once before any call to Init, Plan, Apply, or
// Destroy to establish the configuration. Any previously-set configuration is
// discarded and any saved plan is cleared.
func (wd *WorkingDir) SetConfig(cfg string) error {
	// Each call to SetConfig creates a new directory under our baseDir.
	// We create them within so that our final cleanup step will delete them
	// automatically without any additional tracking.
	configDir, err := ioutil.TempDir(wd.baseDir, "config")
	if err != nil {
		return err
	}
	configFilename := filepath.Join(configDir, "terraform_plugin_test.tf")
	err = ioutil.WriteFile(configFilename, []byte(cfg), 0700)
	if err != nil {
		return err
	}

	tf, err := tfexec.NewTerraform(wd.baseDir, wd.terraformExec)
	if err != nil {
		return err
	}

	var mismatch *tfexec.ErrVersionMismatch
	err = tf.SetDisablePluginTLS(true)
	if err != nil && !errors.As(err, &mismatch) {
		return err
	}
	err = tf.SetSkipProviderVerify(true)
	if err != nil && !errors.As(err, &mismatch) {
		return err
	}

	if p := os.Getenv("TF_ACC_LOG_PATH"); p != "" {
		tf.SetLogPath(p)
	}

	wd.configDir = configDir
	wd.tf = tf

	// Changing configuration invalidates any saved plan.
	err = wd.ClearPlan()
	if err != nil {
		return err
	}
	return nil
}

// RequireSetConfig is a variant of SetConfig that will fail the test via the
// given TestControl if the configuration cannot be set.
func (wd *WorkingDir) RequireSetConfig(t TestControl, cfg string) {
	t.Helper()
	if err := wd.SetConfig(cfg); err != nil {
		t := testingT{t}
		t.Fatalf("failed to set config: %s", err)
	}
}

// ClearState deletes any Terraform state present in the working directory.
//
// Any remote objects tracked by the state are not destroyed first, so this
// will leave them dangling in the remote system.
func (wd *WorkingDir) ClearState() error {
	err := os.Remove(filepath.Join(wd.baseDir, "terraform.tfstate"))
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// RequireClearState is a variant of ClearState that will fail the test via the
// given TestControl if the state cannot be cleared.
func (wd *WorkingDir) RequireClearState(t TestControl) {
	t.Helper()
	if err := wd.ClearState(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to clear state: %s", err)
	}
}

// ClearPlan deletes any saved plan present in the working directory.
func (wd *WorkingDir) ClearPlan() error {
	err := os.Remove(wd.planFilename())
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// RequireClearPlan is a variant of ClearPlan that will fail the test via the
// given TestControl if the plan cannot be cleared.
func (wd *WorkingDir) RequireClearPlan(t TestControl) {
	t.Helper()
	if err := wd.ClearPlan(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to clear plan: %s", err)
	}
}

// Init runs "terraform init" for the given working directory, forcing Terraform
// to use the current version of the plugin under test.
func (wd *WorkingDir) Init() error {
	if wd.configDir == "" {
		return fmt.Errorf("must call SetConfig before Init")
	}

	return wd.tf.Init(context.Background(), tfexec.Dir(wd.configDir))
}

// RequireInit is a variant of Init that will fail the test via the given
// TestControl if init fails.
func (wd *WorkingDir) RequireInit(t TestControl) {
	t.Helper()
	if err := wd.Init(); err != nil {
		t := testingT{t}
		t.Fatalf("init failed: %s", err)
	}
}

func (wd *WorkingDir) planFilename() string {
	return filepath.Join(wd.baseDir, "tfplan")
}

// CreatePlan runs "terraform plan" to create a saved plan file, which if successful
// will then be used for the next call to Apply.
func (wd *WorkingDir) CreatePlan() error {
	return wd.tf.Plan(context.Background(), tfexec.Refresh(false), tfexec.Out("tfplan"), tfexec.Dir(wd.configDir))
}

// RequireCreatePlan is a variant of CreatePlan that will fail the test via
// the given TestControl if plan creation fails.
func (wd *WorkingDir) RequireCreatePlan(t TestControl) {
	t.Helper()
	if err := wd.CreatePlan(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to create plan: %s", err)
	}
}

// CreateDestroyPlan runs "terraform plan -destroy" to create a saved plan
// file, which if successful will then be used for the next call to Apply.
func (wd *WorkingDir) CreateDestroyPlan() error {
	return wd.tf.Plan(context.Background(), tfexec.Refresh(false), tfexec.Out("tfplan"), tfexec.Destroy(true), tfexec.Dir(wd.configDir))
}

// Apply runs "terraform apply". If CreatePlan has previously completed
// successfully and the saved plan has not been cleared in the meantime then
// this will apply the saved plan. Otherwise, it will implicitly create a new
// plan and apply it.
func (wd *WorkingDir) Apply() error {
	args := []tfexec.ApplyOption{tfexec.Refresh(false)}
	if wd.HasSavedPlan() {
		args = append(args, tfexec.DirOrPlan("tfplan"))
	} else {
		// we need to use a relative config dir here or we get an
		// error about Terraform not having any configuration. See
		// https://github.com/hashicorp/terraform-plugin-sdk/issues/495
		// for more info.
		configDir, err := wd.relativeConfigDir()
		if err != nil {
			return err
		}
		args = append(args, tfexec.DirOrPlan(configDir))
	}
	return wd.tf.Apply(context.Background(), args...)
}

// RequireApply is a variant of Apply that will fail the test via
// the given TestControl if the apply operation fails.
func (wd *WorkingDir) RequireApply(t TestControl) {
	t.Helper()
	if err := wd.Apply(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to apply: %s", err)
	}
}

// Destroy runs "terraform destroy". It does not consider or modify any saved
// plan, and is primarily for cleaning up at the end of a test run.
//
// If destroy fails then remote objects might still exist, and continue to
// exist after a particular test is concluded.
func (wd *WorkingDir) Destroy() error {
	return wd.tf.Destroy(context.Background(), tfexec.Refresh(false), tfexec.Dir(wd.configDir))
}

// RequireDestroy is a variant of Destroy that will fail the test via
// the given TestControl if the destroy operation fails.
//
// If destroy fails then remote objects might still exist, and continue to
// exist after a particular test is concluded.
func (wd *WorkingDir) RequireDestroy(t TestControl) {
	t.Helper()
	if err := wd.Destroy(); err != nil {
		t := testingT{t}
		t.Logf("WARNING: destroy failed, so remote objects may still exist and be subject to billing")
		t.Fatalf("failed to destroy: %s", err)
	}
}

// HasSavedPlan returns true if there is a saved plan in the working directory. If
// so, a subsequent call to Apply will apply that saved plan.
func (wd *WorkingDir) HasSavedPlan() bool {
	_, err := os.Stat(wd.planFilename())
	return err == nil
}

// SavedPlan returns an object describing the current saved plan file, if any.
//
// If no plan is saved or if the plan file cannot be read, SavedPlan returns
// an error.
func (wd *WorkingDir) SavedPlan() (*tfjson.Plan, error) {
	if !wd.HasSavedPlan() {
		return nil, fmt.Errorf("there is no current saved plan")
	}

	return wd.tf.ShowPlanFile(context.Background(), wd.planFilename())
}

// RequireSavedPlan is a variant of SavedPlan that will fail the test via
// the given TestControl if the plan cannot be read.
func (wd *WorkingDir) RequireSavedPlan(t TestControl) *tfjson.Plan {
	t.Helper()
	ret, err := wd.SavedPlan()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read saved plan: %s", err)
	}
	return ret
}

// SavedPlanStdout returns a stdout capture of the current saved plan file, if any.
//
// If no plan is saved or if the plan file cannot be read, SavedPlanStdout returns
// an error.
func (wd *WorkingDir) SavedPlanStdout() (string, error) {
	if !wd.HasSavedPlan() {
		return "", fmt.Errorf("there is no current saved plan")
	}

	var ret bytes.Buffer

	wd.tf.SetStdout(&ret)
	defer wd.tf.SetStdout(ioutil.Discard)
	_, err := wd.tf.ShowPlanFile(context.Background(), wd.planFilename())
	if err != nil {
		return "", err
	}

	return ret.String(), nil
}

// RequireSavedPlanStdout is a variant of SavedPlanStdout that will fail the test via
// the given TestControl if the plan cannot be read.
func (wd *WorkingDir) RequireSavedPlanStdout(t TestControl) string {
	t.Helper()
	ret, err := wd.SavedPlanStdout()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read saved plan: %s", err)
	}
	return ret
}

// State returns an object describing the current state.
//
// If the state cannot be read, State returns an error.
func (wd *WorkingDir) State() (*tfjson.State, error) {
	return wd.tf.Show(context.Background())
}

// RequireState is a variant of State that will fail the test via
// the given TestControl if the state cannot be read.
func (wd *WorkingDir) RequireState(t TestControl) *tfjson.State {
	t.Helper()
	ret, err := wd.State()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read state plan: %s", err)
	}
	return ret
}

// Import runs terraform import
func (wd *WorkingDir) Import(resource, id string) error {
	return wd.tf.Import(context.Background(), resource, id, tfexec.Config(wd.configDir))
}

// RequireImport is a variant of Import that will fail the test via
// the given TestControl if the import is non successful.
func (wd *WorkingDir) RequireImport(t TestControl, resource, id string) {
	t.Helper()
	if err := wd.Import(resource, id); err != nil {
		t := testingT{t}
		t.Fatalf("failed to import: %s", err)
	}
}

// Refresh runs terraform refresh
func (wd *WorkingDir) Refresh() error {
	return wd.tf.Refresh(context.Background(), tfexec.State(filepath.Join(wd.baseDir, "terraform.tfstate")), tfexec.Dir(wd.configDir))
}

// RequireRefresh is a variant of Refresh that will fail the test via
// the given TestControl if the refresh is non successful.
func (wd *WorkingDir) RequireRefresh(t TestControl) {
	t.Helper()
	if err := wd.Refresh(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to refresh: %s", err)
	}
}

// Schemas returns an object describing the provider schemas.
//
// If the schemas cannot be read, Schemas returns an error.
func (wd *WorkingDir) Schemas() (*tfjson.ProviderSchemas, error) {
	return wd.tf.ProvidersSchema(context.Background())
}

// RequireSchemas is a variant of Schemas that will fail the test via
// the given TestControl if the schemas cannot be read.
func (wd *WorkingDir) RequireSchemas(t TestControl) *tfjson.ProviderSchemas {
	t.Helper()

	ret, err := wd.Schemas()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read schemas: %s", err)
	}
	return ret
}
<commit_msg>propagate reattach info<commit_after>package tftest

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/hashicorp/terraform-exec/tfexec"
	tfjson "github.com/hashicorp/terraform-json"
)

// WorkingDir represents a distinct working directory that can be used for
// running tests. Each test should construct its own WorkingDir by calling
// NewWorkingDir or RequireNewWorkingDir on its package's singleton
// tftest.Helper.
type WorkingDir struct {
	h *Helper

	// baseDir is the root of the working directory tree
	baseDir string

	// baseArgs is arguments that should be appended to all commands
	baseArgs []string

	// configDir contains the singular config file generated for each test
	configDir string

	// tf is the instance of tfexec.Terraform used for running Terraform commands
	tf *tfexec.Terraform

	// terraformExec is a path to a terraform binary, inherited from Helper
	terraformExec string

	// reattachInfo stores the gRPC socket info required for Terraform's
	// plugin reattach functionality
	reattachInfo string

	env map[string]string
}

// Close deletes the directories and files created to represent the receiving
// working directory. After this method is called, the working directory object
// is invalid and may no longer be used.
func (wd *WorkingDir) Close() error {
	return os.RemoveAll(wd.baseDir)
}

// Setenv sets an environment variable on the WorkingDir.
func (wd *WorkingDir) Setenv(envVar, val string) {
	if wd.env == nil {
		wd.env = map[string]string{}
	}
	wd.env[envVar] = val
}

// Unsetenv removes an environment variable from the WorkingDir.
func (wd *WorkingDir) Unsetenv(envVar string) {
	delete(wd.env, envVar)
}

func (wd *WorkingDir) SetReattachInfo(reattachInfo string) {
	wd.reattachInfo = reattachInfo
}

func (wd *WorkingDir) UnsetReattachInfo() {
	wd.reattachInfo = ""
}

// GetHelper returns the Helper set on the WorkingDir.
func (wd *WorkingDir) GetHelper() *Helper {
	return wd.h
}

func (wd *WorkingDir) relativeConfigDir() (string, error) {
	relPath, err := filepath.Rel(wd.baseDir, wd.configDir)
	if err != nil {
		return "", fmt.Errorf("Error determining relative path of configuration directory: %w", err)
	}
	return relPath, nil
}

// SetConfig sets a new configuration for the working directory.
//
// This must be called at least once before any call to Init, Plan, Apply, or
// Destroy to establish the configuration. Any previously-set configuration is
// discarded and any saved plan is cleared.
func (wd *WorkingDir) SetConfig(cfg string) error {
	// Each call to SetConfig creates a new directory under our baseDir.
	// We create them within so that our final cleanup step will delete them
	// automatically without any additional tracking.
	configDir, err := ioutil.TempDir(wd.baseDir, "config")
	if err != nil {
		return err
	}
	configFilename := filepath.Join(configDir, "terraform_plugin_test.tf")
	err = ioutil.WriteFile(configFilename, []byte(cfg), 0700)
	if err != nil {
		return err
	}

	tf, err := tfexec.NewTerraform(wd.baseDir, wd.terraformExec)
	if err != nil {
		return err
	}

	var mismatch *tfexec.ErrVersionMismatch
	err = tf.SetDisablePluginTLS(true)
	if err != nil && !errors.As(err, &mismatch) {
		return err
	}
	err = tf.SetSkipProviderVerify(true)
	if err != nil && !errors.As(err, &mismatch) {
		return err
	}

	if p := os.Getenv("TF_ACC_LOG_PATH"); p != "" {
		tf.SetLogPath(p)
	}

	wd.configDir = configDir
	wd.tf = tf

	// Changing configuration invalidates any saved plan.
	err = wd.ClearPlan()
	if err != nil {
		return err
	}
	return nil
}

// RequireSetConfig is a variant of SetConfig that will fail the test via the
// given TestControl if the configuration cannot be set.
func (wd *WorkingDir) RequireSetConfig(t TestControl, cfg string) {
	t.Helper()
	if err := wd.SetConfig(cfg); err != nil {
		t := testingT{t}
		t.Fatalf("failed to set config: %s", err)
	}
}

// ClearState deletes any Terraform state present in the working directory.
//
// Any remote objects tracked by the state are not destroyed first, so this
// will leave them dangling in the remote system.
func (wd *WorkingDir) ClearState() error {
	err := os.Remove(filepath.Join(wd.baseDir, "terraform.tfstate"))
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// RequireClearState is a variant of ClearState that will fail the test via the
// given TestControl if the state cannot be cleared.
func (wd *WorkingDir) RequireClearState(t TestControl) {
	t.Helper()
	if err := wd.ClearState(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to clear state: %s", err)
	}
}

// ClearPlan deletes any saved plan present in the working directory.
func (wd *WorkingDir) ClearPlan() error {
	err := os.Remove(wd.planFilename())
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// RequireClearPlan is a variant of ClearPlan that will fail the test via the
// given TestControl if the plan cannot be cleared.
func (wd *WorkingDir) RequireClearPlan(t TestControl) {
	t.Helper()
	if err := wd.ClearPlan(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to clear plan: %s", err)
	}
}

// Init runs "terraform init" for the given working directory, forcing Terraform
// to use the current version of the plugin under test.
func (wd *WorkingDir) Init() error {
	if wd.configDir == "" {
		return fmt.Errorf("must call SetConfig before Init")
	}

	return wd.tf.Init(context.Background(), tfexec.Reattach(wd.reattachInfo), tfexec.Dir(wd.configDir))
}

// RequireInit is a variant of Init that will fail the test via the given
// TestControl if init fails.
func (wd *WorkingDir) RequireInit(t TestControl) {
	t.Helper()
	if err := wd.Init(); err != nil {
		t := testingT{t}
		t.Fatalf("init failed: %s", err)
	}
}

func (wd *WorkingDir) planFilename() string {
	return filepath.Join(wd.baseDir, "tfplan")
}

// CreatePlan runs "terraform plan" to create a saved plan file, which if successful
// will then be used for the next call to Apply.
func (wd *WorkingDir) CreatePlan() error {
	_, err := wd.tf.Plan(context.Background(), tfexec.Reattach(wd.reattachInfo), tfexec.Refresh(false), tfexec.Out("tfplan"), tfexec.Dir(wd.configDir))
	return err
}

// RequireCreatePlan is a variant of CreatePlan that will fail the test via
// the given TestControl if plan creation fails.
func (wd *WorkingDir) RequireCreatePlan(t TestControl) {
	t.Helper()
	if err := wd.CreatePlan(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to create plan: %s", err)
	}
}

// CreateDestroyPlan runs "terraform plan -destroy" to create a saved plan
// file, which if successful will then be used for the next call to Apply.
func (wd *WorkingDir) CreateDestroyPlan() error {
	_, err := wd.tf.Plan(context.Background(), tfexec.Reattach(wd.reattachInfo), tfexec.Refresh(false), tfexec.Out("tfplan"), tfexec.Destroy(true), tfexec.Dir(wd.configDir))
	return err
}

// Apply runs "terraform apply". If CreatePlan has previously completed
// successfully and the saved plan has not been cleared in the meantime then
// this will apply the saved plan. Otherwise, it will implicitly create a new
// plan and apply it.
func (wd *WorkingDir) Apply() error {
	args := []tfexec.ApplyOption{tfexec.Reattach(wd.reattachInfo), tfexec.Refresh(false)}
	if wd.HasSavedPlan() {
		args = append(args, tfexec.DirOrPlan("tfplan"))
	} else {
		// we need to use a relative config dir here or we get an
		// error about Terraform not having any configuration. See
		// https://github.com/hashicorp/terraform-plugin-sdk/issues/495
		// for more info.
		configDir, err := wd.relativeConfigDir()
		if err != nil {
			return err
		}
		args = append(args, tfexec.DirOrPlan(configDir))
	}
	return wd.tf.Apply(context.Background(), args...)
}

// RequireApply is a variant of Apply that will fail the test via
// the given TestControl if the apply operation fails.
func (wd *WorkingDir) RequireApply(t TestControl) {
	t.Helper()
	if err := wd.Apply(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to apply: %s", err)
	}
}

// Destroy runs "terraform destroy". It does not consider or modify any saved
// plan, and is primarily for cleaning up at the end of a test run.
//
// If destroy fails then remote objects might still exist, and continue to
// exist after a particular test is concluded.
func (wd *WorkingDir) Destroy() error {
	return wd.tf.Destroy(context.Background(), tfexec.Reattach(wd.reattachInfo), tfexec.Refresh(false), tfexec.Dir(wd.configDir))
}

// RequireDestroy is a variant of Destroy that will fail the test via
// the given TestControl if the destroy operation fails.
//
// If destroy fails then remote objects might still exist, and continue to
// exist after a particular test is concluded.
func (wd *WorkingDir) RequireDestroy(t TestControl) {
	t.Helper()
	if err := wd.Destroy(); err != nil {
		t := testingT{t}
		t.Logf("WARNING: destroy failed, so remote objects may still exist and be subject to billing")
		t.Fatalf("failed to destroy: %s", err)
	}
}

// HasSavedPlan returns true if there is a saved plan in the working directory. If
// so, a subsequent call to Apply will apply that saved plan.
func (wd *WorkingDir) HasSavedPlan() bool {
	_, err := os.Stat(wd.planFilename())
	return err == nil
}

// SavedPlan returns an object describing the current saved plan file, if any.
//
// If no plan is saved or if the plan file cannot be read, SavedPlan returns
// an error.
func (wd *WorkingDir) SavedPlan() (*tfjson.Plan, error) {
	if !wd.HasSavedPlan() {
		return nil, fmt.Errorf("there is no current saved plan")
	}

	return wd.tf.ShowPlanFile(context.Background(), wd.planFilename(), tfexec.Reattach(wd.reattachInfo))
}

// RequireSavedPlan is a variant of SavedPlan that will fail the test via
// the given TestControl if the plan cannot be read.
func (wd *WorkingDir) RequireSavedPlan(t TestControl) *tfjson.Plan {
	t.Helper()
	ret, err := wd.SavedPlan()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read saved plan: %s", err)
	}
	return ret
}

// SavedPlanStdout returns a stdout capture of the current saved plan file, if any.
//
// If no plan is saved or if the plan file cannot be read, SavedPlanStdout returns
// an error.
func (wd *WorkingDir) SavedPlanStdout() (string, error) {
	if !wd.HasSavedPlan() {
		return "", fmt.Errorf("there is no current saved plan")
	}

	var ret bytes.Buffer

	wd.tf.SetStdout(&ret)
	defer wd.tf.SetStdout(ioutil.Discard)
	_, err := wd.tf.ShowPlanFile(context.Background(), wd.planFilename(), tfexec.Reattach(wd.reattachInfo))
	if err != nil {
		return "", err
	}

	return ret.String(), nil
}

// RequireSavedPlanStdout is a variant of SavedPlanStdout that will fail the test via
// the given TestControl if the plan cannot be read.
func (wd *WorkingDir) RequireSavedPlanStdout(t TestControl) string {
	t.Helper()
	ret, err := wd.SavedPlanStdout()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read saved plan: %s", err)
	}
	return ret
}

// State returns an object describing the current state.
//
// If the state cannot be read, State returns an error.
func (wd *WorkingDir) State() (*tfjson.State, error) {
	return wd.tf.Show(context.Background(), tfexec.Reattach(wd.reattachInfo))
}

// RequireState is a variant of State that will fail the test via
// the given TestControl if the state cannot be read.
func (wd *WorkingDir) RequireState(t TestControl) *tfjson.State {
	t.Helper()
	ret, err := wd.State()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read state plan: %s", err)
	}
	return ret
}

// Import runs terraform import
func (wd *WorkingDir) Import(resource, id string) error {
	return wd.tf.Import(context.Background(), resource, id, tfexec.Config(wd.configDir), tfexec.Reattach(wd.reattachInfo))
}

// RequireImport is a variant of Import that will fail the test via
// the given TestControl if the import is non successful.
func (wd *WorkingDir) RequireImport(t TestControl, resource, id string) {
	t.Helper()
	if err := wd.Import(resource, id); err != nil {
		t := testingT{t}
		t.Fatalf("failed to import: %s", err)
	}
}

// Refresh runs terraform refresh
func (wd *WorkingDir) Refresh() error {
	return wd.tf.Refresh(context.Background(), tfexec.Reattach(wd.reattachInfo), tfexec.State(filepath.Join(wd.baseDir, "terraform.tfstate")), tfexec.Dir(wd.configDir))
}

// RequireRefresh is a variant of Refresh that will fail the test via
// the given TestControl if the refresh is non successful.
func (wd *WorkingDir) RequireRefresh(t TestControl) {
	t.Helper()
	if err := wd.Refresh(); err != nil {
		t := testingT{t}
		t.Fatalf("failed to refresh: %s", err)
	}
}

// Schemas returns an object describing the provider schemas.
//
// If the schemas cannot be read, Schemas returns an error.
func (wd *WorkingDir) Schemas() (*tfjson.ProviderSchemas, error) {
	return wd.tf.ProvidersSchema(context.Background())
}

// RequireSchemas is a variant of Schemas that will fail the test via
// the given TestControl if the schemas cannot be read.
func (wd *WorkingDir) RequireSchemas(t TestControl) *tfjson.ProviderSchemas {
	t.Helper()

	ret, err := wd.Schemas()
	if err != nil {
		t := testingT{t}
		t.Fatalf("failed to read schemas: %s", err)
	}
	return ret
}
<|endoftext|>


<commit_before><commit_msg>Added new urls<commit_after><|endoftext|>


<commit_before>package response

import (
	"sort"
	"strings"

	version "github.com/hashicorp/go-version"
)

// TerraformProvider is the response structure for all required information for
// Terraform to choose a download URL. It must include all versions and all
// platforms for Terraform to perform version and os/arch constraint matching
// locally.
type TerraformProvider struct {
	ID       string `json:"id"`
	Verified bool   `json:"verified"`

	Versions []*TerraformProviderVersion `json:"versions"`
}

// TerraformProviderVersion is the Terraform-specific response structure for a
// provider version.
type TerraformProviderVersion struct {
	Version   string   `json:"version"`
	Protocols []string `json:"protocols"`

	Platforms []*TerraformProviderPlatform `json:"platforms"`
}

// TerraformProviderVersions is the Terraform-specific response structure for an
// array of provider versions
type TerraformProviderVersions struct {
	ID       string                      `json:"id"`
	Versions []*TerraformProviderVersion `json:"versions"`
	Warnings []string                    `json:"warnings"`
}

// TerraformProviderPlatform is the Terraform-specific response structure for a
// provider platform.
type TerraformProviderPlatform struct {
	OS   string `json:"os"`
	Arch string `json:"arch"`
}

// TerraformProviderPlatformLocation is the Terraform-specific response
// structure for a provider platform with all details required to perform a
// download.
type TerraformProviderPlatformLocation struct {
	Protocols           []string `json:"protocols"`
	OS                  string   `json:"os"`
	Arch                string   `json:"arch"`
	Filename            string   `json:"filename"`
	DownloadURL         string   `json:"download_url"`
	ShasumsURL          string   `json:"shasums_url"`
	ShasumsSignatureURL string   `json:"shasums_signature_url"`
	Shasum              string   `json:"shasum"`

	SigningKeys SigningKeyList `json:"signing_keys"`
}

// SigningKeyList is the response structure for a list of signing keys.
type SigningKeyList struct {
	GPGKeys []*GPGKey `json:"gpg_public_keys"`
}

// GPGKey is the response structure for a GPG key.
type GPGKey struct {
	ASCIIArmor string  `json:"ascii_armor"`
	Source     string  `json:"source"`
	SourceURL  *string `json:"source_url"`
}

// Collection type for TerraformProviderVersion
type ProviderVersionCollection []*TerraformProviderVersion

// GPGASCIIArmor returns an ASCII-armor-formatted string for all of the gpg
// keys in the response.
func (signingKeys *SigningKeyList) GPGASCIIArmor() string {
	keys := []string{}

	for _, gpgKey := range signingKeys.GPGKeys {
		keys = append(keys, gpgKey.ASCIIArmor)
	}

	return strings.Join(keys, "\n")
}

// Sort sorts versions from newest to oldest.
func (v ProviderVersionCollection) Sort() {
	sort.Slice(v, func(i, j int) bool {
		versionA, _ := version.NewVersion(v[i].Version)
		versionB, _ := version.NewVersion(v[j].Version)

		return versionA.GreaterThan(versionB)
	})
}
<commit_msg>registry/response: Remove unused verified bool<commit_after>package response

import (
	"sort"
	"strings"

	version "github.com/hashicorp/go-version"
)

// TerraformProvider is the response structure for all required information for
// Terraform to choose a download URL. It must include all versions and all
// platforms for Terraform to perform version and os/arch constraint matching
// locally.
type TerraformProvider struct {
	ID string `json:"id"`

	Versions []*TerraformProviderVersion `json:"versions"`
}

// TerraformProviderVersion is the Terraform-specific response structure for a
// provider version.
type TerraformProviderVersion struct {
	Version   string   `json:"version"`
	Protocols []string `json:"protocols"`

	Platforms []*TerraformProviderPlatform `json:"platforms"`
}

// TerraformProviderVersions is the Terraform-specific response structure for an
// array of provider versions
type TerraformProviderVersions struct {
	ID       string                      `json:"id"`
	Versions []*TerraformProviderVersion `json:"versions"`
	Warnings []string                    `json:"warnings"`
}

// TerraformProviderPlatform is the Terraform-specific response structure for a
// provider platform.
type TerraformProviderPlatform struct {
	OS   string `json:"os"`
	Arch string `json:"arch"`
}

// TerraformProviderPlatformLocation is the Terraform-specific response
// structure for a provider platform with all details required to perform a
// download.
type TerraformProviderPlatformLocation struct {
	Protocols           []string `json:"protocols"`
	OS                  string   `json:"os"`
	Arch                string   `json:"arch"`
	Filename            string   `json:"filename"`
	DownloadURL         string   `json:"download_url"`
	ShasumsURL          string   `json:"shasums_url"`
	ShasumsSignatureURL string   `json:"shasums_signature_url"`
	Shasum              string   `json:"shasum"`

	SigningKeys SigningKeyList `json:"signing_keys"`
}

// SigningKeyList is the response structure for a list of signing keys.
type SigningKeyList struct {
	GPGKeys []*GPGKey `json:"gpg_public_keys"`
}

// GPGKey is the response structure for a GPG key.
type GPGKey struct {
	ASCIIArmor string  `json:"ascii_armor"`
	Source     string  `json:"source"`
	SourceURL  *string `json:"source_url"`
}

// Collection type for TerraformProviderVersion
type ProviderVersionCollection []*TerraformProviderVersion

// GPGASCIIArmor returns an ASCII-armor-formatted string for all of the gpg
// keys in the response.
func (signingKeys *SigningKeyList) GPGASCIIArmor() string {
	keys := []string{}

	for _, gpgKey := range signingKeys.GPGKeys {
		keys = append(keys, gpgKey.ASCIIArmor)
	}

	return strings.Join(keys, "\n")
}

// Sort sorts versions from newest to oldest.
func (v ProviderVersionCollection) Sort() {
	sort.Slice(v, func(i, j int) bool {
		versionA, _ := version.NewVersion(v[i].Version)
		versionB, _ := version.NewVersion(v[j].Version)

		return versionA.GreaterThan(versionB)
	})
}
<|endoftext|>


<commit_before>package hermes

import (
	"fmt"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/PuerkitoBio/fetchbot"
	"github.com/PuerkitoBio/goquery"
)

var (
	mu sync.Mutex // Protect access to dup

	dup = map[string]bool{} // Duplicates table

	settingLinks = map[string]bool{} // Tracking link settings

	ingestionSet []Document // ingestion data TODO make non global

	badLinks []string // bad links TODO make non global
)

// Crawl function that will take a url string and start firing out some crawling functions
// it will return true/false based on the url root it starts with.
func Crawl(settings Settings, linkSettings CustomSettings, u url.URL) ([]Document, bool) {
	// Create the muxer
	mux := fetchbot.NewMux()

	// Handle all errors the same
	mux.HandleErrors(fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
	}))

	// Handle GET requests for html responses, to parse the body and enqueue all links as HEAD
	// requests.
	mux.Response().Method("GET").ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			// Process the body to find the links
			doc, err := goquery.NewDocumentFromResponse(res)
			if err != nil {
				// find the bad links in the documents
				badLinks = append(badLinks, ctx.Cmd.URL().String())
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
				return
			}
			// Enqueue all links as HEAD requests
			enqueueLinks(ctx, doc)
		}))

	// Handle HEAD requests for html responses coming from the source host - we don't want
	// to crawl links from other hosts.
	mux.Response().Method("HEAD").Host(u.Host).ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			if _, err := ctx.Q.SendStringGet(ctx.Cmd.URL().String()); err != nil {
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
			}
		}))

	// Create the Fetcher, handle the logging first, then dispatch to the Muxer
	h := scrapeHandler(mux, linkSettings)
	if settings.StopAtURL != "" || settings.CancelAtURL != "" {
		stopURL := settings.StopAtURL
		if settings.CancelAtURL != "" {
			stopURL = settings.CancelAtURL
		}
		h = stopHandler(stopURL, settings.CancelAtURL != "", scrapeHandler(mux, linkSettings))
	}
	f := fetchbot.New(h)

	// set the fetchbots settings from flag parameters
	f.UserAgent = settings.UserAgent
	f.CrawlDelay = settings.CrawlDelay * time.Second
	f.WorkerIdleTTL = settings.WorkerIdleTTL * time.Second

	// First mem stat print must be right after creating the fetchbot
	if settings.MemStatsInterval > 0 {
		// Print starting stats
		printMemStats(nil)
		// Run at regular intervals
		runMemStats(f, settings.MemStatsInterval)
		// On exit, print ending stats after a GC
		defer func() {
			runtime.GC()
			printMemStats(nil)
		}()
	}

	// Start processing
	q := f.Start()

	// if a stop or cancel is requested after some duration, launch the goroutine
	// that will stop or cancel.
	if settings.StopDuration*time.Minute > 0 || settings.CancelDuration*time.Minute > 0 {
		after := settings.StopDuration * time.Minute
		stopFunc := q.Close
		if settings.CancelDuration != 0 {
			after = settings.CancelDuration * time.Minute
			stopFunc = q.Cancel
		}

		go func() {
			c := time.After(after)
			<-c
			stopFunc()
		}()
	}

	// Enqueue the seed, which is the first entry in the dup map
	dup[linkSettings.RootLink] = true
	_, err := q.SendStringGet(linkSettings.RootLink)
	if err != nil {
		fmt.Printf("[ERR] GET %s - %s\n", linkSettings.RootLink, err)
	}
	q.Block()

	return ingestionSet, true
}

// stopHandler stops the fetcher if the stopurl is reached. Otherwise it dispatches
// the call to the wrapped Handler.
func stopHandler(stopurl string, cancel bool, wrapped fetchbot.Handler) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if ctx.Cmd.URL().String() == stopurl {
			fmt.Printf(">>>>> STOP URL %s\n", ctx.Cmd.URL())
			// generally not a good idea to stop/block from a handler goroutine
			// so do it in a separate goroutine
			go func() {
				if cancel {
					ctx.Q.Cancel()
				} else {
					ctx.Q.Close()
				}
			}()
			return
		}
		wrapped.Handle(ctx, res, err)
	})
}

// scrapeHandler will fire a scraper function on the page if successful response,
// append the scraped document stored for index ingestion
// and dispatches the call to the wrapped Handler.
func scrapeHandler(wrapped fetchbot.Handler, linkSettings CustomSettings) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if err == nil {
			fmt.Printf("[%d] %s %s - %s\n", res.StatusCode, ctx.Cmd.Method(), ctx.Cmd.URL(), res.Header.Get("Content-Type"))
		} else {
			if res.StatusCode == 200 {
				doc, err := goquery.NewDocumentFromResponse(res)
				if err != nil {
					// find the bad links in the documents
					badLinks = append(badLinks, ctx.Cmd.URL().String())
					fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
					return
				}

				// fire scraper
				content := scraper(ctx, doc, linkSettings)
				ingestionSet = append(ingestionSet, content)
			} else {
				fmt.Println("scrapeHandler bad links +1 ========== status code: ", res.StatusCode)
				badLinks = append(badLinks, ctx.Cmd.URL().String())
			}
		}
		wrapped.Handle(ctx, res, err)
	})
}

// enqueueLinks will make sure we are adding links to the queue to be processed for crawling/scraping
// this will pull all the href attributes on pages, check for duplicates and add them to the queue
func enqueueLinks(ctx *fetchbot.Context, doc *goquery.Document) {
	mu.Lock()
	doc.Find("a[href]").Each(func(i int, s *goquery.Selection) {
		val, _ := s.Attr("href")
		// Resolve address
		u, err := ctx.Cmd.URL().Parse(val)
		if err != nil {
			fmt.Printf("error: resolve URL %s - %s\n", val, err)
			return
		}
		// catch the duplicate urls here before trying to add them to the queue
		if !dup[u.String()] {
			if _, err := ctx.Q.SendStringHead(u.String()); err != nil {
				fmt.Printf("error: enqueue head %s - %s\n", u, err)
			} else {
				dup[u.String()] = true
			}
		}
	})
	mu.Unlock()
}

func scraper(ctx *fetchbot.Context, doc *goquery.Document, linkSettings CustomSettings) Document {
	mu.Lock()
	u := ctx.Cmd.URL().String()

	document := scrapeDocument(u, doc, linkSettings.Tags)

	mu.Unlock()
	return document
}

func scrapeDocument(url string, doc *goquery.Document, tags []string) Document {
	var (
		d       Document
		content string
	)
	// scrape page <Title>
	doc.Find("head").Each(func(i int, s *goquery.Selection) {
		d.Title = s.Find("title").Text()
	})

	// scrape page <Description>
	doc.Find("meta").Each(func(i int, s *goquery.Selection) {
		if name, _ := s.Attr("name"); strings.EqualFold(name, "description") {
			description, _ := s.Attr("content")
			d.Description = description
		}
	})

	if len(tags) > 0 {
		for _, tag := range tags {
			text := returnText(doc, tag)
			content += " " + text
			fmt.Printf("Tag %s, returned %s\n", tag, text)
		}
	} else {
		text := returnText(doc, "default")
		content += " " + text
		fmt.Printf("Default tag, returned %s\n", text)
	}

	d.Content = content
	d.Link = url

	fmt.Println(d)

	return d
}

func returnText(doc *goquery.Document, tag string) (text string) {
	doc.Find("body").Each(func(i int, s *goquery.Selection) {
		// default to pulling all the div and p tags else pull custom setting tags
		if tag == "default" {
			text += " " + s.Find("p").Text()
			text += " " + s.Find("div").Text()
		} else {
			text += " " + s.Find(tag).Text()
		}
	})
	return
}
<commit_msg>removed status check<commit_after>package hermes

import (
	"fmt"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/PuerkitoBio/fetchbot"
	"github.com/PuerkitoBio/goquery"
)

var (
	mu sync.Mutex // Protect access to dup

	dup = map[string]bool{} // Duplicates table

	settingLinks = map[string]bool{} // Tracking link settings

	ingestionSet []Document // ingestion data TODO make non global

	badLinks []string // bad links TODO make non global
)

// Crawl function that will take a url string and start firing out some crawling functions
// it will return true/false based on the url root it starts with.
func Crawl(settings Settings, linkSettings CustomSettings, u url.URL) ([]Document, bool) {
	// Create the muxer
	mux := fetchbot.NewMux()

	// Handle all errors the same
	mux.HandleErrors(fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
	}))

	// Handle GET requests for html responses, to parse the body and enqueue all links as HEAD
	// requests.
	mux.Response().Method("GET").ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			// Process the body to find the links
			doc, err := goquery.NewDocumentFromResponse(res)
			if err != nil {
				// find the bad links in the documents
				badLinks = append(badLinks, ctx.Cmd.URL().String())
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
				return
			}
			// Enqueue all links as HEAD requests
			enqueueLinks(ctx, doc)
		}))

	// Handle HEAD requests for html responses coming from the source host - we don't want
	// to crawl links from other hosts.
	mux.Response().Method("HEAD").Host(u.Host).ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			if _, err := ctx.Q.SendStringGet(ctx.Cmd.URL().String()); err != nil {
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
			}
		}))

	// Create the Fetcher, handle the logging first, then dispatch to the Muxer
	h := scrapeHandler(mux, linkSettings)
	if settings.StopAtURL != "" || settings.CancelAtURL != "" {
		stopURL := settings.StopAtURL
		if settings.CancelAtURL != "" {
			stopURL = settings.CancelAtURL
		}
		h = stopHandler(stopURL, settings.CancelAtURL != "", scrapeHandler(mux, linkSettings))
	}
	f := fetchbot.New(h)

	// set the fetchbots settings from flag parameters
	f.UserAgent = settings.UserAgent
	f.CrawlDelay = settings.CrawlDelay * time.Second
	f.WorkerIdleTTL = settings.WorkerIdleTTL * time.Second

	// First mem stat print must be right after creating the fetchbot
	if settings.MemStatsInterval > 0 {
		// Print starting stats
		printMemStats(nil)
		// Run at regular intervals
		runMemStats(f, settings.MemStatsInterval)
		// On exit, print ending stats after a GC
		defer func() {
			runtime.GC()
			printMemStats(nil)
		}()
	}

	// Start processing
	q := f.Start()

	// if a stop or cancel is requested after some duration, launch the goroutine
	// that will stop or cancel.
	if settings.StopDuration*time.Minute > 0 || settings.CancelDuration*time.Minute > 0 {
		after := settings.StopDuration * time.Minute
		stopFunc := q.Close
		if settings.CancelDuration != 0 {
			after = settings.CancelDuration * time.Minute
			stopFunc = q.Cancel
		}

		go func() {
			c := time.After(after)
			<-c
			stopFunc()
		}()
	}

	// Enqueue the seed, which is the first entry in the dup map
	dup[linkSettings.RootLink] = true
	_, err := q.SendStringGet(linkSettings.RootLink)
	if err != nil {
		fmt.Printf("[ERR] GET %s - %s\n", linkSettings.RootLink, err)
	}
	q.Block()

	return ingestionSet, true
}

// stopHandler stops the fetcher if the stopurl is reached. Otherwise it dispatches
// the call to the wrapped Handler.
func stopHandler(stopurl string, cancel bool, wrapped fetchbot.Handler) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if ctx.Cmd.URL().String() == stopurl {
			fmt.Printf(">>>>> STOP URL %s\n", ctx.Cmd.URL())
			// generally not a good idea to stop/block from a handler goroutine
			// so do it in a separate goroutine
			go func() {
				if cancel {
					ctx.Q.Cancel()
				} else {
					ctx.Q.Close()
				}
			}()
			return
		}
		wrapped.Handle(ctx, res, err)
	})
}

// scrapeHandler will fire a scraper function on the page if successful response,
// append the scraped document stored for index ingestion
// and dispatches the call to the wrapped Handler.
func scrapeHandler(wrapped fetchbot.Handler, linkSettings CustomSettings) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if err == nil {
			fmt.Printf("[%d] %s %s - %s\n", res.StatusCode, ctx.Cmd.Method(), ctx.Cmd.URL(), res.Header.Get("Content-Type"))
		} else {
			doc, err := goquery.NewDocumentFromResponse(res)
			if err != nil {
				// find the bad links in the documents
				badLinks = append(badLinks, ctx.Cmd.URL().String())
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
				return
			}

			// fire scraper
			content := scraper(ctx, doc, linkSettings)
			ingestionSet = append(ingestionSet, content)
		}
		wrapped.Handle(ctx, res, err)
	})
}

// enqueueLinks will make sure we are adding links to the queue to be processed for crawling/scraping
// this will pull all the href attributes on pages, check for duplicates and add them to the queue
func enqueueLinks(ctx *fetchbot.Context, doc *goquery.Document) {
	mu.Lock()
	doc.Find("a[href]").Each(func(i int, s *goquery.Selection) {
		val, _ := s.Attr("href")
		// Resolve address
		u, err := ctx.Cmd.URL().Parse(val)
		if err != nil {
			fmt.Printf("error: resolve URL %s - %s\n", val, err)
			return
		}
		// catch the duplicate urls here before trying to add them to the queue
		if !dup[u.String()] {
			if _, err := ctx.Q.SendStringHead(u.String()); err != nil {
				fmt.Printf("error: enqueue head %s - %s\n", u, err)
			} else {
				dup[u.String()] = true
			}
		}
	})
	mu.Unlock()
}

func scraper(ctx *fetchbot.Context, doc *goquery.Document, linkSettings CustomSettings) Document {
	mu.Lock()
	u := ctx.Cmd.URL().String()

	document := scrapeDocument(u, doc, linkSettings.Tags)

	mu.Unlock()
	return document
}

func scrapeDocument(url string, doc *goquery.Document, tags []string) Document {
	var (
		d       Document
		content string
	)
	// scrape page <Title>
	doc.Find("head").Each(func(i int, s *goquery.Selection) {
		d.Title = s.Find("title").Text()
	})

	// scrape page <Description>
	doc.Find("meta").Each(func(i int, s *goquery.Selection) {
		if name, _ := s.Attr("name"); strings.EqualFold(name, "description") {
			description, _ := s.Attr("content")
			d.Description = description
		}
	})

	if len(tags) > 0 {
		for _, tag := range tags {
			text := returnText(doc, tag)
			content += " " + text
			fmt.Printf("Tag %s, returned %s\n", tag, text)
		}
	} else {
		text := returnText(doc, "default")
		content += " " + text
		fmt.Printf("Default tag, returned %s\n", text)
	}

	d.Content = content
	d.Link = url

	fmt.Println(d)

	return d
}

func returnText(doc *goquery.Document, tag string) (text string) {
	doc.Find("body").Each(func(i int, s *goquery.Selection) {
		// default to pulling all the div and p tags else pull custom setting tags
		if tag == "default" {
			text += " " + s.Find("p").Text()
			text += " " + s.Find("div").Text()
		} else {
			text += " " + s.Find(tag).Text()
		}
	})
	return
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"time"
)

var version = "0.3.0"
var favicon []byte
var index []byte
var robots []byte

var runningProcs = 0
var totalTimeoutRequests = 0
var totalTCPRequests = 0
var totalIndexRequests = 0

var http10StatusOK = []byte("HTTP/1.0 200 OK\r\n")
var http11StatusOK = []byte("HTTP/1.1 200 OK\r\n")

//plain-text
var invalidFormat = []byte("invalid_format")
var done = []byte("done")
var invalidCommand = []byte("Unknown command. Type 'info'\n")

//int
var one = []byte("1")

//json
var jsonTrue = []byte("true")

//js
var jsTrue = []byte("_settimeoutio=true")

//css
var cssShow = []byte(".settimeoutio {display: block;}")

//callback
var emptyCallback = []string{""}

var serverName = "settimeout/" + version

func main() {
	addr := flag.String("addr", ":80", "HTTP address to listen on (empty to disable)")
	tcpAddr := flag.String("tcpaddr", ":5103", "Socket address to listen on (empty to disable)")
	statsAddr := flag.String("statsaddr", "127.0.0.1:5104", "Socket address to listen on for stats (empty to disable)")
	flag.Parse()

	var err error
	favicon, err = Asset("assets/favicon.ico")
	if err != nil {
		log.Fatal("Failed to read favicon.ico: " + err.Error())
	}
	index, err = Asset("assets/index.html")
	if err != nil {
		log.Fatal("Failed to read index.html: " + err.Error())
	}
	robots, err = Asset("assets/robots.txt")
	if err != nil {
		log.Fatal("Failed to read robots.txt: " + err.Error())
	}

	go startHTTPServer(*addr)
	go startTCPServer(*tcpAddr)
	go startStatsTCPServer(*statsAddr)

	//wait for a SIGINT
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c
}

func startTCPServer(addr string) {
	if addr == "" {
		return
	}
	listen, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("Failed to start TCP server: " + err.Error())
	}
	for {
		conn, err := listen.Accept()
		if err != nil {
			// handle error
			continue
		}
		go socketHandler(conn)
	}
}

func startStatsTCPServer(addr string) {
	if addr == "" {
		return
	}
	listen, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("Failed to start stats TCP server: " + err.Error())
	}
	for {
		conn, err := listen.Accept()
		if err != nil {
			// handle error
			continue
		}
		go func() {
			defer conn.Close()
			for {
				conn.SetReadDeadline(time.Now().Add(3 * time.Second))
				buf := bufio.NewReader(conn)
				str, err := buf.ReadString('\n')
				if err != nil {
					if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
						//if its a i/o timeout then just continue the loop again
						continue
					}
					return
				}
				str = strings.Trim(str, " \r\n\t")
				switch str {
				case "close", "quit", "exit":
					return
				case "info":
					infos := make([]string, 4)
					infos[0] = "Running Requests: " + strconv.Itoa(runningProcs)
					infos[1] = "Total HTTP Index Requests: " + strconv.Itoa(totalIndexRequests)
					infos[2] = "Total TCP Requests: " + strconv.Itoa(totalTCPRequests)
					infos[3] = "Total HTTP Timeout Requests: " + strconv.Itoa(totalTimeoutRequests)
					conn.Write([]byte(strings.Join(infos, "\n")))
				default:
					conn.Write(invalidCommand)
				}
			}
		}()
	}
}

func startHTTPServer(addr string) {
	if addr == "" {
		return
	}
	s := &http.Server{
		Addr:        addr,
		Handler:     http.HandlerFunc(httpHandler),
		ReadTimeout: time.Duration(5) * time.Second,
	}
	err := s.ListenAndServe()
	if err != nil {
		log.Fatal("Failed to start HTTP server: " + err.Error())
	}
}

func socketHandler(conn net.Conn) {
	defer conn.Close()
	conn.SetReadDeadline(time.Now().Add(3 * time.Second))
	buf := bufio.NewReader(conn)
	str, err := buf.ReadString('\n')
	if err != nil {
		return
	}
	totalTCPRequests++
	d, err := parseDurationString(str)
	if err != nil {
		conn.Write(invalidFormat)
		return
	}
	waitForTimeoutClose(conn, d, done)
}

func waitForTimeoutClose(conn net.Conn, timeout time.Duration, resp []byte) {
	if timeout.Nanoseconds() <= 0 {
		conn.Write(resp)
		return
	}

	runningProcs++
	closedCh := make(chan int)
	go func() {

		defer func() {
			runningProcs--
		}()

		singleChar := make([]byte, 1)
		for {
			conn.SetReadDeadline(time.Now().Add(5 * time.Second))
			_, err := conn.Read(singleChar)
			//if we got a read error and it wasn't a timeout then close the connection
			if netErr, ok := err.(net.Error); err != nil && (!ok || !netErr.Timeout()) {
				select {
				case closedCh <- 1:
				default:
				}
			}
			return
		}
	}()
	select {
	case <-closedCh:
		return
	case <-time.After(timeout):
		conn.Write(done)
	}
}

func httpHandler(w http.ResponseWriter, req *http.Request) {
	reqHeader := w.Header()
	reqHeader.Set("Access-Control-Allow-Origin", "*")
	reqHeader.Set("Cache-Control", "no-store, no-cache, must-revalidate, post-check=0, pre-check=0")
	reqHeader.Set("Server", serverName)
	if req.Method != "GET" && req.Method != "POST" {
		//assume its the browser just checking headers
		w.WriteHeader(http.StatusOK)
		return
	}
	str := strings.TrimLeft(req.URL.Path, "/")

	switch str {
	case "": //index page
		reqHeader.Set("Cache-Control", "max-age=3600") //1 hour
		w.Write(index)
		totalIndexRequests++
	case "favicon.ico":
		reqHeader.Set("Cache-Control", "max-age=31536000") //1 year
		w.Write(favicon)
	case "robots.txt":
		reqHeader.Set("Cache-Control", "max-age=31536000") //1 year
		w.Write(robots)
	default:
		totalTimeoutRequests++
		d, err := parseDurationString(str)
		query := req.URL.Query()
		exists := func(key string) bool { _, ok := query[key]; return ok }
		var resp []byte
		var contentType string
		switch {
		case exists("js"):
			resp = jsTrue
			contentType = "application/javascript"
		case exists("callback"):
			callback, _ := query["callback"]
			if len(callback) == 0 {
				callback = emptyCallback
			}
			resp = []byte(callback[0] + "(true)")
			contentType = "application/javascript"
		case exists("int"):
			resp = one
			contentType = "text/plain"
		case exists("json"):
			resp = jsonTrue
			contentType = "application/json"
		case exists("css"):
			resp = cssShow
			contentType = "text/css"
		default:
			resp = done
			contentType = "text/plain"
		}

		reqHeader.Set("Content-Type", contentType+"; charset=utf-8")
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		hj, ok := w.(http.Hijacker)
		if !ok {
			log.Println("Connection doesn't support hijacking for some reason")
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		//now hijack the connection so we can see if they disconnected
		conn, bufrw, err := hj.Hijack()
		if err != nil {
			log.Println("Failed to hijack http connection: " + err.Error())
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		defer conn.Close()

		reqHeader.Set("Content-Length", strconv.Itoa(len(resp)))
		reqHeader.Set("Connection", "close")

		//send back headers
		statusLine := http10StatusOK
		if req.ProtoAtLeast(1, 1) {
			statusLine = http11StatusOK
		}
		bufrw.Write(statusLine)
		reqHeader.Write(bufrw)

		//if they're going to be waiting more than 3 seconds immediately
		//send back headers to prevent timeout
		if d.Seconds() > 3 {
			bufrw.Flush()
		}
		waitForTimeoutClose(conn, d, resp)
	}
}

func parseDurationString(timeStr string) (time.Duration, error) {
	timeStr = strings.Trim(timeStr, " \r\n\t")
	var d time.Duration
	ms, err := strconv.ParseUint(timeStr, 10, 64)
	if err != nil {
		d, err = time.ParseDuration(timeStr)
		if err != nil {
			return d, err
		}
	} else {
		d = time.Duration(ms) * time.Millisecond
	}
	return d, nil
}
<commit_msg>Woops, forgot CRLF and added a closeAndWait like http has<commit_after>package main

import (
	"bufio"
	"flag"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"time"
)

var version = "0.3.1"
var favicon []byte
var index []byte
var robots []byte
var readDeadline = 5 * time.Second

var runningProcs = 0
var totalTimeoutRequests = 0
var totalTCPRequests = 0
var totalIndexRequests = 0

var http10StatusOK = []byte("HTTP/1.0 200 OK\r\n")
var http11StatusOK = []byte("HTTP/1.1 200 OK\r\n")
var crlf = []byte("\r\n")

//plain-text
var invalidFormat = []byte("invalid_format")
var done = []byte("done")
var invalidCommand = []byte("Unknown command. Type 'info'\n")

//int
var one = []byte("1")

//json
var jsonTrue = []byte("true")

//js
var jsTrue = []byte("_settimeoutio=true")

//css
var cssShow = []byte(".settimeoutio {display: block;}")

//callback
var emptyCallback = []string{""}

var serverName = "settimeout/" + version

func main() {
	addr := flag.String("addr", ":80", "HTTP address to listen on (empty to disable)")
	tcpAddr := flag.String("tcpaddr", ":5103", "Socket address to listen on (empty to disable)")
	statsAddr := flag.String("statsaddr", "127.0.0.1:5104", "Socket address to listen on for stats (empty to disable)")
	flag.Parse()

	var err error
	favicon, err = Asset("assets/favicon.ico")
	if err != nil {
		log.Fatal("Failed to read favicon.ico: " + err.Error())
	}
	index, err = Asset("assets/index.html")
	if err != nil {
		log.Fatal("Failed to read index.html: " + err.Error())
	}
	robots, err = Asset("assets/robots.txt")
	if err != nil {
		log.Fatal("Failed to read robots.txt: " + err.Error())
	}

	go startHTTPServer(*addr)
	go startTCPServer(*tcpAddr)
	go startStatsTCPServer(*statsAddr)

	//wait for a SIGINT
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c
}

func startTCPServer(addr string) {
	if addr == "" {
		return
	}
	listen, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("Failed to start TCP server: " + err.Error())
	}
	for {
		conn, err := listen.Accept()
		if err != nil {
			// handle error
			continue
		}
		go socketHandler(conn)
	}
}

func startStatsTCPServer(addr string) {
	if addr == "" {
		return
	}
	listen, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("Failed to start stats TCP server: " + err.Error())
	}
	for {
		conn, err := listen.Accept()
		if err != nil {
			// handle error
			continue
		}
		go func() {
			defer conn.Close()
			for {
				conn.SetReadDeadline(time.Now().Add(readDeadline))
				buf := bufio.NewReader(conn)
				str, err := buf.ReadString('\n')
				if err != nil {
					if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
						//if its a i/o timeout then just continue the loop again
						continue
					}
					return
				}
				str = strings.Trim(str, " \r\n\t")
				switch str {
				case "close", "quit", "exit":
					return
				case "info":
					infos := make([]string, 4)
					infos[0] = "Running Requests: " + strconv.Itoa(runningProcs)
					infos[1] = "Total HTTP Index Requests: " + strconv.Itoa(totalIndexRequests)
					infos[2] = "Total TCP Requests: " + strconv.Itoa(totalTCPRequests)
					infos[3] = "Total HTTP Timeout Requests: " + strconv.Itoa(totalTimeoutRequests)
					conn.Write([]byte(strings.Join(infos, "\n")))
				default:
					conn.Write(invalidCommand)
				}
			}
		}()
	}
}

func startHTTPServer(addr string) {
	if addr == "" {
		return
	}
	s := &http.Server{
		Addr:        addr,
		Handler:     http.HandlerFunc(httpHandler),
		ReadTimeout: readDeadline,
		//don't set a write deadline since it would set it from NOW which is not what we want
	}
	err := s.ListenAndServe()
	if err != nil {
		log.Fatal("Failed to start HTTP server: " + err.Error())
	}
}

func socketHandler(conn net.Conn) {
	defer closeAndWait(conn)
	conn.SetReadDeadline(time.Now().Add(readDeadline))
	buf := bufio.NewReader(conn)
	str, err := buf.ReadString('\n')
	if err != nil {
		return
	}
	totalTCPRequests++
	d, err := parseDurationString(str)
	if err != nil {
		conn.Write(invalidFormat)
		return
	}
	if waitForTimeoutClose(conn, d) {
		conn.Write(done)
	}
}

//returns if we should write or not becuase were leaving the writing to the caller
func waitForTimeoutClose(c net.Conn, timeout time.Duration) bool {
	if timeout.Nanoseconds() <= 0 {
		return true
	}

	runningProcs++
	closedCh := make(chan int)
	go func() {

		defer func() {
			runningProcs--
		}()

		singleChar := make([]byte, 1)
		for {
			c.SetReadDeadline(time.Now().Add(readDeadline))
			_, err := c.Read(singleChar)
			//if we got a read error and it wasn't a timeout then close the connection
			if netErr, ok := err.(net.Error); err != nil && (!ok || !netErr.Timeout()) {
				select {
				case closedCh <- 1:
				default:
				}
			}
			return
		}
	}()
	select {
	case <-closedCh:
		return false
	case <-time.After(timeout):
		return true
	}
}

func closeAndWait(c net.Conn) {
	defer c.Close()
	//from server.go which is from http://golang.org/issue/3595
	//close writing side first before actually closing the connection
	//and sending a RST
	if tcp, ok := c.(*net.TCPConn); ok {
		tcp.CloseWrite()
	}
	//using 500 milliseconds which is from rstAvoidanceDelay in server.go
	//rough estimate at the longest latency we'll have around the world
	time.Sleep(500 * time.Millisecond)
}

func httpHandler(w http.ResponseWriter, req *http.Request) {
	reqHeader := w.Header()
	reqHeader.Set("Access-Control-Allow-Origin", "*")
	reqHeader.Set("Cache-Control", "no-store, no-cache, must-revalidate, post-check=0, pre-check=0")
	reqHeader.Set("Server", serverName)
	if req.Method != "GET" && req.Method != "POST" {
		//assume its the browser just checking headers
		w.WriteHeader(http.StatusOK)
		return
	}
	str := strings.TrimLeft(req.URL.Path, "/")

	switch str {
	case "": //index page
		reqHeader.Set("Cache-Control", "max-age=3600") //1 hour
		w.Write(index)
		totalIndexRequests++
	case "favicon.ico":
		reqHeader.Set("Cache-Control", "max-age=31536000") //1 year
		w.Write(favicon)
	case "robots.txt":
		reqHeader.Set("Cache-Control", "max-age=31536000") //1 year
		w.Write(robots)
	default:
		totalTimeoutRequests++
		d, err := parseDurationString(str)
		query := req.URL.Query()
		exists := func(key string) bool { _, ok := query[key]; return ok }
		var resp []byte
		var contentType string
		switch {
		case exists("js"):
			resp = jsTrue
			contentType = "application/javascript"
		case exists("callback"):
			callback, _ := query["callback"]
			if len(callback) == 0 {
				callback = emptyCallback
			}
			resp = []byte(callback[0] + "(true)")
			contentType = "application/javascript"
		case exists("int"):
			resp = one
			contentType = "text/plain"
		case exists("json"):
			resp = jsonTrue
			contentType = "application/json"
		case exists("css"):
			resp = cssShow
			contentType = "text/css"
		default:
			resp = done
			contentType = "text/plain"
		}

		reqHeader.Set("Content-Type", contentType+"; charset=utf-8")
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		hj, ok := w.(http.Hijacker)
		if !ok {
			log.Println("Connection doesn't support hijacking for some reason")
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		//now hijack the connection so we can see if they disconnected
		conn, bufrw, err := hj.Hijack()
		if err != nil {
			log.Println("Failed to hijack http connection: " + err.Error())
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		defer closeAndWait(conn)

		reqHeader.Set("Content-Length", strconv.Itoa(len(resp)))
		reqHeader.Set("Connection", "close")

		//send back headers
		statusLine := http10StatusOK
		if req.ProtoAtLeast(1, 1) {
			statusLine = http11StatusOK
		}
		bufrw.Write(statusLine)
		reqHeader.Write(bufrw)
		bufrw.Write(crlf)

		//if they're going to be waiting more than 3 seconds immediately
		//send back headers to prevent timeout
		if d.Seconds() > 3 {
			bufrw.Flush()
		}
		if waitForTimeoutClose(conn, d) {
			bufrw.Write(resp)
			bufrw.Flush()
		}
	}
}

func parseDurationString(timeStr string) (time.Duration, error) {
	timeStr = strings.Trim(timeStr, " \r\n\t")
	var d time.Duration
	ms, err := strconv.ParseUint(timeStr, 10, 64)
	if err != nil {
		d, err = time.ParseDuration(timeStr)
		if err != nil {
			return d, err
		}
	} else {
		d = time.Duration(ms) * time.Millisecond
	}
	return d, nil
}
<|endoftext|>


<commit_before>// Copyright 2016, RadiantBlue Technologies, Inc.
//
// 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 main

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"github.com/venicegeo/pzsvc-exec/pzse"
	"github.com/venicegeo/pzsvc-exec/pzsvc"
)

func main() {

	s := pzsvc.Session{AppName: "pzsvc-taskworker", SessionID: "startup", LogRootDir: "pzsvc-exec"}
	pzsvc.LogAudit(s, s.AppName, "startup", s.AppName)

	if len(os.Args) < 2 {
		pzsvc.LogSimpleErr(s, "error: Insufficient parameters.  You must specify a config file.", nil)
		return
	}

	// First argument after the base call should be the path to the config file.
	// ReadFile returns the contents of the file as a byte buffer.
	configBuf, err := ioutil.ReadFile(os.Args[1])
	if err != nil {
		pzsvc.LogSimpleErr(s, "pzsvc-taskworker error in reading config: ", err)
		return
	}
	var configObj pzse.ConfigType
	err = json.Unmarshal(configBuf, &configObj)
	if err != nil {
		pzsvc.LogSimpleErr(s, "pzsvc-taskworker error in unmarshalling config: ", err)
		return
	}

	s.LogAudit = configObj.LogAudit
	if configObj.LogAudit {
		pzsvc.LogInfo(s, "Config: Audit logging enabled.")
	} else {
		pzsvc.LogInfo(s, "Config: Audit logging disabled.")
	}

	s.PzAddr = configObj.PzAddr
	if configObj.PzAddrEnVar != "" {
		newAddr := os.Getenv(configObj.PzAddrEnVar)
		if newAddr != "" {
			s.PzAddr = newAddr
		}
	}
	if s.PzAddr == "" {
		pzsvc.LogSimpleErr(s, "Config: Cannot work tasks.  Must have either a valid PzAddr, or a valid and populated PzAddrEnVar.", nil)
		return
	}

	if configObj.SvcName == "" {
		pzsvc.LogSimpleErr(s, "Config: Cannot work tasks without service name.", nil)
		return
	}

	if configObj.APIKeyEnVar == "" {
		pzsvc.LogSimpleErr(s, "Config: Cannot work tasks without valid APIKeyEnVar.", nil)
		return
	}
	apiKey := os.Getenv(configObj.APIKeyEnVar)
	if apiKey == "" {
		pzsvc.LogSimpleErr(s, "No API key at APIKeyEnVar.  Cannot work.", nil)
		return
	}
	s.PzAuth = "Basic " + base64.StdEncoding.EncodeToString([]byte(apiKey+":"))

	if configObj.NumProcs == 0 {
		pzsvc.LogInfo(s, "Config: No Proc number specified.  Defaulting to one at a time.")
		configObj.NumProcs = 1
	}

	if configObj.Port == 0 {
		pzsvc.LogInfo(s, "Config: No target Port specified.  Defaulting to 8080.")
		configObj.Port = 8080
	}

	svcID := ""
	for i := 0; svcID == "" && i < 10; i++ {
		svcID, err = pzsvc.FindMySvc(s, configObj.SvcName, configObj.PzAddr, s.PzAuth)
		if err != nil {
			pzsvc.LogSimpleErr(s, "Taskworker could not find Pz Service ID.  Initial Error: ", err)
			return
		}
		if svcID == "" && i < 9 {
			pzsvc.LogInfo(s, "Could not find service.  Will sleep and wait.")
			time.Sleep(15 * time.Second)
		}
	}
	if svcID == "" {
		pzsvc.LogSimpleErr(s, "Taskworker could not find Pz Service ID.  No error, just no service.", err)
		return
	}

	pzsvc.LogInfo(s, "Found target service.  ServiceID: "+svcID)
	for i := 0; i < configObj.NumProcs; i++ {
		go workerThread(s, configObj, svcID)
	}
	select {} //blocks forever

}

// WorkBody exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkBody struct {
	Content string `json:"content"`
}

// WorkDataInputs exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkDataInputs struct {
	Body WorkBody `json:"body"`
}

// WorkInData exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkInData struct {
	DataInputs WorkDataInputs `json:"dataInputs"`
}

// WorkSvcData exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkSvcData struct {
	Data  WorkInData `json:"data"`
	JobID string     `json:"jobId"`
}

// WorkOutData exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkOutData struct {
	SvcData WorkSvcData `json:"serviceData"`
}

func workerThread(s pzsvc.Session, configObj pzse.ConfigType, svcID string) {

	var (
		err       error
		failCount int
	)
	workAddr := fmt.Sprintf("http://localhost:%d/execute", configObj.Port)

	s.SessionID, err = pzsvc.PsuUUID()
	if err != nil {
		s.SessionID = "FailedSessionInit"
		pzsvc.LogSimpleErr(s, "psuUUID error: ", err)
		panic("Worker thread failed on uid generation.  Something is very wrong: " + err.Error())
	}
	pzsvc.LogInfo(s, "Worker thread initiated.")

	for {
		var pzJobObj struct {
			Data WorkOutData `json:"data"`
		}
		pzJobObj.Data = WorkOutData{SvcData: WorkSvcData{JobID: "", Data: WorkInData{DataInputs: WorkDataInputs{Body: WorkBody{Content: ""}}}}}

		byts, pErr := pzsvc.RequestKnownJSON("POST", "", configObj.PzAddr+"/service/"+svcID+"/task", s.PzAuth, &pzJobObj)
		if pErr != nil {
			pErr.Log(s, "Taskworker worker thread: error getting new task:")
			failCount++
			time.Sleep(time.Duration(10*failCount) * time.Second)
			continue
		}
		inpStr := pzJobObj.Data.SvcData.Data.DataInputs.Body.Content
		jobID := pzJobObj.Data.SvcData.JobID
		if inpStr != "" {
			pzsvc.LogInfo(s, "New Task Grabbed.  JobID: "+jobID)
			failCount = 0

			var outpByts []byte
			if configObj.JwtSecAuthURL != "" {
				// TODO: once JWT conversion exists as an option, handle it here.
				// jwtBody = content
				// call JwtSecAuthURL.  send jwtBody.  get response
				// outBody = response (more or less)
			}

			var respObj pzse.OutStruct
			pzsvc.LogAuditBuf(s, s.UserID, "http request - calling pzsvc-exec", inpStr, workAddr)
			outpByts, pErr := pzsvc.RequestKnownJSON("POST", inpStr, workAddr, "", &respObj)
			if pErr != nil {
				pErr.Log(s, "Error calling pzsvc-exec")
				sendExecResult(s, configObj.PzAddr, s.PzAuth, svcID, jobID, "Fail", nil)
			} else {
				pzsvc.LogAuditBuf(s, workAddr, "http response from pzsvc-exec", string(outpByts), s.UserID)
				sendExecResult(s, configObj.PzAddr, s.PzAuth, svcID, jobID, "Success", outpByts)
			}
			time.Sleep(10 * time.Second)

		} else {
			pzsvc.LogInfo(s, "No Task.  Sleeping now.  input: "+string(byts))
			time.Sleep(60 * time.Second)
		}
	}

}

func sendExecResult(s pzsvc.Session, pzAddr, pzAuth, svcID, jobID, status string, resJSON []byte) {
	outAddr := pzAddr + `/service/` + svcID + `/task/` + jobID

	pzsvc.LogInfo(s, "Sending Exec Results.  Status: "+status+".")
	if resJSON != nil {
		dataID, err := pzsvc.Ingest(s, "Output", "text", "pzsvc-taskworker", "", resJSON, nil)
		if err == nil {
			outStr := `{ "status" : "` + status + `", "result" : { "type" : "data", "dataId" : "` + dataID + `" } }`
			pzsvc.SubmitSinglePart("POST", outStr, outAddr, s.PzAuth)
			return
		}
		pzsvc.LogInfo(s, "Send Exec Results: Ingest failed.")
		status = "Fail"
	}

	outStr := `{ "status" : "` + status + `" }`
	pzsvc.SubmitSinglePart("POST", outStr, outAddr, s.PzAuth)
}
<commit_msg>fixing typoish thing win taskworker<commit_after>// Copyright 2016, RadiantBlue Technologies, Inc.
//
// 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 main

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"github.com/venicegeo/pzsvc-exec/pzse"
	"github.com/venicegeo/pzsvc-exec/pzsvc"
)

func main() {

	s := pzsvc.Session{AppName: "pzsvc-taskworker", SessionID: "startup", LogRootDir: "pzsvc-exec"}
	pzsvc.LogAudit(s, s.AppName, "startup", s.AppName)

	if len(os.Args) < 2 {
		pzsvc.LogSimpleErr(s, "error: Insufficient parameters.  You must specify a config file.", nil)
		return
	}

	// First argument after the base call should be the path to the config file.
	// ReadFile returns the contents of the file as a byte buffer.
	configBuf, err := ioutil.ReadFile(os.Args[1])
	if err != nil {
		pzsvc.LogSimpleErr(s, "pzsvc-taskworker error in reading config: ", err)
		return
	}
	var configObj pzse.ConfigType
	err = json.Unmarshal(configBuf, &configObj)
	if err != nil {
		pzsvc.LogSimpleErr(s, "pzsvc-taskworker error in unmarshalling config: ", err)
		return
	}

	s.LogAudit = configObj.LogAudit
	if configObj.LogAudit {
		pzsvc.LogInfo(s, "Config: Audit logging enabled.")
	} else {
		pzsvc.LogInfo(s, "Config: Audit logging disabled.")
	}

	s.PzAddr = configObj.PzAddr
	if configObj.PzAddrEnVar != "" {
		newAddr := os.Getenv(configObj.PzAddrEnVar)
		if newAddr != "" {
			s.PzAddr = newAddr
		}
	}
	if s.PzAddr == "" {
		pzsvc.LogSimpleErr(s, "Config: Cannot work tasks.  Must have either a valid PzAddr, or a valid and populated PzAddrEnVar.", nil)
		return
	}

	if configObj.SvcName == "" {
		pzsvc.LogSimpleErr(s, "Config: Cannot work tasks without service name.", nil)
		return
	}

	if configObj.APIKeyEnVar == "" {
		pzsvc.LogSimpleErr(s, "Config: Cannot work tasks without valid APIKeyEnVar.", nil)
		return
	}
	apiKey := os.Getenv(configObj.APIKeyEnVar)
	if apiKey == "" {
		pzsvc.LogSimpleErr(s, "No API key at APIKeyEnVar.  Cannot work.", nil)
		return
	}
	s.PzAuth = "Basic " + base64.StdEncoding.EncodeToString([]byte(apiKey+":"))

	if configObj.NumProcs == 0 {
		pzsvc.LogInfo(s, "Config: No Proc number specified.  Defaulting to one at a time.")
		configObj.NumProcs = 1
	}

	if configObj.Port == 0 {
		pzsvc.LogInfo(s, "Config: No target Port specified.  Defaulting to 8080.")
		configObj.Port = 8080
	}

	svcID := ""
	for i := 0; svcID == "" && i < 10; i++ {
		svcID, err = pzsvc.FindMySvc(s, configObj.SvcName, s.PzAddr, s.PzAuth)
		if err != nil {
			pzsvc.LogSimpleErr(s, "Taskworker could not find Pz Service ID.  Initial Error: ", err)
			return
		}
		if svcID == "" && i < 9 {
			pzsvc.LogInfo(s, "Could not find service.  Will sleep and wait.")
			time.Sleep(15 * time.Second)
		}
	}
	if svcID == "" {
		pzsvc.LogSimpleErr(s, "Taskworker could not find Pz Service ID.  No error, just no service.", err)
		return
	}

	pzsvc.LogInfo(s, "Found target service.  ServiceID: "+svcID)
	for i := 0; i < configObj.NumProcs; i++ {
		go workerThread(s, configObj, svcID)
	}
	select {} //blocks forever

}

// WorkBody exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkBody struct {
	Content string `json:"content"`
}

// WorkDataInputs exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkDataInputs struct {
	Body WorkBody `json:"body"`
}

// WorkInData exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkInData struct {
	DataInputs WorkDataInputs `json:"dataInputs"`
}

// WorkSvcData exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkSvcData struct {
	Data  WorkInData `json:"data"`
	JobID string     `json:"jobId"`
}

// WorkOutData exists as part of the response format of the Piazza job manager task request endpoint.
// specifically, it's one layer of the bit we care about.
type WorkOutData struct {
	SvcData WorkSvcData `json:"serviceData"`
}

func workerThread(s pzsvc.Session, configObj pzse.ConfigType, svcID string) {

	var (
		err       error
		failCount int
	)
	workAddr := fmt.Sprintf("http://localhost:%d/execute", configObj.Port)

	s.SessionID, err = pzsvc.PsuUUID()
	if err != nil {
		s.SessionID = "FailedSessionInit"
		pzsvc.LogSimpleErr(s, "psuUUID error: ", err)
		panic("Worker thread failed on uid generation.  Something is very wrong: " + err.Error())
	}
	pzsvc.LogInfo(s, "Worker thread initiated.")

	for {
		var pzJobObj struct {
			Data WorkOutData `json:"data"`
		}
		pzJobObj.Data = WorkOutData{SvcData: WorkSvcData{JobID: "", Data: WorkInData{DataInputs: WorkDataInputs{Body: WorkBody{Content: ""}}}}}

		byts, pErr := pzsvc.RequestKnownJSON("POST", "", s.PzAddr+"/service/"+svcID+"/task", s.PzAuth, &pzJobObj)
		if pErr != nil {
			pErr.Log(s, "Taskworker worker thread: error getting new task:")
			failCount++
			time.Sleep(time.Duration(10*failCount) * time.Second)
			continue
		}
		inpStr := pzJobObj.Data.SvcData.Data.DataInputs.Body.Content
		jobID := pzJobObj.Data.SvcData.JobID
		if inpStr != "" {
			pzsvc.LogInfo(s, "New Task Grabbed.  JobID: "+jobID)
			failCount = 0

			var outpByts []byte
			if configObj.JwtSecAuthURL != "" {
				// TODO: once JWT conversion exists as an option, handle it here.
				// jwtBody = content
				// call JwtSecAuthURL.  send jwtBody.  get response
				// outBody = response (more or less)
			}

			var respObj pzse.OutStruct
			pzsvc.LogAuditBuf(s, s.UserID, "http request - calling pzsvc-exec", inpStr, workAddr)
			outpByts, pErr := pzsvc.RequestKnownJSON("POST", inpStr, workAddr, "", &respObj)
			if pErr != nil {
				pErr.Log(s, "Error calling pzsvc-exec")
				sendExecResult(s, s.PzAddr, s.PzAuth, svcID, jobID, "Fail", nil)
			} else {
				pzsvc.LogAuditBuf(s, workAddr, "http response from pzsvc-exec", string(outpByts), s.UserID)
				sendExecResult(s, s.PzAddr, s.PzAuth, svcID, jobID, "Success", outpByts)
			}
			time.Sleep(10 * time.Second)

		} else {
			pzsvc.LogInfo(s, "No Task.  Sleeping now.  input: "+string(byts))
			time.Sleep(60 * time.Second)
		}
	}

}

func sendExecResult(s pzsvc.Session, pzAddr, pzAuth, svcID, jobID, status string, resJSON []byte) {
	outAddr := pzAddr + `/service/` + svcID + `/task/` + jobID

	pzsvc.LogInfo(s, "Sending Exec Results.  Status: "+status+".")
	if resJSON != nil {
		dataID, err := pzsvc.Ingest(s, "Output", "text", "pzsvc-taskworker", "", resJSON, nil)
		if err == nil {
			outStr := `{ "status" : "` + status + `", "result" : { "type" : "data", "dataId" : "` + dataID + `" } }`
			pzsvc.SubmitSinglePart("POST", outStr, outAddr, s.PzAuth)
			return
		}
		pzsvc.LogInfo(s, "Send Exec Results: Ingest failed.")
		status = "Fail"
	}

	outStr := `{ "status" : "` + status + `" }`
	pzsvc.SubmitSinglePart("POST", outStr, outAddr, s.PzAuth)
}
<|endoftext|>


<commit_before>package ws

import (
	"net/http"

	"github.com/gorilla/websocket"
)

func Web(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusInternalServerError)
		return
	}

	if origin := r.Header.Get("Origin"); origin != "http://127.0.0.1" {
		http.Error(w, "Origin not allowed", http.StatusForbidden)
		return
	}

	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", http.StatusBadRequest)
		return
	} else if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	for {
		messageType, p, err := ws.ReadMessage()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if err = ws.WriteMessage(messageType, p); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
}

func Runner(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusInternalServerError)
		return
	}
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", http.StatusBadRequest)
		return
	} else if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	for {
		messageType, p, err := ws.ReadMessage()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if err = ws.WriteMessage(messageType, p); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
}
<commit_msg>Added comments to exported func and commented blocking code now tests pass<commit_after>package ws

import (
	"net/http"

	"github.com/gorilla/websocket"
)

// Web handler websocket from web page side
func Web(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusInternalServerError)
		return
	}
	// TODO: FIX: "http://" used here! Maybe set "origin" option in /etc/apollo.conf
	if origin := r.Header.Get("Origin"); origin != "http://"+r.Host {
		http.Error(w, "Origin not allowed", http.StatusForbidden)
		return
	}
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", http.StatusBadRequest)
		return
	} else if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	ws.Close()
	/**for {
		messageType, p, err := ws.ReadMessage()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if err = ws.WriteMessage(messageType, p); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}**/
}

// Runner handler websocket from runner side
func Runner(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusInternalServerError)
		return
	}
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", http.StatusBadRequest)
		return
	} else if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	ws.Close()
	/**for {
		messageType, p, err := ws.ReadMessage()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if err = ws.WriteMessage(messageType, p); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}**/
}
<|endoftext|>