<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/streadway/amqp"
	"github.com/taskcluster/pulse-go/pulse"
	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-client-go/queue"
	"github.com/taskcluster/taskcluster-client-go/queueevents"
)

var (
	expiry queue.Time
	// all tests can share taskGroupId so we can view all test tasks in same
	// graph later for troubleshooting
	taskGroupId string = slugid.Nice()
)

func setup(t *testing.T) {
	// some basic setup...
	cwd, err := os.Getwd()
	if err != nil {
		t.Fatalf("Test failed during setup phase!")
	}
	TaskUser.HomeDir = filepath.Join(cwd, "test")

	expiry = queue.Time(time.Now().Add(time.Minute * 1))
}

func validateArtifacts(
	t *testing.T,
	payloadArtifacts []struct {
		Expires queue.Time `json:"expires"`
		Path    string     `json:"path"`
		Type    string     `json:"type"`
	},
	expected []Artifact) {

	// to test, create a dummy task run with given artifacts
	// and then call PayloadArtifacts() method to see what
	// artifacts would get uploaded...
	tr := &TaskRun{
		Payload: GenericWorkerPayload{
			Artifacts: payloadArtifacts,
		},
	}
	artifacts := tr.PayloadArtifacts()

	// compare expected vs actual artifacts by converting artifacts to strings...
	if fmt.Sprintf("%q", artifacts) != fmt.Sprintf("%q", expected) {
		t.Fatalf("Expected different artifacts to be generated...\nExpected:\n%q\nActual:\n%q", expected, artifacts)
	}
}

// See the test/SampleArtifacts subdirectory of this project. This simulates
// adding it as a directory artifact in a task payload, and checks that all
// files underneath this directory are discovered and created as s3 artifacts.
func TestDirectoryArtifacts(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "SampleArtifacts",
			Type:    "directory",
		}},

		// what we expect to discover on file system
		[]Artifact{
			S3Artifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/%%%/v/X",
					Expires:       expiry,
				},
				MimeType: "application/octet-stream",
			},
			S3Artifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/_/X.txt",
					Expires:       expiry,
				},
				MimeType: "text/plain; charset=utf-8",
			},
			S3Artifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/b/c/d.jpg",
					Expires:       expiry,
				},
				MimeType: "image/jpeg",
			},
		})
}

// Task payload specifies a file artifact which doesn't exist on worker
func TestMissingFileArtifact(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "TestMissingFileArtifact/no_such_file",
			Type:    "file",
		}},

		// what we expect to discover on file system
		[]Artifact{
			ErrorArtifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "TestMissingFileArtifact/no_such_file",
					Expires:       expiry,
				},
				Message: "Could not read file '" + filepath.Join(TaskUser.HomeDir, "TestMissingFileArtifact", "no_such_file") + "'",
				Reason:  "file-missing-on-worker",
			},
		})
}

// Task payload specifies a directory artifact which doesn't exist on worker
func TestMissingDirectoryArtifact(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "TestMissingDirectoryArtifact/no_such_dir",
			Type:    "directory",
		}},

		// what we expect to discover on file system
		[]Artifact{
			ErrorArtifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "TestMissingDirectoryArtifact/no_such_dir",
					Expires:       expiry,
				},
				Message: "Could not read directory '" + filepath.Join(TaskUser.HomeDir, "TestMissingDirectoryArtifact", "no_such_dir") + "'",
				Reason:  "file-missing-on-worker",
			},
		})
}

// Task payload specifies a file artifact which is actually a directory on worker
func TestFileArtifactIsDirectory(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "SampleArtifacts/b/c",
			Type:    "file",
		}},

		// what we expect to discover on file system
		[]Artifact{
			ErrorArtifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/b/c",
					Expires:       expiry,
				},
				Message: "File artifact '" + filepath.Join(TaskUser.HomeDir, "SampleArtifacts", "b", "c") + "' exists as a directory, not a file, on the worker",
				Reason:  "invalid-resource-on-worker",
			},
		})
}

// Task payload specifies a directory artifact which is a regular file on worker
func TestDirectoryArtifactIsFile(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "SampleArtifacts/b/c/d.jpg",
			Type:    "directory",
		}},

		// what we expect to discover on file system
		[]Artifact{
			ErrorArtifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/b/c/d.jpg",
					Expires:       expiry,
				},
				Message: "Directory artifact '" + filepath.Join(TaskUser.HomeDir, "SampleArtifacts", "b", "c", "d.jpg") + "' exists as a file, not a directory, on the worker",
				Reason:  "invalid-resource-on-worker",
			},
		})
}

func TestUpload(t *testing.T) {

	// check we have all the env vars we need to run this test
	clientId := os.Getenv("TASKCLUSTER_CLIENT_ID")
	accessToken := os.Getenv("TASKCLUSTER_ACCESS_TOKEN")
	certificate := os.Getenv("TASKCLUSTER_CERTIFICATE")
	if clientId == "" || accessToken == "" {
		t.Skip("Skipping test since TASKCLUSTER_CLIENT_ID and/or TASKCLUSTER_ACCESS_TOKEN env vars not set")
	}

	pulseUsername := os.Getenv("PULSE_USERNAME")
	pulsePassword := os.Getenv("PULSE_PASSWORD")
	if pulseUsername == "" || pulsePassword == "" {
		t.Skip("Skipping test since PULSE_USERNAME and/or PULSE_PASSWORD env vars are not set")
	}

	// define a unique workerType/provisionerId combination for this session
	provisionerId := "test-provisioner"
	// this should be sufficiently unique
	workerType := slugid.Nice()
	taskId := slugid.Nice()

	// configure the worker
	config = Config{
		AccessToken:                accessToken,
		Certificate:                certificate,
		ClientId:                   clientId,
		Debug:                      "*",
		ProvisionerId:              provisionerId,
		RefreshUrlsPrematurelySecs: 310,
		WorkerGroup:                "test-worker-group",
		WorkerId:                   "test-worker-id",
		WorkerType:                 workerType,
	}

	// get the worker started
	// killWorkerChan := runWorker()
	runWorker()

	artifactCreatedMessages := make(map[string]*queueevents.ArtifactCreatedMessage)
	// size 1 so that we don't block writing on taskCompleted
	artifactsCreatedChan := make(chan bool, 1)
	taskCompleted := make(chan bool)
	// timeout after 30 seconds - that should be plenty
	timeoutTimer := time.NewTimer(time.Second * 30)

	// start a listener for published artifacts
	// (uses PULSE_USERNAME, PULSE_PASSWORD and prod url)
	pulseConn := pulse.NewConnection("", "", "")
	pulseConn.Consume(
		"", // anonymous queue
		func(message interface{}, delivery amqp.Delivery) {
			switch message.(type) {
			case *queueevents.ArtifactCreatedMessage:
				a := message.(*queueevents.ArtifactCreatedMessage)
				artifactCreatedMessages[a.Artifact.Name] = a
				// finish after 3 artifacts have been created
				if len(artifactCreatedMessages) == 3 {
					// killWorkerChan <- true
					// pulseConn.AMQPConn.Close()
					artifactsCreatedChan <- true
				}
			case *queueevents.TaskCompletedMessage:
				taskCompleted <- true
			}
		},
		1,    // prefetch
		true, // auto-ack
		queueevents.ArtifactCreated{
			TaskId:        taskId,
			WorkerType:    workerType,
			ProvisionerId: provisionerId,
		},
		queueevents.TaskCompleted{
			TaskId:        taskId,
			WorkerType:    workerType,
			ProvisionerId: provisionerId,
		},
	)

	// create dummy task
	myQueue := queue.New(clientId, accessToken)
	myQueue.Certificate = certificate

	created := time.Now()
	// deadline in one days' time
	deadline := created.AddDate(0, 0, 1)
	// expiry in one month, in case we need test results
	expires := created.AddDate(0, 1, 0)

	td := &queue.TaskDefinitionRequest{
		Created:  queue.Time(created),
		Deadline: queue.Time(deadline),
		Expires:  queue.Time(expires),
		Extra:    json.RawMessage(`{}`),
		Metadata: struct {
			Description string `json:"description"`
			Name        string `json:"name"`
			Owner       string `json:"owner"`
			Source      string `json:"source"`
		}{
			Description: "Test task",
			Name:        "[TC] TestUpload",
			Owner:       "pmoore@mozilla.com",
			Source:      "https://github.com/taskcluster/generic-worker/blob/master/artifacts_test.go",
		},
		Payload: json.RawMessage(`

		{
			"command": [
				[
					"echo",
					"hello world!"
				]
			],
			"maxRunTime": 7200,
			"artifacts": [
				{
					"path": "SampleArtifacts/_/X.txt",
					"expires": "` + queue.Time(expires).String() + `",
					"type": "file"
				}
			]
		}

		`),
		ProvisionerId: provisionerId,
		Retries:       1,
		Routes:        []string{},
		SchedulerId:   "test-scheduler",
		Scopes:        []string{},
		Tags:          json.RawMessage(`{"createdForUser":"pmoore@mozilla.com"}`),
		Priority:      "normal",
		TaskGroupId:   taskGroupId,
		WorkerType:    workerType,
	}

	_, cs := myQueue.CreateTask(taskId, td)

	if cs.Error != nil {
		t.Fatalf("Suffered error when posting task to Queue in test setup:\n%s", cs.Error)
	}

	expectedArtifacts := map[string]string{
		"public/logs/all_commands.log":   "hello world!\n",
		"public/logs/command_000000.log": "hello world!\n",
		"SampleArtifacts/_/X.txt":        "test artifact\n",
	}

	// wait for task to complete, so we know artifact upload also completed
	select {
	case <-timeoutTimer.C:
		t.Fatalf("Test timed out waiting for artifacts to be published")
	case <-taskCompleted:
	}

	// now check artifact metadata is ok
	select {
	case <-timeoutTimer.C:
		t.Fatalf("Test timed out waiting for artifacts to be published")
	case <-taskCompleted:
	case <-artifactsCreatedChan:
		for artifact, _ := range expectedArtifacts {
			if A1 := artifactCreatedMessages[artifact]; A1 != nil {
				if A1.Artifact.ContentType != "text/plain; charset=utf-8" {
					t.Errorf("Artifact %s should have mime type 'text/plain; charset=utf-8' but has '%s'", artifact, A1.Artifact.ContentType)
				}
				if A1.Artifact.Expires.String() != queue.Time(expires).String() {
					t.Errorf("Artifact %s should have expiry '%s' but has '%s'", artifact, queue.Time(expires), A1.Artifact.Expires)
				}
			} else {
				t.Errorf("Artifact '%s' not created", artifact)
			}
		}
	}

	// now check content was uploaded to Amazon, and is correct
	for artifact, content := range expectedArtifacts {
		cs = myQueue.GetLatestArtifact(taskId, artifact)
		if cs.Error != nil {
			t.Fatalf("Error trying to fetch artifacts from Amazon...")
		}
		if cs.HttpResponseBody != content {
			t.Errorf("Artifact '%s': Was expecting content '%s' but found '%s'", artifact, content, cs.HttpResponseBody)
		}
	}
}
<commit_msg>cleaned up<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/streadway/amqp"
	"github.com/taskcluster/pulse-go/pulse"
	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-client-go/queue"
	"github.com/taskcluster/taskcluster-client-go/queueevents"
)

var (
	expiry queue.Time
	// all tests can share taskGroupId so we can view all test tasks in same
	// graph later for troubleshooting
	taskGroupId string = slugid.Nice()
)

func setup(t *testing.T) {
	// some basic setup...
	cwd, err := os.Getwd()
	if err != nil {
		t.Fatalf("Test failed during setup phase!")
	}
	TaskUser.HomeDir = filepath.Join(cwd, "test")

	expiry = queue.Time(time.Now().Add(time.Minute * 1))
}

func validateArtifacts(
	t *testing.T,
	payloadArtifacts []struct {
		Expires queue.Time `json:"expires"`
		Path    string     `json:"path"`
		Type    string     `json:"type"`
	},
	expected []Artifact) {

	// to test, create a dummy task run with given artifacts
	// and then call PayloadArtifacts() method to see what
	// artifacts would get uploaded...
	tr := &TaskRun{
		Payload: GenericWorkerPayload{
			Artifacts: payloadArtifacts,
		},
	}
	artifacts := tr.PayloadArtifacts()

	// compare expected vs actual artifacts by converting artifacts to strings...
	if fmt.Sprintf("%q", artifacts) != fmt.Sprintf("%q", expected) {
		t.Fatalf("Expected different artifacts to be generated...\nExpected:\n%q\nActual:\n%q", expected, artifacts)
	}
}

// See the test/SampleArtifacts subdirectory of this project. This simulates
// adding it as a directory artifact in a task payload, and checks that all
// files underneath this directory are discovered and created as s3 artifacts.
func TestDirectoryArtifacts(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "SampleArtifacts",
			Type:    "directory",
		}},

		// what we expect to discover on file system
		[]Artifact{
			S3Artifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/%%%/v/X",
					Expires:       expiry,
				},
				MimeType: "application/octet-stream",
			},
			S3Artifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/_/X.txt",
					Expires:       expiry,
				},
				MimeType: "text/plain; charset=utf-8",
			},
			S3Artifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/b/c/d.jpg",
					Expires:       expiry,
				},
				MimeType: "image/jpeg",
			},
		})
}

// Task payload specifies a file artifact which doesn't exist on worker
func TestMissingFileArtifact(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "TestMissingFileArtifact/no_such_file",
			Type:    "file",
		}},

		// what we expect to discover on file system
		[]Artifact{
			ErrorArtifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "TestMissingFileArtifact/no_such_file",
					Expires:       expiry,
				},
				Message: "Could not read file '" + filepath.Join(TaskUser.HomeDir, "TestMissingFileArtifact", "no_such_file") + "'",
				Reason:  "file-missing-on-worker",
			},
		})
}

// Task payload specifies a directory artifact which doesn't exist on worker
func TestMissingDirectoryArtifact(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "TestMissingDirectoryArtifact/no_such_dir",
			Type:    "directory",
		}},

		// what we expect to discover on file system
		[]Artifact{
			ErrorArtifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "TestMissingDirectoryArtifact/no_such_dir",
					Expires:       expiry,
				},
				Message: "Could not read directory '" + filepath.Join(TaskUser.HomeDir, "TestMissingDirectoryArtifact", "no_such_dir") + "'",
				Reason:  "file-missing-on-worker",
			},
		})
}

// Task payload specifies a file artifact which is actually a directory on worker
func TestFileArtifactIsDirectory(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "SampleArtifacts/b/c",
			Type:    "file",
		}},

		// what we expect to discover on file system
		[]Artifact{
			ErrorArtifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/b/c",
					Expires:       expiry,
				},
				Message: "File artifact '" + filepath.Join(TaskUser.HomeDir, "SampleArtifacts", "b", "c") + "' exists as a directory, not a file, on the worker",
				Reason:  "invalid-resource-on-worker",
			},
		})
}

// Task payload specifies a directory artifact which is a regular file on worker
func TestDirectoryArtifactIsFile(t *testing.T) {

	setup(t)
	validateArtifacts(t,

		// what appears in task payload
		[]struct {
			Expires queue.Time `json:"expires"`
			Path    string     `json:"path"`
			Type    string     `json:"type"`
		}{{
			Expires: expiry,
			Path:    "SampleArtifacts/b/c/d.jpg",
			Type:    "directory",
		}},

		// what we expect to discover on file system
		[]Artifact{
			ErrorArtifact{
				BaseArtifact: BaseArtifact{
					CanonicalPath: "SampleArtifacts/b/c/d.jpg",
					Expires:       expiry,
				},
				Message: "Directory artifact '" + filepath.Join(TaskUser.HomeDir, "SampleArtifacts", "b", "c", "d.jpg") + "' exists as a file, not a directory, on the worker",
				Reason:  "invalid-resource-on-worker",
			},
		})
}

func TestUpload(t *testing.T) {

	// check we have all the env vars we need to run this test
	clientId := os.Getenv("TASKCLUSTER_CLIENT_ID")
	accessToken := os.Getenv("TASKCLUSTER_ACCESS_TOKEN")
	certificate := os.Getenv("TASKCLUSTER_CERTIFICATE")
	if clientId == "" || accessToken == "" {
		t.Skip("Skipping test since TASKCLUSTER_CLIENT_ID and/or TASKCLUSTER_ACCESS_TOKEN env vars not set")
	}

	pulseUsername := os.Getenv("PULSE_USERNAME")
	pulsePassword := os.Getenv("PULSE_PASSWORD")
	if pulseUsername == "" || pulsePassword == "" {
		t.Skip("Skipping test since PULSE_USERNAME and/or PULSE_PASSWORD env vars are not set")
	}

	// define a unique workerType/provisionerId combination for this session
	provisionerId := "test-provisioner"
	// this should be sufficiently unique
	workerType := slugid.Nice()
	taskId := slugid.Nice()

	// configure the worker
	config = Config{
		AccessToken:                accessToken,
		Certificate:                certificate,
		ClientId:                   clientId,
		Debug:                      "*",
		ProvisionerId:              provisionerId,
		RefreshUrlsPrematurelySecs: 310,
		WorkerGroup:                "test-worker-group",
		WorkerId:                   "test-worker-id",
		WorkerType:                 workerType,
	}

	// get the worker started
	// killWorkerChan := runWorker()
	runWorker()

	artifactCreatedMessages := make(map[string]*queueevents.ArtifactCreatedMessage)
	// size 1 so that we don't block writing on taskCompleted
	artifactsCreatedChan := make(chan bool, 1)
	taskCompleted := make(chan bool)
	// timeout after 30 seconds - that should be plenty
	timeoutTimer := time.NewTimer(time.Second * 30)

	// start a listener for published artifacts
	// (uses PULSE_USERNAME, PULSE_PASSWORD and prod url)
	pulseConn := pulse.NewConnection("", "", "")
	pulseConn.Consume(
		"", // anonymous queue
		func(message interface{}, delivery amqp.Delivery) {
			switch message.(type) {
			case *queueevents.ArtifactCreatedMessage:
				a := message.(*queueevents.ArtifactCreatedMessage)
				artifactCreatedMessages[a.Artifact.Name] = a
				// finish after 3 artifacts have been created
				if len(artifactCreatedMessages) == 3 {
					// killWorkerChan <- true
					// pulseConn.AMQPConn.Close()
					artifactsCreatedChan <- true
				}
			case *queueevents.TaskCompletedMessage:
				taskCompleted <- true
			}
		},
		1,    // prefetch
		true, // auto-ack
		queueevents.ArtifactCreated{
			TaskId:        taskId,
			WorkerType:    workerType,
			ProvisionerId: provisionerId,
		},
		queueevents.TaskCompleted{
			TaskId:        taskId,
			WorkerType:    workerType,
			ProvisionerId: provisionerId,
		},
	)

	// create dummy task
	myQueue := queue.New(clientId, accessToken)
	myQueue.Certificate = certificate

	created := time.Now()
	// deadline in one days' time
	deadline := created.AddDate(0, 0, 1)
	// expiry in one month, in case we need test results
	expires := created.AddDate(0, 1, 0)

	td := &queue.TaskDefinitionRequest{
		Created:  queue.Time(created),
		Deadline: queue.Time(deadline),
		Expires:  queue.Time(expires),
		Extra:    json.RawMessage(`{}`),
		Metadata: struct {
			Description string `json:"description"`
			Name        string `json:"name"`
			Owner       string `json:"owner"`
			Source      string `json:"source"`
		}{
			Description: "Test task",
			Name:        "[TC] TestUpload",
			Owner:       "pmoore@mozilla.com",
			Source:      "https://github.com/taskcluster/generic-worker/blob/master/artifacts_test.go",
		},
		Payload: json.RawMessage(`

		{
			"command": [
				[
					"echo",
					"hello world!"
				]
			],
			"maxRunTime": 7200,
			"artifacts": [
				{
					"path": "SampleArtifacts/_/X.txt",
					"expires": "` + queue.Time(expires).String() + `",
					"type": "file"
				}
			]
		}

		`),
		ProvisionerId: provisionerId,
		Retries:       1,
		Routes:        []string{},
		SchedulerId:   "test-scheduler",
		Scopes:        []string{},
		Tags:          json.RawMessage(`{"createdForUser":"pmoore@mozilla.com"}`),
		Priority:      "normal",
		TaskGroupId:   taskGroupId,
		WorkerType:    workerType,
	}

	_, cs := myQueue.CreateTask(taskId, td)

	if cs.Error != nil {
		t.Fatalf("Suffered error when posting task to Queue in test setup:\n%s", cs.Error)
	}

	expectedArtifacts := map[string]string{
		"public/logs/all_commands.log":   "hello world!\n",
		"public/logs/command_000000.log": "hello world!\n",
		"SampleArtifacts/_/X.txt":        "test artifact\n",
	}

	// wait for task to complete, so we know artifact upload also completed
	select {
	case <-timeoutTimer.C:
		t.Fatalf("Test timed out waiting for artifacts to be published")
	case <-taskCompleted:
	}

	// now check artifact metadata is ok
	select {
	case <-timeoutTimer.C:
		t.Fatalf("Test timed out waiting for artifacts to be published")
	case <-artifactsCreatedChan:
		for artifact, _ := range expectedArtifacts {
			if a := artifactCreatedMessages[artifact]; a != nil {
				if a.Artifact.ContentType != "text/plain; charset=utf-8" {
					t.Errorf("Artifact %s should have mime type 'text/plain; charset=utf-8' but has '%s'", artifact, a.Artifact.ContentType)
				}
				if a.Artifact.Expires.String() != queue.Time(expires).String() {
					t.Errorf("Artifact %s should have expiry '%s' but has '%s'", artifact, queue.Time(expires), a.Artifact.Expires)
				}
			} else {
				t.Errorf("Artifact '%s' not created", artifact)
			}
		}
	}

	// now check content was uploaded to Amazon, and is correct
	for artifact, content := range expectedArtifacts {
		cs = myQueue.GetLatestArtifact(taskId, artifact)
		if cs.Error != nil {
			t.Fatalf("Error trying to fetch artifacts from Amazon...")
		}
		if cs.HttpResponseBody != content {
			t.Errorf("Artifact '%s': Was expecting content '%s' but found '%s'", artifact, content, cs.HttpResponseBody)
		}
	}
}
<|endoftext|>


<commit_before>package ltsv

import (
	"bufio"
	"io"
)

// Writer is LTSV writer.
type Writer struct {
	writer *bufio.Writer
	labels []string
}

// NewWriter creates new LTSV writer.
func NewWriter(w io.Writer, labels ...string) *Writer {
	return &Writer{
		writer: bufio.NewWriter(w),
		labels: labels,
	}
}

// Write writes a LTSV line.
func (w *Writer) Write(values ...string) error {
	for i, l := range w.labels {
		if i != 0 {
			w.writer.WriteRune('\t')
		}
		w.writer.WriteString(l)
		w.writer.WriteRune(':')
		if i < len(values) {
			w.writer.WriteString(values[i])
		}
	}
	w.writer.WriteRune('\n')
	return w.writer.Flush()
}

func Write(w io.StringWriter, props []Property) error {
	for i, p := range props {
		if i != 0 {
			_, err := w.WriteString("\t")
			if err != nil {
				return err
			}
		}
		_, err := w.WriteString(p.Label)
		if err != nil {
			return err
		}
		_, err = w.WriteString(":")
		if err != nil {
			return err
		}
		_, err = w.WriteString(p.Value)
		if err != nil {
			return err
		}
	}
	_, err := w.WriteString("\n")
	return err
}
<commit_msg>escape values of LTSV<commit_after>package ltsv

import (
	"bufio"
	"bytes"
	"io"
	"strings"
)

// Writer is LTSV writer.
type Writer struct {
	writer *bufio.Writer
	labels []string
}

// NewWriter creates new LTSV writer.
func NewWriter(w io.Writer, labels ...string) *Writer {
	return &Writer{
		writer: bufio.NewWriter(w),
		labels: labels,
	}
}

// Write writes a LTSV line.
func (w *Writer) Write(values ...string) error {
	for i, l := range w.labels {
		if i != 0 {
			w.writer.WriteRune('\t')
		}
		w.writer.WriteString(l)
		w.writer.WriteRune(':')
		if i < len(values) {
			w.writer.WriteString(escape(values[i]))
		}
	}
	w.writer.WriteRune('\n')
	return w.writer.Flush()
}

func escape(s string) string {
	if !strings.ContainsAny(s, "\\\t\n\r") {
		return s
	}
	bb := &bytes.Buffer{}
	for _, r := range s {
		switch r {
		case '\\':
			bb.WriteString(`\\`)
		case '\t':
			bb.WriteString(`\t`)
		case '\n':
			bb.WriteString(`\n`)
		case '\r':
			bb.WriteString(`\r`)
		default:
			bb.WriteRune(r)
		}
	}
	return bb.String()
}

func Write(w io.StringWriter, props []Property) error {
	for i, p := range props {
		if i != 0 {
			_, err := w.WriteString("\t")
			if err != nil {
				return err
			}
		}
		_, err := w.WriteString(p.Label)
		if err != nil {
			return err
		}
		_, err = w.WriteString(":")
		if err != nil {
			return err
		}
		_, err = w.WriteString(p.Value)
		if err != nil {
			return err
		}
	}
	_, err := w.WriteString("\n")
	return err
}
<|endoftext|>


<commit_before>package list

import (
	"bytes"
	"fmt"
	"os"
	"regexp"
	"strings"
	"text/tabwriter"

	"github.com/blacknon/lssh/conf"
	runewidth "github.com/mattn/go-runewidth"
	termbox "github.com/nsf/termbox-go"
)

type ListArrayInfo struct {
	Name    string
	Connect string
	Note    string
}

// Draw Line
func drawLine(x, y int, str string, colorNum int, backColorNum int) {
	color := termbox.Attribute(colorNum + 1)
	backColor := termbox.Attribute(backColorNum + 1)
	// View Multi-Byte
	for _, c := range str {
		termbox.SetCell(x, y, c, color, backColor)
		x += runewidth.RuneWidth(c)
	}
}

// Draw List
func draw(serverNameList []string, selectCursor int, searchText string) {
	headLine := 2
	defaultColor := 255
	defaultBackColor := 255
	pronpt := "lssh>>"
	termbox.Clear(termbox.Attribute(defaultColor+1), termbox.Attribute(defaultBackColor+1))

	// Get Terminal Size
	_, height := termbox.Size()
	lineHeight := height - headLine

	// Set View List Range
	viewFirstLine := (selectCursor/lineHeight)*lineHeight + 1
	viewLastLine := viewFirstLine + lineHeight
	var serverViewList []string
	if viewLastLine > len(serverNameList) {
		serverViewList = serverNameList[viewFirstLine:]
	} else {
		serverViewList = serverNameList[viewFirstLine:viewLastLine]
	}
	selectViewCursor := selectCursor - viewFirstLine + 1

	// View Head
	drawLine(0, 0, pronpt, 3, defaultBackColor)
	drawLine(len(pronpt), 0, searchText, defaultColor, defaultBackColor)
	drawLine(headLine, 1, serverNameList[0], 3, defaultBackColor)

	// View List
	for k, v := range serverViewList {
		cursorColor := defaultColor
		cursorBackColor := defaultBackColor
		if k == selectViewCursor {
			// Select line color
			cursorColor = 0
			cursorBackColor = 2
		}

		viewListData := v
		drawLine(2, k+2, viewListData, cursorColor, cursorBackColor)
		k += 1
	}

	// Multi-Byte SetCursor
	x := 0
	for _, c := range searchText {
		x += runewidth.RuneWidth(c)
	}
	termbox.SetCursor(len(pronpt)+x, 0)
	termbox.Flush()
}

// Create View List Data (use text/tabwriter)
func getListData(serverNameList []string, serverList conf.Config) (listData []string) {
	buffer := &bytes.Buffer{}
	tabWriterBuffer := new(tabwriter.Writer)
	tabWriterBuffer.Init(buffer, 0, 4, 8, ' ', 0)
	fmt.Fprintln(tabWriterBuffer, "ServerName \tConnect Infomation \tNote \t")

	serverName := ""
	connectInfomation := ""
	serverNote := ""
	for _, v := range serverNameList {
		serverName = v
		connectInfomation = serverList.Server[v].User + "@" + serverList.Server[v].Addr
		serverNote = serverList.Server[v].Note
		fmt.Fprintln(tabWriterBuffer, serverName+"\t"+connectInfomation+"\t"+serverNote+"\t")
	}
	tabWriterBuffer.Flush()
	line, err := buffer.ReadString('\n')
	for err == nil {
		str := strings.Replace(line, "\t", " ", -1)
		listData = append(listData, str)
		line, err = buffer.ReadString('\n')
	}
	return listData
}

func insertRune(text string, inputRune rune) (returnText string) {
	returnText = text + string(inputRune)
	return
}

func deleteRune(text string) (returnText string) {
	s := text
	sc := []rune(s)
	returnText = string(sc[:(len(sc) - 1)])
	return
}

func getFilterListData(searchText string, listData []string) (returnListData []string) {
	// SearchText Bounds Space
	searchWords := strings.Fields(searchText)
	searchWordMeta := ""
	re := regexp.MustCompile(searchWordMeta)
	r := listData[1:]
	line := ""
	loopListData := []string{}
	returnListData = append(returnListData, listData[0])

	// if No searchWords
	if len(searchWords) == 0 {
		returnListData = listData
		return returnListData
	}

	for i := 0; i < len(searchWords); i += 1 {
		searchWordMeta = regexp.QuoteMeta(strings.ToLower(searchWords[i]))
		re = regexp.MustCompile(searchWordMeta)
		loopListData = []string{}

		for j := 0; j < len(r); j += 1 {
			line += string(r[j])
			if re.MatchString(strings.ToLower(line)) {
				loopListData = append(loopListData, line)
			}
			line = ""
		}
		r = loopListData

	}
	returnListData = append(returnListData, loopListData...)
	return returnListData
}

//func filterListData(searchKeyword string,listData []string) (returnListData []string){
//	line := ""
//	for i := 0; i < len(listData); i += 1 {
//		line += string(listData[i])
//		if re.MatchString(strings.ToLower(line)) {
//			returnListData = append(returnListData, line)
//		}
//		line = ""
//	}
//}

func pollEvent(serverNameList []string, serverList conf.Config) (lineData string) {
	defer termbox.Close()
	listData := getListData(serverNameList, serverList)
	selectline := 0
	headLine := 2

	_, height := termbox.Size()
	lineHeight := height - headLine

	searchText := ""

	filterListData := getFilterListData(searchText, listData)
	draw(filterListData, selectline, searchText)
	for {
		switch ev := termbox.PollEvent(); ev.Type {

		// Get Key Event
		case termbox.EventKey:
			switch ev.Key {
			// ESC or Ctrl + C Key (Exit)
			case termbox.KeyEsc, termbox.KeyCtrlC:
				termbox.Close()
				os.Exit(0)

			// AllowUp Key
			case termbox.KeyArrowUp:
				if selectline > 0 {
					selectline -= 1
				}
				draw(filterListData, selectline, searchText)

			// AllowDown Key
			case termbox.KeyArrowDown:
				if selectline < len(filterListData)-headLine {
					selectline += 1
				}
				draw(filterListData, selectline, searchText)

			// AllowRight Key
			case termbox.KeyArrowRight:
				if ((selectline+lineHeight)/lineHeight)*lineHeight <= len(filterListData) {
					selectline = ((selectline + lineHeight) / lineHeight) * lineHeight
				}
				draw(filterListData, selectline, searchText)

			// AllowLeft Key
			case termbox.KeyArrowLeft:
				if ((selectline-lineHeight)/lineHeight)*lineHeight >= 0 {
					selectline = ((selectline - lineHeight) / lineHeight) * lineHeight
				}

				draw(filterListData, selectline, searchText)

			// Enter Key
			case termbox.KeyEnter:
				lineData = strings.Fields(filterListData[selectline+1])[0]
				return

			// BackSpace Key
			case termbox.KeyBackspace, termbox.KeyBackspace2:
				if len(searchText) > 0 {
					searchText = deleteRune(searchText)
					filterListData = getFilterListData(searchText, listData)
					if selectline > len(filterListData) {
						selectline = len(filterListData)
					}
					if selectline < 0 {
						selectline = 0
					}
					draw(filterListData, selectline, searchText)
				}

			// Space Key
			case termbox.KeySpace:
				searchText = searchText + " "
				draw(filterListData, selectline, searchText)

			// Other Key
			default:
				if ev.Ch != 0 {
					searchText = insertRune(searchText, ev.Ch)
					filterListData = getFilterListData(searchText, listData)
					if selectline > len(filterListData)-headLine {
						selectline = len(filterListData) - headLine
					}
					draw(filterListData, selectline, searchText)
				}
			}
		default:
			draw(filterListData, selectline, searchText)
		}
	}
}

func DrawList(serverNameList []string, serverList conf.Config) (lineName string) {
	err := termbox.Init()
	if err != nil {
		panic(err)
	}

	lineName = pollEvent(serverNameList, serverList)
	return lineName
}
<commit_msg>Change variable declaration<commit_after>package list

import (
	"bytes"
	"fmt"
	"os"
	"regexp"
	"strings"
	"text/tabwriter"

	"github.com/blacknon/lssh/conf"
	runewidth "github.com/mattn/go-runewidth"
	termbox "github.com/nsf/termbox-go"
)

type ListArrayInfo struct {
	Name    string
	Connect string
	Note    string
}

func drawLine(x, y int, str string, colorNum int, backColorNum int) {
	color := termbox.Attribute(colorNum + 1)
	backColor := termbox.Attribute(backColorNum + 1)
	// View Multi-Byte
	for _, c := range str {
		termbox.SetCell(x, y, c, color, backColor)
		x += runewidth.RuneWidth(c)
	}
}

// Draw List
func draw(serverNameList []string, selectCursor int, searchText string) {
	headLine := 2
	leftMargin := 2
	defaultColor := 255
	defaultBackColor := 255
	pronpt := "lssh>>"
	termbox.Clear(termbox.Attribute(defaultColor+1), termbox.Attribute(defaultBackColor+1))

	// Get Terminal Size
	_, height := termbox.Size()
	lineHeight := height - headLine

	// Set View List Range
	viewFirstLine := (selectCursor/lineHeight)*lineHeight + 1
	viewLastLine := viewFirstLine + lineHeight
	var serverViewList []string
	if viewLastLine > len(serverNameList) {
		serverViewList = serverNameList[viewFirstLine:]
	} else {
		serverViewList = serverNameList[viewFirstLine:viewLastLine]
	}
	selectViewCursor := selectCursor - viewFirstLine + 1

	// View Head
	drawLine(0, 0, pronpt, 3, defaultBackColor)
	drawLine(len(pronpt), 0, searchText, defaultColor, defaultBackColor)
	drawLine(leftMargin, 1, serverNameList[0], 3, defaultBackColor)

	// View List
	for listKey, listValue := range serverViewList {
		cursorColor := defaultColor
		cursorBackColor := defaultBackColor
		if listKey == selectViewCursor {
			// Select line color
			cursorColor = 0
			cursorBackColor = 2
		}

		viewListData := listValue
		drawLine(leftMargin, listKey+headLine, viewListData, cursorColor, cursorBackColor)
		listKey += 1
	}

	// Multi-Byte SetCursor
	x := 0
	for _, c := range searchText {
		x += runewidth.RuneWidth(c)
	}
	termbox.SetCursor(len(pronpt)+x, 0)
	termbox.Flush()
}

// Create View List Data (use text/tabwriter)
func getListData(serverNameList []string, serverList conf.Config) (listData []string) {
	buffer := &bytes.Buffer{}
	tabWriterBuffer := new(tabwriter.Writer)
	tabWriterBuffer.Init(buffer, 0, 4, 8, ' ', 0)
	fmt.Fprintln(tabWriterBuffer, "ServerName \tConnect Infomation \tNote \t")

	for _, v := range serverNameList {
		serverName := v
		connectInfomation := serverList.Server[v].User + "@" + serverList.Server[v].Addr
		serverNote := serverList.Server[v].Note
		fmt.Fprintln(tabWriterBuffer, serverName+"\t"+connectInfomation+"\t"+serverNote+"\t")
	}
	tabWriterBuffer.Flush()
	line, err := buffer.ReadString('\n')
	for err == nil {
		str := strings.Replace(line, "\t", " ", -1)
		listData = append(listData, str)
		line, err = buffer.ReadString('\n')
	}
	return listData
}

func insertRune(text string, inputRune rune) (returnText string) {
	returnText = text + string(inputRune)
	return
}

func deleteRune(text string) (returnText string) {
	s := text
	sc := []rune(s)
	returnText = string(sc[:(len(sc) - 1)])
	return
}

func getFilterListData(searchText string, listData []string) (returnListData []string) {
	// SearchText Bounds Space
	searchWords := strings.Fields(searchText)
	r := listData[1:]
	line := ""
	loopListData := []string{}
	returnListData = append(returnListData, listData[0])

	// if No searchWords
	if len(searchWords) == 0 {
		returnListData = listData
		return returnListData
	}

	for i := 0; i < len(searchWords); i += 1 {
		searchWordMeta := regexp.QuoteMeta(strings.ToLower(searchWords[i]))
		re := regexp.MustCompile(searchWordMeta)
		loopListData = []string{}

		for j := 0; j < len(r); j += 1 {
			line += string(r[j])
			if re.MatchString(strings.ToLower(line)) {
				loopListData = append(loopListData, line)
			}
			line = ""
		}
		r = loopListData
	}
	returnListData = append(returnListData, loopListData...)
	return returnListData
}

func pollEvent(serverNameList []string, serverList conf.Config) (lineData string) {
	defer termbox.Close()
	listData := getListData(serverNameList, serverList)
	selectline := 0
	headLine := 2

	_, height := termbox.Size()
	lineHeight := height - headLine

	searchText := ""

	filterListData := getFilterListData(searchText, listData)
	draw(filterListData, selectline, searchText)
	for {
		switch ev := termbox.PollEvent(); ev.Type {

		// Get Key Event
		case termbox.EventKey:
			switch ev.Key {
			// ESC or Ctrl + C Key (Exit)
			case termbox.KeyEsc, termbox.KeyCtrlC:
				termbox.Close()
				os.Exit(0)

			// AllowUp Key
			case termbox.KeyArrowUp:
				if selectline > 0 {
					selectline -= 1
				}
				draw(filterListData, selectline, searchText)

			// AllowDown Key
			case termbox.KeyArrowDown:
				if selectline < len(filterListData)-headLine {
					selectline += 1
				}
				draw(filterListData, selectline, searchText)

			// AllowRight Key
			case termbox.KeyArrowRight:
				if ((selectline+lineHeight)/lineHeight)*lineHeight <= len(filterListData) {
					selectline = ((selectline + lineHeight) / lineHeight) * lineHeight
				}
				draw(filterListData, selectline, searchText)

			// AllowLeft Key
			case termbox.KeyArrowLeft:
				if ((selectline-lineHeight)/lineHeight)*lineHeight >= 0 {
					selectline = ((selectline - lineHeight) / lineHeight) * lineHeight
				}

				draw(filterListData, selectline, searchText)

			// Enter Key
			case termbox.KeyEnter:
				lineData = strings.Fields(filterListData[selectline+1])[0]
				return

			// BackSpace Key
			case termbox.KeyBackspace, termbox.KeyBackspace2:
				if len(searchText) > 0 {
					searchText = deleteRune(searchText)
					filterListData = getFilterListData(searchText, listData)
					if selectline > len(filterListData) {
						selectline = len(filterListData)
					}
					if selectline < 0 {
						selectline = 0
					}
					draw(filterListData, selectline, searchText)
				}

			// Space Key
			case termbox.KeySpace:
				searchText = searchText + " "
				draw(filterListData, selectline, searchText)

			// Other Key
			default:
				if ev.Ch != 0 {
					searchText = insertRune(searchText, ev.Ch)
					filterListData = getFilterListData(searchText, listData)
					if selectline > len(filterListData)-headLine {
						selectline = len(filterListData) - headLine
					}
					draw(filterListData, selectline, searchText)
				}
			}
		default:
			draw(filterListData, selectline, searchText)
		}
	}
}

func DrawList(serverNameList []string, serverList conf.Config) (lineName string) {
	err := termbox.Init()
	if err != nil {
		panic(err)
	}

	lineName = pollEvent(serverNameList, serverList)
	return lineName
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/device"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/state"
	storagePools "github.com/lxc/lxd/lxd/storage"
	storageDrivers "github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/idmap"
	"github.com/lxc/lxd/shared/logger"
	"github.com/lxc/lxd/shared/version"
)

func init() {
	// Expose storageVolumeMount to the device package as StorageVolumeMount.
	device.StorageVolumeMount = storageVolumeMount

	// Expose storageVolumeUmount to the device package as StorageVolumeUmount.
	device.StorageVolumeUmount = storageVolumeUmount

	// Expose storageRootFSApplyQuota to the device package as StorageRootFSApplyQuota.
	device.StorageRootFSApplyQuota = storageRootFSApplyQuota
}

// Simply cache used to storage the activated drivers on this LXD instance. This
// allows us to avoid querying the database everytime and API call is made.
var storagePoolDriversCacheVal atomic.Value
var storagePoolDriversCacheLock sync.Mutex

func readStoragePoolDriversCache() map[string]string {
	drivers := storagePoolDriversCacheVal.Load()
	if drivers == nil {
		return map[string]string{}
	}

	return drivers.(map[string]string)
}

func storagePoolVolumeAttachPrepare(s *state.State, poolName string, volumeName string, volumeType int, c instance.Container) error {
	// Load the DB records
	poolID, pool, err := s.Cluster.StoragePoolGet(poolName)
	if err != nil {
		return err
	}

	_, volume, err := s.Cluster.StoragePoolNodeVolumeGetTypeByProject("default", volumeName, volumeType, poolID)
	if err != nil {
		return err
	}

	poolVolumePut := volume.Writable()

	// Check if unmapped
	if shared.IsTrue(poolVolumePut.Config["security.unmapped"]) {
		// No need to look at containers and maps for unmapped volumes
		return nil
	}

	// Get the on-disk idmap for the volume
	var lastIdmap *idmap.IdmapSet
	if poolVolumePut.Config["volatile.idmap.last"] != "" {
		lastIdmap, err = idmapsetFromString(poolVolumePut.Config["volatile.idmap.last"])
		if err != nil {
			logger.Errorf("Failed to unmarshal last idmapping: %s", poolVolumePut.Config["volatile.idmap.last"])
			return err
		}
	}

	var nextIdmap *idmap.IdmapSet
	nextJsonMap := "[]"
	if !shared.IsTrue(poolVolumePut.Config["security.shifted"]) {
		// Get the container's idmap
		if c.IsRunning() {
			nextIdmap, err = c.CurrentIdmap()
		} else {
			nextIdmap, err = c.NextIdmap()
		}
		if err != nil {
			return err
		}

		if nextIdmap != nil {
			nextJsonMap, err = idmapsetToJSON(nextIdmap)
			if err != nil {
				return err
			}
		}
	}
	poolVolumePut.Config["volatile.idmap.next"] = nextJsonMap

	// Get mountpoint of storage volume
	remapPath := storagePools.GetStoragePoolVolumeMountPoint(poolName, volumeName)

	if !nextIdmap.Equals(lastIdmap) {
		logger.Debugf("Shifting storage volume")

		if !shared.IsTrue(poolVolumePut.Config["security.shifted"]) {
			volumeUsedBy, err := storagePoolVolumeUsedByInstancesGet(s, "default", poolName, volumeName)
			if err != nil {
				return err
			}

			if len(volumeUsedBy) > 1 {
				for _, ctName := range volumeUsedBy {
					instt, err := instance.LoadByProjectAndName(s, c.Project(), ctName)
					if err != nil {
						continue
					}

					if instt.Type() != instancetype.Container {
						continue
					}

					ct := instt.(instance.Container)

					var ctNextIdmap *idmap.IdmapSet
					if ct.IsRunning() {
						ctNextIdmap, err = ct.CurrentIdmap()
					} else {
						ctNextIdmap, err = ct.NextIdmap()
					}
					if err != nil {
						return fmt.Errorf("Failed to retrieve idmap of container")
					}

					if !nextIdmap.Equals(ctNextIdmap) {
						return fmt.Errorf("Idmaps of container %v and storage volume %v are not identical", ctName, volumeName)
					}
				}
			} else if len(volumeUsedBy) == 1 {
				// If we're the only one who's attached that container
				// we can shift the storage volume.
				// I'm not sure if we want some locking here.
				if volumeUsedBy[0] != c.Name() {
					return fmt.Errorf("idmaps of container and storage volume are not identical")
				}
			}
		}

		// Unshift rootfs
		if lastIdmap != nil {
			var err error

			if pool.Driver == "zfs" {
				err = lastIdmap.UnshiftRootfs(remapPath, storageDrivers.ShiftZFSSkipper)
			} else {
				err = lastIdmap.UnshiftRootfs(remapPath, nil)
			}

			if err != nil {
				logger.Errorf("Failed to unshift \"%s\"", remapPath)
				return err
			}

			logger.Debugf("Unshifted \"%s\"", remapPath)
		}

		// Shift rootfs
		if nextIdmap != nil {
			var err error

			if pool.Driver == "zfs" {
				err = nextIdmap.ShiftRootfs(remapPath, storageDrivers.ShiftZFSSkipper)
			} else {
				err = nextIdmap.ShiftRootfs(remapPath, nil)
			}

			if err != nil {
				logger.Errorf("Failed to shift \"%s\"", remapPath)
				return err
			}

			logger.Debugf("Shifted \"%s\"", remapPath)
		}
		logger.Debugf("Shifted storage volume")
	}

	jsonIdmap := "[]"
	if nextIdmap != nil {
		var err error
		jsonIdmap, err = idmapsetToJSON(nextIdmap)
		if err != nil {
			logger.Errorf("Failed to marshal idmap")
			return err
		}
	}

	// Update last idmap
	poolVolumePut.Config["volatile.idmap.last"] = jsonIdmap

	err = s.Cluster.StoragePoolVolumeUpdateByProject("default", volumeName, volumeType, poolID, poolVolumePut.Description, poolVolumePut.Config)
	if err != nil {
		return err
	}

	return nil
}

func resetContainerDiskIdmap(container instance.Container, srcIdmap *idmap.IdmapSet) error {
	dstIdmap, err := container.DiskIdmap()
	if err != nil {
		return err
	}

	if dstIdmap == nil {
		dstIdmap = new(idmap.IdmapSet)
	}

	if !srcIdmap.Equals(dstIdmap) {
		var jsonIdmap string
		if srcIdmap != nil {
			idmapBytes, err := json.Marshal(srcIdmap.Idmap)
			if err != nil {
				return err
			}
			jsonIdmap = string(idmapBytes)
		} else {
			jsonIdmap = "[]"
		}

		err := container.VolatileSet(map[string]string{"volatile.last_state.idmap": jsonIdmap})
		if err != nil {
			return err
		}
	}

	return nil
}

func setupStorageDriver(s *state.State, forceCheck bool) error {
	pools, err := s.Cluster.StoragePoolsNotPending()
	if err != nil {
		if err == db.ErrNoSuchObject {
			logger.Debugf("No existing storage pools detected")
			return nil
		}
		logger.Debugf("Failed to retrieve existing storage pools")
		return err
	}

	// In case the daemon got killed during upgrade we will already have a
	// valid storage pool entry but it might have gotten messed up and so we
	// cannot perform StoragePoolCheck(). This case can be detected by
	// looking at the patches db: If we already have a storage pool defined
	// but the upgrade somehow got messed up then there will be no
	// "storage_api" entry in the db.
	if len(pools) > 0 && !forceCheck {
		appliedPatches, err := s.Node.Patches()
		if err != nil {
			return err
		}

		if !shared.StringInSlice("storage_api", appliedPatches) {
			logger.Warnf("Incorrectly applied \"storage_api\" patch, skipping storage pool initialization as it might be corrupt")
			return nil
		}

	}

	for _, poolName := range pools {
		logger.Debugf("Initializing and checking storage pool %q", poolName)
		errPrefix := fmt.Sprintf("Failed initializing storage pool %q", poolName)

		pool, err := storagePools.GetPoolByName(s, poolName)
		if err != nil {
			return errors.Wrap(err, errPrefix)
		}

		_, err = pool.Mount()
		if err != nil {
			return errors.Wrap(err, errPrefix)
		}
	}

	// Update the storage drivers cache in api_1.0.go.
	storagePoolDriversCacheUpdate(s)
	return nil
}

func storagePoolDriversCacheUpdate(s *state.State) {
	// Get a list of all storage drivers currently in use
	// on this LXD instance. Only do this when we do not already have done
	// this once to avoid unnecessarily querying the db. All subsequent
	// updates of the cache will be done when we create or delete storage
	// pools in the db. Since this is a rare event, this cache
	// implementation is a classic frequent-read, rare-update case so
	// copy-on-write semantics without locking in the read case seems
	// appropriate. (Should be cheaper then querying the db all the time,
	// especially if we keep adding more storage drivers.)

	drivers, err := s.Cluster.StoragePoolsGetDrivers()
	if err != nil && err != db.ErrNoSuchObject {
		return
	}

	data := map[string]string{}

	// Get the driver info.
	info := storageDrivers.SupportedDrivers(s)
	for _, entry := range info {
		if shared.StringInSlice(entry.Name, drivers) {
			data[entry.Name] = entry.Version
		}
	}

	// Prepare the cache entries.
	backends := []string{}
	for k, v := range data {
		backends = append(backends, fmt.Sprintf("%s %s", k, v))
	}

	// Update the user agent.
	version.UserAgentStorageBackends(backends)

	storagePoolDriversCacheLock.Lock()
	storagePoolDriversCacheVal.Store(data)
	storagePoolDriversCacheLock.Unlock()

	return
}

// storageVolumeMount initialises a new storage interface and checks the pool and volume are
// mounted. If they are not then they are mounted.
func storageVolumeMount(state *state.State, poolName string, volumeName string, volumeTypeName string, inst instance.Instance) error {
	c, ok := inst.(instance.Container)
	if !ok {
		return fmt.Errorf("Received non-LXC container instance")
	}

	volumeType, err := storagePools.VolumeTypeNameToType(volumeTypeName)
	if err != nil {
		return err
	}

	pool, err := storagePools.GetPoolByName(state, poolName)
	if err != nil {
		return err
	}

	// Mount the storage volume.
	ourMount, err := pool.MountCustomVolume(volumeName, nil)
	if err != nil {
		return err
	}

	revert := true
	if ourMount {
		defer func() {
			if !revert {
				return
			}

			pool.UnmountCustomVolume(volumeName, nil)
		}()
	}

	// Custom storage volumes do not currently support projects, so hardcode "default" project.
	err = storagePoolVolumeAttachPrepare(state, poolName, volumeName, volumeType, c)
	if err != nil {
		return err
	}

	revert = false
	return nil
}

// storageVolumeUmount unmounts a storage volume on a pool.
func storageVolumeUmount(state *state.State, poolName string, volumeName string, volumeType int) error {
	pool, err := storagePools.GetPoolByName(state, poolName)
	if err != nil {
		return err
	}

	_, err = pool.UnmountCustomVolume(volumeName, nil)
	if err != nil {
		return err
	}

	return nil
}

// storageRootFSApplyQuota applies a quota to an instance if it can, if it cannot then it will
// return false indicating that the quota needs to be stored in volatile to be applied on next boot.
func storageRootFSApplyQuota(state *state.State, inst instance.Instance, size string) error {
	pool, err := storagePools.GetPoolByInstance(state, inst)
	if err != nil {
		return err
	}

	err = pool.SetInstanceQuota(inst, size, nil)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>lxd/storage: idmap.JSONUnmarshal usage<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/device"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/state"
	storagePools "github.com/lxc/lxd/lxd/storage"
	storageDrivers "github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/idmap"
	"github.com/lxc/lxd/shared/logger"
	"github.com/lxc/lxd/shared/version"
)

func init() {
	// Expose storageVolumeMount to the device package as StorageVolumeMount.
	device.StorageVolumeMount = storageVolumeMount

	// Expose storageVolumeUmount to the device package as StorageVolumeUmount.
	device.StorageVolumeUmount = storageVolumeUmount

	// Expose storageRootFSApplyQuota to the device package as StorageRootFSApplyQuota.
	device.StorageRootFSApplyQuota = storageRootFSApplyQuota
}

// Simply cache used to storage the activated drivers on this LXD instance. This
// allows us to avoid querying the database everytime and API call is made.
var storagePoolDriversCacheVal atomic.Value
var storagePoolDriversCacheLock sync.Mutex

func readStoragePoolDriversCache() map[string]string {
	drivers := storagePoolDriversCacheVal.Load()
	if drivers == nil {
		return map[string]string{}
	}

	return drivers.(map[string]string)
}

func storagePoolVolumeAttachPrepare(s *state.State, poolName string, volumeName string, volumeType int, c instance.Container) error {
	// Load the DB records
	poolID, pool, err := s.Cluster.StoragePoolGet(poolName)
	if err != nil {
		return err
	}

	_, volume, err := s.Cluster.StoragePoolNodeVolumeGetTypeByProject("default", volumeName, volumeType, poolID)
	if err != nil {
		return err
	}

	poolVolumePut := volume.Writable()

	// Check if unmapped
	if shared.IsTrue(poolVolumePut.Config["security.unmapped"]) {
		// No need to look at containers and maps for unmapped volumes
		return nil
	}

	// Get the on-disk idmap for the volume
	var lastIdmap *idmap.IdmapSet
	if poolVolumePut.Config["volatile.idmap.last"] != "" {
		lastIdmap, err = idmap.JSONUnmarshal(poolVolumePut.Config["volatile.idmap.last"])
		if err != nil {
			logger.Errorf("Failed to unmarshal last idmapping: %s", poolVolumePut.Config["volatile.idmap.last"])
			return err
		}
	}

	var nextIdmap *idmap.IdmapSet
	nextJsonMap := "[]"
	if !shared.IsTrue(poolVolumePut.Config["security.shifted"]) {
		// Get the container's idmap
		if c.IsRunning() {
			nextIdmap, err = c.CurrentIdmap()
		} else {
			nextIdmap, err = c.NextIdmap()
		}
		if err != nil {
			return err
		}

		if nextIdmap != nil {
			nextJsonMap, err = idmapsetToJSON(nextIdmap)
			if err != nil {
				return err
			}
		}
	}
	poolVolumePut.Config["volatile.idmap.next"] = nextJsonMap

	// Get mountpoint of storage volume
	remapPath := storagePools.GetStoragePoolVolumeMountPoint(poolName, volumeName)

	if !nextIdmap.Equals(lastIdmap) {
		logger.Debugf("Shifting storage volume")

		if !shared.IsTrue(poolVolumePut.Config["security.shifted"]) {
			volumeUsedBy, err := storagePoolVolumeUsedByInstancesGet(s, "default", poolName, volumeName)
			if err != nil {
				return err
			}

			if len(volumeUsedBy) > 1 {
				for _, ctName := range volumeUsedBy {
					instt, err := instance.LoadByProjectAndName(s, c.Project(), ctName)
					if err != nil {
						continue
					}

					if instt.Type() != instancetype.Container {
						continue
					}

					ct := instt.(instance.Container)

					var ctNextIdmap *idmap.IdmapSet
					if ct.IsRunning() {
						ctNextIdmap, err = ct.CurrentIdmap()
					} else {
						ctNextIdmap, err = ct.NextIdmap()
					}
					if err != nil {
						return fmt.Errorf("Failed to retrieve idmap of container")
					}

					if !nextIdmap.Equals(ctNextIdmap) {
						return fmt.Errorf("Idmaps of container %v and storage volume %v are not identical", ctName, volumeName)
					}
				}
			} else if len(volumeUsedBy) == 1 {
				// If we're the only one who's attached that container
				// we can shift the storage volume.
				// I'm not sure if we want some locking here.
				if volumeUsedBy[0] != c.Name() {
					return fmt.Errorf("idmaps of container and storage volume are not identical")
				}
			}
		}

		// Unshift rootfs
		if lastIdmap != nil {
			var err error

			if pool.Driver == "zfs" {
				err = lastIdmap.UnshiftRootfs(remapPath, storageDrivers.ShiftZFSSkipper)
			} else {
				err = lastIdmap.UnshiftRootfs(remapPath, nil)
			}

			if err != nil {
				logger.Errorf("Failed to unshift \"%s\"", remapPath)
				return err
			}

			logger.Debugf("Unshifted \"%s\"", remapPath)
		}

		// Shift rootfs
		if nextIdmap != nil {
			var err error

			if pool.Driver == "zfs" {
				err = nextIdmap.ShiftRootfs(remapPath, storageDrivers.ShiftZFSSkipper)
			} else {
				err = nextIdmap.ShiftRootfs(remapPath, nil)
			}

			if err != nil {
				logger.Errorf("Failed to shift \"%s\"", remapPath)
				return err
			}

			logger.Debugf("Shifted \"%s\"", remapPath)
		}
		logger.Debugf("Shifted storage volume")
	}

	jsonIdmap := "[]"
	if nextIdmap != nil {
		var err error
		jsonIdmap, err = idmapsetToJSON(nextIdmap)
		if err != nil {
			logger.Errorf("Failed to marshal idmap")
			return err
		}
	}

	// Update last idmap
	poolVolumePut.Config["volatile.idmap.last"] = jsonIdmap

	err = s.Cluster.StoragePoolVolumeUpdateByProject("default", volumeName, volumeType, poolID, poolVolumePut.Description, poolVolumePut.Config)
	if err != nil {
		return err
	}

	return nil
}

func resetContainerDiskIdmap(container instance.Container, srcIdmap *idmap.IdmapSet) error {
	dstIdmap, err := container.DiskIdmap()
	if err != nil {
		return err
	}

	if dstIdmap == nil {
		dstIdmap = new(idmap.IdmapSet)
	}

	if !srcIdmap.Equals(dstIdmap) {
		var jsonIdmap string
		if srcIdmap != nil {
			idmapBytes, err := json.Marshal(srcIdmap.Idmap)
			if err != nil {
				return err
			}
			jsonIdmap = string(idmapBytes)
		} else {
			jsonIdmap = "[]"
		}

		err := container.VolatileSet(map[string]string{"volatile.last_state.idmap": jsonIdmap})
		if err != nil {
			return err
		}
	}

	return nil
}

func setupStorageDriver(s *state.State, forceCheck bool) error {
	pools, err := s.Cluster.StoragePoolsNotPending()
	if err != nil {
		if err == db.ErrNoSuchObject {
			logger.Debugf("No existing storage pools detected")
			return nil
		}
		logger.Debugf("Failed to retrieve existing storage pools")
		return err
	}

	// In case the daemon got killed during upgrade we will already have a
	// valid storage pool entry but it might have gotten messed up and so we
	// cannot perform StoragePoolCheck(). This case can be detected by
	// looking at the patches db: If we already have a storage pool defined
	// but the upgrade somehow got messed up then there will be no
	// "storage_api" entry in the db.
	if len(pools) > 0 && !forceCheck {
		appliedPatches, err := s.Node.Patches()
		if err != nil {
			return err
		}

		if !shared.StringInSlice("storage_api", appliedPatches) {
			logger.Warnf("Incorrectly applied \"storage_api\" patch, skipping storage pool initialization as it might be corrupt")
			return nil
		}

	}

	for _, poolName := range pools {
		logger.Debugf("Initializing and checking storage pool %q", poolName)
		errPrefix := fmt.Sprintf("Failed initializing storage pool %q", poolName)

		pool, err := storagePools.GetPoolByName(s, poolName)
		if err != nil {
			return errors.Wrap(err, errPrefix)
		}

		_, err = pool.Mount()
		if err != nil {
			return errors.Wrap(err, errPrefix)
		}
	}

	// Update the storage drivers cache in api_1.0.go.
	storagePoolDriversCacheUpdate(s)
	return nil
}

func storagePoolDriversCacheUpdate(s *state.State) {
	// Get a list of all storage drivers currently in use
	// on this LXD instance. Only do this when we do not already have done
	// this once to avoid unnecessarily querying the db. All subsequent
	// updates of the cache will be done when we create or delete storage
	// pools in the db. Since this is a rare event, this cache
	// implementation is a classic frequent-read, rare-update case so
	// copy-on-write semantics without locking in the read case seems
	// appropriate. (Should be cheaper then querying the db all the time,
	// especially if we keep adding more storage drivers.)

	drivers, err := s.Cluster.StoragePoolsGetDrivers()
	if err != nil && err != db.ErrNoSuchObject {
		return
	}

	data := map[string]string{}

	// Get the driver info.
	info := storageDrivers.SupportedDrivers(s)
	for _, entry := range info {
		if shared.StringInSlice(entry.Name, drivers) {
			data[entry.Name] = entry.Version
		}
	}

	// Prepare the cache entries.
	backends := []string{}
	for k, v := range data {
		backends = append(backends, fmt.Sprintf("%s %s", k, v))
	}

	// Update the user agent.
	version.UserAgentStorageBackends(backends)

	storagePoolDriversCacheLock.Lock()
	storagePoolDriversCacheVal.Store(data)
	storagePoolDriversCacheLock.Unlock()

	return
}

// storageVolumeMount initialises a new storage interface and checks the pool and volume are
// mounted. If they are not then they are mounted.
func storageVolumeMount(state *state.State, poolName string, volumeName string, volumeTypeName string, inst instance.Instance) error {
	c, ok := inst.(instance.Container)
	if !ok {
		return fmt.Errorf("Received non-LXC container instance")
	}

	volumeType, err := storagePools.VolumeTypeNameToType(volumeTypeName)
	if err != nil {
		return err
	}

	pool, err := storagePools.GetPoolByName(state, poolName)
	if err != nil {
		return err
	}

	// Mount the storage volume.
	ourMount, err := pool.MountCustomVolume(volumeName, nil)
	if err != nil {
		return err
	}

	revert := true
	if ourMount {
		defer func() {
			if !revert {
				return
			}

			pool.UnmountCustomVolume(volumeName, nil)
		}()
	}

	// Custom storage volumes do not currently support projects, so hardcode "default" project.
	err = storagePoolVolumeAttachPrepare(state, poolName, volumeName, volumeType, c)
	if err != nil {
		return err
	}

	revert = false
	return nil
}

// storageVolumeUmount unmounts a storage volume on a pool.
func storageVolumeUmount(state *state.State, poolName string, volumeName string, volumeType int) error {
	pool, err := storagePools.GetPoolByName(state, poolName)
	if err != nil {
		return err
	}

	_, err = pool.UnmountCustomVolume(volumeName, nil)
	if err != nil {
		return err
	}

	return nil
}

// storageRootFSApplyQuota applies a quota to an instance if it can, if it cannot then it will
// return false indicating that the quota needs to be stored in volatile to be applied on next boot.
func storageRootFSApplyQuota(state *state.State, inst instance.Instance, size string) error {
	pool, err := storagePools.GetPoolByInstance(state, inst)
	if err != nil {
		return err
	}

	err = pool.SetInstanceQuota(inst, size, nil)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package list

import (
	"bytes"
	"fmt"
	"math"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/tabwriter"

	"github.com/blacknon/lssh/conf"
	runewidth "github.com/mattn/go-runewidth"
	termbox "github.com/nsf/termbox-go"
)

type ListArrayInfo struct {
	Name    string
	Connect string
	Note    string
}

func drawLine(x, y int, str string, colorNum int, backColorNum int) {
	color := termbox.Attribute(colorNum + 1)
	backColor := termbox.Attribute(backColorNum + 1)
	// View Multi-Byte
	for _, char := range str {
		termbox.SetCell(x, y, char, color, backColor)
		x += runewidth.RuneWidth(char)
	}
}

// toggle select line (multi select)
func toggleList(selectedList []string, newLine string) (toggledSelectedList []string) {
	//
	//result := []int{}
	addFlag := true
	for _, selectedLine := range selectedList {
		if selectedLine != newLine {
			toggledSelectedList = append(toggledSelectedList, selectedLine)
		} else {
			addFlag = false
		}
	}
	if addFlag == true {
		toggledSelectedList = append(toggledSelectedList, newLine)
	}
	return
}

func allToggle(allFlag bool, selectedList []string, addList []string) (allSelectedList []string) {
	// selectedList in allSelectedList
	for _, selectedLine := range selectedList {
		allSelectedList = append(allSelectedList, selectedLine)
	}

	// allFlag is False
	if allFlag == false {
		for _, addLine := range addList {
			addData := strings.Fields(addLine)[0]
			allSelectedList = append(allSelectedList, addData)
		}
		return
	} else {
		for _, addLine := range addList {
			addData := strings.Fields(addLine)[0]
			allSelectedList = toggleList(allSelectedList, addData)
		}
		return
	}
}

func drawFilterLine(x, y int, str string, colorNum int, backColorNum int, keywordColorNum int, searchText string) {
	// SearchText Bounds Space
	searchWords := strings.Fields(searchText)

	for i := 0; i < len(searchWords); i += 1 {
		searchLowLine := strings.ToLower(str)
		searchKeyword := strings.ToLower(searchWords[i])
		searchKeywordLen := len(searchKeyword)
		searchKeywordCount := strings.Count(searchLowLine, searchKeyword)

		charLocation := 0
		for j := 0; j < searchKeywordCount; j += 1 {
			searchLineData := ""

			// Countermeasure "slice bounds out of range"
			if charLocation < len(str) {
				searchLineData = str[charLocation:]
			}
			searchLineDataStr := string(searchLineData)
			searchKeywordIndex := strings.Index(strings.ToLower(searchLineDataStr), searchKeyword)

			charLocation = charLocation + searchKeywordIndex
			keyword := ""

			// Countermeasure "slice bounds out of range"
			if charLocation < len(str) {
				keyword = str[charLocation : charLocation+searchKeywordLen]
			}

			// Get Multibyte Charctor Location
			multibyteStrCheckLine := str[:charLocation]
			multiByteCharLocation := 0
			for _, multiByteChar := range multibyteStrCheckLine {
				multiByteCharLocation += runewidth.RuneWidth(multiByteChar)
			}

			drawLine(x+multiByteCharLocation, y, keyword, keywordColorNum, backColorNum)
			charLocation = charLocation + searchKeywordLen
		}
	}
}

// Draw List
func draw(serverNameList []string, lineData []string, selectCursor int, searchText string) {
	headLine := 2
	leftMargin := 2
	defaultColor := 255
	defaultBackColor := 255
	pronpt := "lssh>>"
	termbox.Clear(termbox.Attribute(defaultColor+1), termbox.Attribute(defaultBackColor+1))

	// Get Terminal Size
	_, height := termbox.Size()
	lineHeight := height - headLine

	// Set View List Range
	viewFirstLine := (selectCursor/lineHeight)*lineHeight + 1
	viewLastLine := viewFirstLine + lineHeight
	var serverViewList []string
	if viewLastLine > len(serverNameList) {
		serverViewList = serverNameList[viewFirstLine:]
	} else {
		serverViewList = serverNameList[viewFirstLine:viewLastLine]
	}
	selectViewCursor := selectCursor - viewFirstLine + 1

	// View Head
	drawLine(0, 0, pronpt, 3, defaultBackColor)
	drawLine(len(pronpt), 0, searchText, defaultColor, defaultBackColor)
	drawLine(leftMargin, 1, serverNameList[0], 3, defaultBackColor)

	// Get View List Max Length
	serverValueMaxLength := float64(0)
	for _, viewListValue := range serverViewList {
		serverValueMaxLength = math.Max(serverValueMaxLength, float64(len(viewListValue)))
	}

	// View List
	for listKey, listValue := range serverViewList {
		paddingListValue := fmt.Sprintf("%-"+strconv.FormatFloat(serverValueMaxLength, 'g', -1, 64)+"s", listValue)
		// Set cursor color
		cursorColor := defaultColor
		cursorBackColor := defaultBackColor
		keywordColor := 5

		for _, selectedLine := range lineData {
			if strings.Split(listValue, " ")[0] == selectedLine {
				cursorColor = 0
				cursorBackColor = 6
			}
		}

		if listKey == selectViewCursor {
			// Select line color
			cursorColor = 0
			cursorBackColor = 2
		}

		// Draw filter line
		drawLine(leftMargin, listKey+headLine, paddingListValue, cursorColor, cursorBackColor)

		// Keyword Highlight
		drawFilterLine(leftMargin, listKey+headLine, paddingListValue, cursorColor, cursorBackColor, keywordColor, searchText)
		listKey += 1
	}

	// Multi-Byte SetCursor
	x := 0
	for _, c := range searchText {
		x += runewidth.RuneWidth(c)
	}
	termbox.SetCursor(len(pronpt)+x, 0)
	termbox.Flush()
}

// Create View List Data (use text/tabwriter)
func getListData(serverNameList []string, serverList conf.Config) (listData []string) {
	buffer := &bytes.Buffer{}
	tabWriterBuffer := new(tabwriter.Writer)
	tabWriterBuffer.Init(buffer, 0, 4, 8, ' ', 0)
	fmt.Fprintln(tabWriterBuffer, "ServerName \tConnect Infomation \tNote \t")

	// Check serverNote length
	//serverNoteMaxLength := float64(0)
	//for _, key := range serverNameList {
	//	serverNoteMaxLength = math.Max(serverNoteMaxLength, float64(len(serverList.Server[key].Note)))
	//}

	// Create list table
	for _, key := range serverNameList {
		serverName := key
		connectInfomation := serverList.Server[key].User + "@" + serverList.Server[key].Addr
		serverNote := serverList.Server[key].Note
		//serverNote := fmt.Sprintf("%-"+strconv.FormatFloat(serverNoteMaxLength, 'g', -1, 64)+"s", serverList.Server[key].Note)
		//serverNote := fmt.Sprintf("%-"+string(serverNoteMaxLength)+"s", serverList.Server[key].Note)
		//aaa := "%" + string(3) + "s"
		//serverNote := fmt.Sprintf(aaa, serverList.Server[key].Note)

		fmt.Fprintln(tabWriterBuffer, serverName+"\t"+connectInfomation+"\t"+serverNote)
		// fmt.Fprintln(tabWriterBuffer, serverName+"\t"+connectInfomation+"\t"+serverNote+"\t)

	}
	tabWriterBuffer.Flush()
	line, err := buffer.ReadString('\n')
	for err == nil {
		str := strings.Replace(line, "\t", " ", -1)
		listData = append(listData, str)
		line, err = buffer.ReadString('\n')
	}
	return listData
}

func insertRune(text string, inputRune rune) (returnText string) {
	returnText = text + string(inputRune)
	return
}

func deleteRune(text string) (returnText string) {
	s := text
	sc := []rune(s)
	returnText = string(sc[:(len(sc) - 1)])
	return
}

func getFilterListData(searchText string, listData []string) (returnListData []string) {
	// SearchText Bounds Space
	searchWords := strings.Fields(searchText)
	r := listData[1:]
	line := ""
	loopListData := []string{}
	returnListData = append(returnListData, listData[0])

	// if No searchWords
	if len(searchWords) == 0 {
		returnListData = listData
		return returnListData
	}

	for i := 0; i < len(searchWords); i += 1 {
		searchWordMeta := regexp.QuoteMeta(strings.ToLower(searchWords[i]))
		re := regexp.MustCompile(searchWordMeta)
		loopListData = []string{}

		for j := 0; j < len(r); j += 1 {
			line += string(r[j])
			if re.MatchString(strings.ToLower(line)) {
				loopListData = append(loopListData, line)
			}
			line = ""
		}
		r = loopListData
	}
	returnListData = append(returnListData, loopListData...)
	return returnListData
}

func pollEvent(serverNameList []string, cmdFlag bool, serverList conf.Config) (lineData []string) {
	defer termbox.Close()
	listData := getListData(serverNameList, serverList)
	selectline := 0
	headLine := 2

	_, height := termbox.Size()
	lineHeight := height - headLine

	searchText := ""
	allFlag := false

	filterListData := getFilterListData(searchText, listData)
	draw(filterListData, lineData, selectline, searchText)
	for {
		switch ev := termbox.PollEvent(); ev.Type {

		// Get Key Event
		case termbox.EventKey:
			switch ev.Key {
			// ESC or Ctrl + C Key (Exit)
			case termbox.KeyEsc, termbox.KeyCtrlC:
				termbox.Close()
				os.Exit(0)

			// AllowUp Key
			case termbox.KeyArrowUp:
				if selectline > 0 {
					selectline -= 1
				}
				draw(filterListData, lineData, selectline, searchText)

			// AllowDown Key
			case termbox.KeyArrowDown:
				if selectline < len(filterListData)-headLine {
					selectline += 1
				}
				draw(filterListData, lineData, selectline, searchText)

			// AllowRight Key
			case termbox.KeyArrowRight:
				if ((selectline+lineHeight)/lineHeight)*lineHeight <= len(filterListData) {
					selectline = ((selectline + lineHeight) / lineHeight) * lineHeight
				}
				draw(filterListData, lineData, selectline, searchText)

			// AllowLeft Key
			case termbox.KeyArrowLeft:
				if ((selectline-lineHeight)/lineHeight)*lineHeight >= 0 {
					selectline = ((selectline - lineHeight) / lineHeight) * lineHeight
				}

				draw(filterListData, lineData, selectline, searchText)

			// Ctrl + x Key(select)
			case termbox.KeyCtrlX:
				if cmdFlag == true {
					lineData = toggleList(lineData, strings.Fields(filterListData[selectline+1])[0])
				}

				draw(filterListData, lineData, selectline, searchText)

			// Ctrl + a Key(all select)
			case termbox.KeyCtrlA:
				if cmdFlag == true {
					lineData = allToggle(allFlag, lineData, filterListData[1:])
				}

				// allFlag Toggle
				if allFlag == false {
					allFlag = true
				} else {
					allFlag = false
				}

				draw(filterListData, lineData, selectline, searchText)

			// Enter Key
			case termbox.KeyEnter:
				if len(lineData) == 0 {
					lineData = append(lineData, strings.Fields(filterListData[selectline+1])[0])
				}
				return

			// BackSpace Key
			case termbox.KeyBackspace, termbox.KeyBackspace2:
				if len(searchText) > 0 {
					searchText = deleteRune(searchText)
					filterListData = getFilterListData(searchText, listData)
					if selectline > len(filterListData) {
						selectline = len(filterListData)
					}
					if selectline < 0 {
						selectline = 0
					}
					allFlag = false
					draw(filterListData, lineData, selectline, searchText)
				}

			// Space Key
			case termbox.KeySpace:
				searchText = searchText + " "
				draw(filterListData, lineData, selectline, searchText)

			// Other Key
			default:
				if ev.Ch != 0 {
					searchText = insertRune(searchText, ev.Ch)
					filterListData = getFilterListData(searchText, listData)
					if selectline > len(filterListData)-headLine {
						selectline = len(filterListData) - headLine
					}
					allFlag = false
					draw(filterListData, lineData, selectline, searchText)
				}
			}
		default:
			draw(filterListData, lineData, selectline, searchText)
		}
	}
}

func DrawList(serverNameList []string, cmdFlag bool, serverList conf.Config) (lineName []string) {
	if err := termbox.Init(); err != nil {
		panic(err)
	}

	lineName = pollEvent(serverNameList, cmdFlag, serverList)
	return lineName
}
<commit_msg>select line note length 1000 set<commit_after>package list

import (
	"bytes"
	"fmt"
	"os"
	"regexp"
	"strings"
	"text/tabwriter"

	"github.com/blacknon/lssh/conf"
	runewidth "github.com/mattn/go-runewidth"
	termbox "github.com/nsf/termbox-go"
)

type ListArrayInfo struct {
	Name    string
	Connect string
	Note    string
}

func drawLine(x, y int, str string, colorNum int, backColorNum int) {
	color := termbox.Attribute(colorNum + 1)
	backColor := termbox.Attribute(backColorNum + 1)
	// View Multi-Byte
	for _, char := range str {
		termbox.SetCell(x, y, char, color, backColor)
		x += runewidth.RuneWidth(char)
	}
}

// toggle select line (multi select)
func toggleList(selectedList []string, newLine string) (toggledSelectedList []string) {
	//result := []int{}
	addFlag := true
	for _, selectedLine := range selectedList {
		if selectedLine != newLine {
			toggledSelectedList = append(toggledSelectedList, selectedLine)
		} else {
			addFlag = false
		}
	}
	if addFlag == true {
		toggledSelectedList = append(toggledSelectedList, newLine)
	}
	return
}

func allToggle(allFlag bool, selectedList []string, addList []string) (allSelectedList []string) {
	// selectedList in allSelectedList
	for _, selectedLine := range selectedList {
		allSelectedList = append(allSelectedList, selectedLine)
	}

	// allFlag is False
	if allFlag == false {
		for _, addLine := range addList {
			addData := strings.Fields(addLine)[0]
			allSelectedList = append(allSelectedList, addData)
		}
		return
	} else {
		for _, addLine := range addList {
			addData := strings.Fields(addLine)[0]
			allSelectedList = toggleList(allSelectedList, addData)
		}
		return
	}
}

func drawFilterLine(x, y int, str string, colorNum int, backColorNum int, keywordColorNum int, searchText string) {
	// SearchText Bounds Space
	searchWords := strings.Fields(searchText)

	for i := 0; i < len(searchWords); i += 1 {
		searchLowLine := strings.ToLower(str)
		searchKeyword := strings.ToLower(searchWords[i])
		searchKeywordLen := len(searchKeyword)
		searchKeywordCount := strings.Count(searchLowLine, searchKeyword)

		charLocation := 0
		for j := 0; j < searchKeywordCount; j += 1 {
			searchLineData := ""

			// Countermeasure "slice bounds out of range"
			if charLocation < len(str) {
				searchLineData = str[charLocation:]
			}
			searchLineDataStr := string(searchLineData)
			searchKeywordIndex := strings.Index(strings.ToLower(searchLineDataStr), searchKeyword)

			charLocation = charLocation + searchKeywordIndex
			keyword := ""

			// Countermeasure "slice bounds out of range"
			if charLocation < len(str) {
				keyword = str[charLocation : charLocation+searchKeywordLen]
			}

			// Get Multibyte Charctor Location
			multibyteStrCheckLine := str[:charLocation]
			multiByteCharLocation := 0
			for _, multiByteChar := range multibyteStrCheckLine {
				multiByteCharLocation += runewidth.RuneWidth(multiByteChar)
			}

			drawLine(x+multiByteCharLocation, y, keyword, keywordColorNum, backColorNum)
			charLocation = charLocation + searchKeywordLen
		}
	}
}

// Draw List
func draw(serverNameList []string, lineData []string, selectCursor int, searchText string) {
	headLine := 2
	leftMargin := 2
	defaultColor := 255
	defaultBackColor := 255
	pronpt := "lssh>>"
	termbox.Clear(termbox.Attribute(defaultColor+1), termbox.Attribute(defaultBackColor+1))

	// Get Terminal Size
	_, height := termbox.Size()
	lineHeight := height - headLine

	// Set View List Range
	viewFirstLine := (selectCursor/lineHeight)*lineHeight + 1
	viewLastLine := viewFirstLine + lineHeight
	var serverViewList []string
	if viewLastLine > len(serverNameList) {
		serverViewList = serverNameList[viewFirstLine:]
	} else {
		serverViewList = serverNameList[viewFirstLine:viewLastLine]
	}
	selectViewCursor := selectCursor - viewFirstLine + 1

	// View Head
	drawLine(0, 0, pronpt, 3, defaultBackColor)
	drawLine(len(pronpt), 0, searchText, defaultColor, defaultBackColor)
	drawLine(leftMargin, 1, serverNameList[0], 3, defaultBackColor)

	// View List
	for listKey, listValue := range serverViewList {
		paddingListValue := fmt.Sprintf("%-1000s", listValue)
		// Set cursor color
		cursorColor := defaultColor
		cursorBackColor := defaultBackColor
		keywordColor := 5

		for _, selectedLine := range lineData {
			if strings.Split(listValue, " ")[0] == selectedLine {
				cursorColor = 0
				cursorBackColor = 6
			}
		}

		if listKey == selectViewCursor {
			// Select line color
			cursorColor = 0
			cursorBackColor = 2
		}

		// Draw filter line
		drawLine(leftMargin, listKey+headLine, paddingListValue, cursorColor, cursorBackColor)

		// Keyword Highlight
		drawFilterLine(leftMargin, listKey+headLine, paddingListValue, cursorColor, cursorBackColor, keywordColor, searchText)
		listKey += 1
	}

	// Multi-Byte SetCursor
	x := 0
	for _, c := range searchText {
		x += runewidth.RuneWidth(c)
	}
	termbox.SetCursor(len(pronpt)+x, 0)
	termbox.Flush()
}

// Create View List Data (use text/tabwriter)
func getListData(serverNameList []string, serverList conf.Config) (listData []string) {
	buffer := &bytes.Buffer{}
	tabWriterBuffer := new(tabwriter.Writer)
	tabWriterBuffer.Init(buffer, 0, 4, 8, ' ', 0)
	fmt.Fprintln(tabWriterBuffer, "ServerName \tConnect Infomation \tNote \t")

	// Create list table
	for _, key := range serverNameList {
		serverName := key
		connectInfomation := serverList.Server[key].User + "@" + serverList.Server[key].Addr
		serverNote := serverList.Server[key].Note

		fmt.Fprintln(tabWriterBuffer, serverName+"\t"+connectInfomation+"\t"+serverNote)
	}

	tabWriterBuffer.Flush()
	line, err := buffer.ReadString('\n')
	for err == nil {
		str := strings.Replace(line, "\t", " ", -1)
		listData = append(listData, str)
		line, err = buffer.ReadString('\n')
	}
	return listData
}

func insertRune(text string, inputRune rune) (returnText string) {
	returnText = text + string(inputRune)
	return
}

func deleteRune(text string) (returnText string) {
	s := text
	sc := []rune(s)
	returnText = string(sc[:(len(sc) - 1)])
	return
}

func getFilterListData(searchText string, listData []string) (returnListData []string) {
	// SearchText Bounds Space
	searchWords := strings.Fields(searchText)
	r := listData[1:]
	line := ""
	loopListData := []string{}
	returnListData = append(returnListData, listData[0])

	// if No searchWords
	if len(searchWords) == 0 {
		returnListData = listData
		return returnListData
	}

	for i := 0; i < len(searchWords); i += 1 {
		searchWordMeta := regexp.QuoteMeta(strings.ToLower(searchWords[i]))
		re := regexp.MustCompile(searchWordMeta)
		loopListData = []string{}

		for j := 0; j < len(r); j += 1 {
			line += string(r[j])
			if re.MatchString(strings.ToLower(line)) {
				loopListData = append(loopListData, line)
			}
			line = ""
		}
		r = loopListData
	}
	returnListData = append(returnListData, loopListData...)
	return returnListData
}

func pollEvent(serverNameList []string, cmdFlag bool, serverList conf.Config) (lineData []string) {
	defer termbox.Close()
	listData := getListData(serverNameList, serverList)
	selectline := 0
	headLine := 2

	_, height := termbox.Size()
	lineHeight := height - headLine

	searchText := ""
	allFlag := false

	filterListData := getFilterListData(searchText, listData)
	draw(filterListData, lineData, selectline, searchText)
	for {
		switch ev := termbox.PollEvent(); ev.Type {

		// Get Key Event
		case termbox.EventKey:
			switch ev.Key {
			// ESC or Ctrl + C Key (Exit)
			case termbox.KeyEsc, termbox.KeyCtrlC:
				termbox.Close()
				os.Exit(0)

			// AllowUp Key
			case termbox.KeyArrowUp:
				if selectline > 0 {
					selectline -= 1
				}
				draw(filterListData, lineData, selectline, searchText)

			// AllowDown Key
			case termbox.KeyArrowDown:
				if selectline < len(filterListData)-headLine {
					selectline += 1
				}
				draw(filterListData, lineData, selectline, searchText)

			// AllowRight Key
			case termbox.KeyArrowRight:
				if ((selectline+lineHeight)/lineHeight)*lineHeight <= len(filterListData) {
					selectline = ((selectline + lineHeight) / lineHeight) * lineHeight
				}
				draw(filterListData, lineData, selectline, searchText)

			// AllowLeft Key
			case termbox.KeyArrowLeft:
				if ((selectline-lineHeight)/lineHeight)*lineHeight >= 0 {
					selectline = ((selectline - lineHeight) / lineHeight) * lineHeight
				}

				draw(filterListData, lineData, selectline, searchText)

			// Ctrl + x Key(select)
			case termbox.KeyCtrlX:
				if cmdFlag == true {
					lineData = toggleList(lineData, strings.Fields(filterListData[selectline+1])[0])
				}

				draw(filterListData, lineData, selectline, searchText)

			// Ctrl + a Key(all select)
			case termbox.KeyCtrlA:
				if cmdFlag == true {
					lineData = allToggle(allFlag, lineData, filterListData[1:])
				}

				// allFlag Toggle
				if allFlag == false {
					allFlag = true
				} else {
					allFlag = false
				}

				draw(filterListData, lineData, selectline, searchText)

			// Enter Key
			case termbox.KeyEnter:
				if len(lineData) == 0 {
					lineData = append(lineData, strings.Fields(filterListData[selectline+1])[0])
				}
				return

			// BackSpace Key
			case termbox.KeyBackspace, termbox.KeyBackspace2:
				if len(searchText) > 0 {
					searchText = deleteRune(searchText)
					filterListData = getFilterListData(searchText, listData)
					if selectline > len(filterListData) {
						selectline = len(filterListData)
					}
					if selectline < 0 {
						selectline = 0
					}
					allFlag = false
					draw(filterListData, lineData, selectline, searchText)
				}

			// Space Key
			case termbox.KeySpace:
				searchText = searchText + " "
				draw(filterListData, lineData, selectline, searchText)

			// Other Key
			default:
				if ev.Ch != 0 {
					searchText = insertRune(searchText, ev.Ch)
					filterListData = getFilterListData(searchText, listData)
					if selectline > len(filterListData)-headLine {
						selectline = len(filterListData) - headLine
					}
					allFlag = false
					draw(filterListData, lineData, selectline, searchText)
				}
			}
		default:
			draw(filterListData, lineData, selectline, searchText)
		}
	}
}

func DrawList(serverNameList []string, cmdFlag bool, serverList conf.Config) (lineName []string) {
	if err := termbox.Init(); err != nil {
		panic(err)
	}

	lineName = pollEvent(serverNameList, cmdFlag, serverList)
	return lineName
}
<|endoftext|>


<commit_before>package evaluation

import (
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func TestMetrics(t *testing.T) {
	Convey("Quantities derived from a confusion matrix", t, func() {
		confusionMat := make(ConfusionMatrix)
		confusionMat["a"] = make(map[string]int)
		confusionMat["b"] = make(map[string]int)
		confusionMat["a"]["a"] = 75
		confusionMat["a"]["b"] = 5
		confusionMat["b"]["a"] = 10
		confusionMat["b"]["b"] = 10

		Convey("True Positives", func() {
			So(GetTruePositives("a", confusionMat), ShouldAlmostEqual, 75, 1)
			So(GetTruePositives("b", confusionMat), ShouldAlmostEqual, 10, 1)
		})

		Convey("True Negatives", func() {
			So(GetTrueNegatives("a", confusionMat), ShouldAlmostEqual, 10, 1)
			So(GetTrueNegatives("b", confusionMat), ShouldAlmostEqual, 75, 1)
		})

		Convey("False Positives", func() {
			So(GetFalsePositives("a", confusionMat), ShouldAlmostEqual, 10, 1)
			So(GetFalsePositives("b", confusionMat), ShouldAlmostEqual, 5, 1)
		})

		Convey("False Negatives", func() {
			So(GetFalseNegatives("a", confusionMat), ShouldAlmostEqual, 5, 1)
			So(GetFalseNegatives("b", confusionMat), ShouldAlmostEqual, 10, 1)
		})

		Convey("Precision", func() {
			So(GetPrecision("a", confusionMat), ShouldAlmostEqual, 0.88, 0.01)
			So(GetPrecision("b", confusionMat), ShouldAlmostEqual, 0.666, 0.01)
		})

		Convey("Recall", func() {
			So(GetRecall("a", confusionMat), ShouldAlmostEqual, 0.94, 0.01)
			So(GetRecall("b", confusionMat), ShouldAlmostEqual, 0.50, 0.01)
		})

		Convey("MicroPrecision", func() {
			So(GetMicroPrecision(confusionMat), ShouldAlmostEqual, 0.85, 0.01)
		})

		Convey("MicroRecall", func() {
			So(GetMicroRecall(confusionMat), ShouldAlmostEqual, 0.85, 0.01)
		})

		Convey("MacroPrecision", func() {
			So(GetMacroPrecision(confusionMat), ShouldAlmostEqual, 0.775, 0.01)
		})

		Convey("MacroRecall", func() {
			So(GetMacroRecall(confusionMat), ShouldAlmostEqual, 0.719, 0.01)
		})

		Convey("F1Score", func() {
			So(GetF1Score("a", confusionMat), ShouldAlmostEqual, 0.91, 0.1)
			So(GetF1Score("b", confusionMat), ShouldAlmostEqual, 0.571, 0.01)
		})

		Convey("Accuracy", func() {
			So(GetAccuracy(confusionMat), ShouldAlmostEqual, 0.85, 0.1)
		})
	})
}
<commit_msg>Add test cases for C0<commit_after>package evaluation

import (
	"github.com/sjwhitworth/golearn/base"
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func TestMetrics(t *testing.T) {
	Convey("Quantities derived from a confusion matrix", t, func() {
		confusionMat := make(ConfusionMatrix)
		confusionMat["a"] = make(map[string]int)
		confusionMat["b"] = make(map[string]int)
		confusionMat["a"]["a"] = 75
		confusionMat["a"]["b"] = 5
		confusionMat["b"]["a"] = 10
		confusionMat["b"]["b"] = 10

		Convey("True Positives", func() {
			So(GetTruePositives("a", confusionMat), ShouldAlmostEqual, 75, 1)
			So(GetTruePositives("b", confusionMat), ShouldAlmostEqual, 10, 1)
		})

		Convey("True Negatives", func() {
			So(GetTrueNegatives("a", confusionMat), ShouldAlmostEqual, 10, 1)
			So(GetTrueNegatives("b", confusionMat), ShouldAlmostEqual, 75, 1)
		})

		Convey("False Positives", func() {
			So(GetFalsePositives("a", confusionMat), ShouldAlmostEqual, 10, 1)
			So(GetFalsePositives("b", confusionMat), ShouldAlmostEqual, 5, 1)
		})

		Convey("False Negatives", func() {
			So(GetFalseNegatives("a", confusionMat), ShouldAlmostEqual, 5, 1)
			So(GetFalseNegatives("b", confusionMat), ShouldAlmostEqual, 10, 1)
		})

		Convey("Precision", func() {
			So(GetPrecision("a", confusionMat), ShouldAlmostEqual, 0.88, 0.01)
			So(GetPrecision("b", confusionMat), ShouldAlmostEqual, 0.666, 0.01)
		})

		Convey("Recall", func() {
			So(GetRecall("a", confusionMat), ShouldAlmostEqual, 0.94, 0.01)
			So(GetRecall("b", confusionMat), ShouldAlmostEqual, 0.50, 0.01)
		})

		Convey("MicroPrecision", func() {
			So(GetMicroPrecision(confusionMat), ShouldAlmostEqual, 0.85, 0.01)
		})

		Convey("MicroRecall", func() {
			So(GetMicroRecall(confusionMat), ShouldAlmostEqual, 0.85, 0.01)
		})

		Convey("MacroPrecision", func() {
			So(GetMacroPrecision(confusionMat), ShouldAlmostEqual, 0.775, 0.01)
		})

		Convey("MacroRecall", func() {
			So(GetMacroRecall(confusionMat), ShouldAlmostEqual, 0.719, 0.01)
		})

		Convey("F1Score", func() {
			So(GetF1Score("a", confusionMat), ShouldAlmostEqual, 0.91, 0.1)
			So(GetF1Score("b", confusionMat), ShouldAlmostEqual, 0.571, 0.01)
		})

		Convey("Accuracy", func() {
			So(GetAccuracy(confusionMat), ShouldAlmostEqual, 0.85, 0.1)
		})

		Convey("Get Summary", func() {
			output := GetSummary(confusionMat)
			So(output, ShouldStartWith, "Reference Class")
			So(output, ShouldContainSubstring, "True Positives")
			So(output, ShouldContainSubstring, "False Positives")
			So(output, ShouldContainSubstring, "True Negatives")
			So(output, ShouldContainSubstring, "Precision")
			So(output, ShouldContainSubstring, "Recall")
			So(output, ShouldContainSubstring, "F1 Score")
			So(output, ShouldContainSubstring, "------")
			So(output, ShouldContainSubstring, "Overall accuracy:")
		})

		Convey("Show Confusion Matrix", func() {
			output := ShowConfusionMatrix(confusionMat)
			So(output, ShouldStartWith, "Reference Class")
			So(output, ShouldContainSubstring, "---------------")
		})

		Convey("Get Confusion Matrix", func() {
			X, _ := base.ParseCSVToInstances("../examples/datasets/iris_headers.csv", true)
			Y, _ := base.ParseCSVToInstances("../examples/datasets/exam.csv", true)
			Convey("Nomarl ref and gen matrices", func() {
				C, _ := GetConfusionMatrix(X, X)
				ret := make(map[string]map[string]int)
				So(C, ShouldHaveSameTypeAs, ret)
			})
			Convey("Row count mismatch", func() {
				_, err := GetConfusionMatrix(X, Y)
				So(err.Error(), ShouldStartWith, "Row count mismatch:")
			})
		})
	})
}
<|endoftext|>


<commit_before>package assets

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"gnd.la/log"
	"hash/fnv"
	"io"
	"os"
	"path"
	"regexp"
	"strings"
)

const (
	// If you define your own code types, use numbers > 1000
	CodeTypeCss        = 1
	CodeTypeJavascript = 2
)

var (
	ErrNoAssets = errors.New("no assets to compile")
	urlRe       = regexp.MustCompile("i?url\\s*?\\((.*?)\\)")
)

type CodeAsset interface {
	Asset
	CodeType() int
	Code() (string, error)
}

type CodeAssetList []CodeAsset

func (c CodeAssetList) Names() []string {
	var names []string
	for _, v := range c {
		names = append(names, v.AssetName())
	}
	return names
}

func (c CodeAssetList) CompiledName(ext string, o Options) (string, error) {
	if len(c) == 0 {
		return "", nil
	}
	h := fnv.New32a()
	for _, v := range c {
		code, err := v.Code()
		if err != nil {
			return "", err
		}
		io.WriteString(h, code)
	}
	io.WriteString(h, o.String())
	sum := hex.EncodeToString(h.Sum(nil))
	name := c[0].AssetName()
	if ext == "" {
		ext = path.Ext(name)
	} else {
		ext = "." + ext
	}
	return path.Join(path.Dir(name), "asset-"+sum+ext), nil
}

func Compile(m Manager, assets []Asset, opts Options) ([]Asset, error) {
	if len(assets) == 0 {
		return nil, ErrNoAssets
	}
	var ctype int
	codeAssets := make(CodeAssetList, len(assets))
	for ii, v := range assets {
		c, ok := v.(CodeAsset)
		if !ok {
			return nil, fmt.Errorf("asset %q (type %T) does not implement CodeAsset and can't be compiled", v.AssetName(), v)
		}
		if ctype == 0 {
			ctype = c.CodeType()
		} else if ctype != c.CodeType() {
			return nil, fmt.Errorf("asset %q has different code type %d (first asset is of type %d)", v.AssetName(), c.CodeType(), ctype)
		}
		codeAssets[ii] = c
	}
	compiler := compilers[ctype]
	if compiler == nil {
		return nil, fmt.Errorf("no compiler for code type %d", ctype)
	}
	// Prepare the code, changing relative paths if required
	name, err := codeAssets.CompiledName(compiler.Ext(), opts)
	if err != nil {
		return nil, err
	}
	dir := path.Dir(name)
	var code []string
	for _, v := range codeAssets {
		c, err := v.Code()
		if err != nil {
			return nil, fmt.Errorf("error getting code for asset %q: %s", v.AssetName(), err)
		}
		if vd := path.Dir(v.AssetName()); vd != dir {
			if ctype == CodeTypeCss {
				log.Debugf("asset %q will move from %v to %v, rewriting relative paths...", v.AssetName(), vd, dir)
				c = replaceRelativePaths(c, vd, dir)
			} else {
				log.Warningf("asset %q will move from %v to %v, relative paths might not work", v.AssetName(), vd, dir)
			}
		}
		code = append(code, c)
	}
	// Check if the code has been already compiled
	if _, _, err := m.Load(name); err == nil {
		log.Debugf("%s already compiled into %s and up to date", codeAssets.Names(), name)
	} else {
		log.Debugf("Compiling %v", codeAssets.Names())
		// Compile to a buf first. We don't want to create
		// the file if the compilation fails
		var buf bytes.Buffer
		reader := strings.NewReader(strings.Join(code, "\n\n"))
		if err := compiler.Compile(reader, &buf, m, opts); err != nil {
			return nil, err
		}
		w, err := m.Create(name)
		if err == nil {
			s := makeLinksCacheable(m, dir, buf.Bytes())
			if _, err := io.Copy(w, strings.NewReader(s)); err != nil {
				w.Close()
				return nil, err
			}
			if err := w.Close(); err != nil {
				return nil, err
			}
		} else {
			// If the file exists, is up to date
			if !os.IsExist(err) {
				return nil, err
			}
		}
	}
	asset, err := compiler.Asset(name, m, opts)
	if err != nil {
		return nil, err
	}
	return []Asset{asset}, nil
}

func makeLinksCacheable(m Manager, dir string, b []byte) string {
	css := string(b)
	return replaceCssUrls(css, func(s string) string {
		var suffix string
		if sep := strings.IndexAny(s, "?#"); sep >= 0 {
			suffix = s[sep:]
			s = s[:sep]
		}
		p := path.Join(dir, s)
		base := m.URL(p)
		if strings.Contains(base, "?") && suffix != "" && suffix[0] == '?' {
			suffix = "&" + suffix[1:]
		}
		repl := base + suffix
		return repl
	})
}

func replaceRelativePaths(code string, dir string, final string) string {
	count := strings.Count(final, "/") + 1
	return replaceCssUrls(code, func(s string) string {
		old := path.Join(dir, s)
		return strings.Repeat("../", count) + old
	})
}

func replaceCssUrls(code string, f func(string) string) string {
	return urlRe.ReplaceAllStringFunc(code, func(s string) string {
		r := urlRe.FindStringSubmatch(s)
		p := r[1]
		quote := ""
		if len(p) > 0 && (p[0] == '\'' || p[0] == '"') {
			quote = string(p[0])
			p = p[1 : len(p)-1]
		}
		if !urlIsRelative(p) {
			return s
		}
		repl := f(p)
		if repl == p {
			return s
		}
		return fmt.Sprintf("url(%s%s%s)", quote, repl, quote)
	})
}

func urlIsRelative(u string) bool {
	return !strings.HasPrefix(u, "//") && !strings.HasPrefix(u, "http://") &&
		!strings.HasPrefix(u, "https://") && !strings.HasPrefix(u, "data:")
}
<commit_msg>Name bundles bundle.<hash>.<ext><commit_after>package assets

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"gnd.la/log"
	"hash/fnv"
	"io"
	"os"
	"path"
	"regexp"
	"strings"
)

const (
	// If you define your own code types, use numbers > 1000
	CodeTypeCss        = 1
	CodeTypeJavascript = 2
)

var (
	ErrNoAssets = errors.New("no assets to compile")
	urlRe       = regexp.MustCompile("i?url\\s*?\\((.*?)\\)")
)

type CodeAsset interface {
	Asset
	CodeType() int
	Code() (string, error)
}

type CodeAssetList []CodeAsset

func (c CodeAssetList) Names() []string {
	var names []string
	for _, v := range c {
		names = append(names, v.AssetName())
	}
	return names
}

func (c CodeAssetList) CompiledName(ext string, o Options) (string, error) {
	if len(c) == 0 {
		return "", nil
	}
	h := fnv.New32a()
	for _, v := range c {
		code, err := v.Code()
		if err != nil {
			return "", err
		}
		io.WriteString(h, code)
	}
	io.WriteString(h, o.String())
	sum := hex.EncodeToString(h.Sum(nil))
	name := c[0].AssetName()
	if ext == "" {
		ext = path.Ext(name)
	} else {
		ext = "." + ext
	}
	return path.Join(path.Dir(name), "bundle."+sum+ext), nil
}

func Compile(m Manager, assets []Asset, opts Options) ([]Asset, error) {
	if len(assets) == 0 {
		return nil, ErrNoAssets
	}
	var ctype int
	codeAssets := make(CodeAssetList, len(assets))
	for ii, v := range assets {
		c, ok := v.(CodeAsset)
		if !ok {
			return nil, fmt.Errorf("asset %q (type %T) does not implement CodeAsset and can't be compiled", v.AssetName(), v)
		}
		if ctype == 0 {
			ctype = c.CodeType()
		} else if ctype != c.CodeType() {
			return nil, fmt.Errorf("asset %q has different code type %d (first asset is of type %d)", v.AssetName(), c.CodeType(), ctype)
		}
		codeAssets[ii] = c
	}
	compiler := compilers[ctype]
	if compiler == nil {
		return nil, fmt.Errorf("no compiler for code type %d", ctype)
	}
	// Prepare the code, changing relative paths if required
	name, err := codeAssets.CompiledName(compiler.Ext(), opts)
	if err != nil {
		return nil, err
	}
	dir := path.Dir(name)
	var code []string
	for _, v := range codeAssets {
		c, err := v.Code()
		if err != nil {
			return nil, fmt.Errorf("error getting code for asset %q: %s", v.AssetName(), err)
		}
		if vd := path.Dir(v.AssetName()); vd != dir {
			if ctype == CodeTypeCss {
				log.Debugf("asset %q will move from %v to %v, rewriting relative paths...", v.AssetName(), vd, dir)
				c = replaceRelativePaths(c, vd, dir)
			} else {
				log.Warningf("asset %q will move from %v to %v, relative paths might not work", v.AssetName(), vd, dir)
			}
		}
		code = append(code, c)
	}
	// Check if the code has been already compiled
	if _, _, err := m.Load(name); err == nil {
		log.Debugf("%s already compiled into %s and up to date", codeAssets.Names(), name)
	} else {
		log.Debugf("Compiling %v", codeAssets.Names())
		// Compile to a buf first. We don't want to create
		// the file if the compilation fails
		var buf bytes.Buffer
		reader := strings.NewReader(strings.Join(code, "\n\n"))
		if err := compiler.Compile(reader, &buf, m, opts); err != nil {
			return nil, err
		}
		w, err := m.Create(name)
		if err == nil {
			s := makeLinksCacheable(m, dir, buf.Bytes())
			if _, err := io.Copy(w, strings.NewReader(s)); err != nil {
				w.Close()
				return nil, err
			}
			if err := w.Close(); err != nil {
				return nil, err
			}
		} else {
			// If the file exists, is up to date
			if !os.IsExist(err) {
				return nil, err
			}
		}
	}
	asset, err := compiler.Asset(name, m, opts)
	if err != nil {
		return nil, err
	}
	return []Asset{asset}, nil
}

func makeLinksCacheable(m Manager, dir string, b []byte) string {
	css := string(b)
	return replaceCssUrls(css, func(s string) string {
		var suffix string
		if sep := strings.IndexAny(s, "?#"); sep >= 0 {
			suffix = s[sep:]
			s = s[:sep]
		}
		p := path.Join(dir, s)
		base := m.URL(p)
		if strings.Contains(base, "?") && suffix != "" && suffix[0] == '?' {
			suffix = "&" + suffix[1:]
		}
		repl := base + suffix
		return repl
	})
}

func replaceRelativePaths(code string, dir string, final string) string {
	count := strings.Count(final, "/") + 1
	return replaceCssUrls(code, func(s string) string {
		old := path.Join(dir, s)
		return strings.Repeat("../", count) + old
	})
}

func replaceCssUrls(code string, f func(string) string) string {
	return urlRe.ReplaceAllStringFunc(code, func(s string) string {
		r := urlRe.FindStringSubmatch(s)
		p := r[1]
		quote := ""
		if len(p) > 0 && (p[0] == '\'' || p[0] == '"') {
			quote = string(p[0])
			p = p[1 : len(p)-1]
		}
		if !urlIsRelative(p) {
			return s
		}
		repl := f(p)
		if repl == p {
			return s
		}
		return fmt.Sprintf("url(%s%s%s)", quote, repl, quote)
	})
}

func urlIsRelative(u string) bool {
	return !strings.HasPrefix(u, "//") && !strings.HasPrefix(u, "http://") &&
		!strings.HasPrefix(u, "https://") && !strings.HasPrefix(u, "data:")
}
<|endoftext|>