<commit_before>package gotoml

import (
	"bytes"
	"io/ioutil"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/ota42y/gotoml/example"
)

func isValidToml(t *testing.T, b bool) {
	if !b {
		t.Error("invalid toml")
	}
}

func TestExampleData(t *testing.T) {
	// example/normal.go is valid toml struct?
	var n example.Normal

	file, err := os.Open("example/normal.toml")
	if err != nil {
		t.Error(err)
	}
	// read toml file
	if _, err := toml.DecodeReader(file, &n); err != nil {
		t.Error(err)
	}

	isValidToml(t, n.Name == "Dream Sensation")

	d, err := time.Parse("2006-01-02T15:04:05Z", "2015-01-31T18:08:03Z")
	if err != nil {
		t.Error(err)
	}
	isValidToml(t, n.Date == d)

	dates := make([]time.Time, 2)
	dates[0], err = time.Parse("2006-01-02T15:04:05Z", "2015-01-31T18:08:03Z")
	if err != nil {
		t.Error(err)
	}
	dates[1], err = time.Parse("2006-01-02T15:04:05Z", "2015-02-01T16:08:03Z")
	if err != nil {
		t.Error(err)
	}

	isValidToml(t, n.Dates[0] == dates[0])
	isValidToml(t, n.Dates[1] == dates[1])

	isValidToml(t, n.Num == 64)

	isValidToml(t, n.Numbers[0] == 1)
	isValidToml(t, n.Numbers[1] == 2)
	isValidToml(t, n.Numbers[2] == 3)

	isValidToml(t, n.Group.Name == "μ's")
	isValidToml(t, n.Group.School == "National Society of Otonoki")

	isValidToml(t, n.Site.Name == "Saitama Super Arena")
	isValidToml(t, n.Site.Num == 36500)
}

func TestGenerateHead(t *testing.T) {
	pkgName := "main"

	w := new(bytes.Buffer)
	g := newGenerator()
	g.generateHead(w, pkgName)

	expect := `package main

`
	actual := string(w.Bytes())

	if expect != actual {
		t.Errorf("generateHead expect %s but %s", expect, actual)
	}

	g = newGenerator()
	g.usingPackages[""] = true
	g.usingPackages["time"] = true
	w = new(bytes.Buffer)
	g.generateHead(w, pkgName)

	expect = `package main
import (
"time"
)

`
	actual = string(w.Bytes())

	if expect != actual {
		t.Errorf("generateHead expect %s but %s", expect, actual)
	}

}

func TestGenerate(t *testing.T) {
	structName := "Normal"
	pkgName := "example"

	file, err := os.Open("example/normal.toml")
	if err != nil {
		t.Error(err)
	}

	b, err := ioutil.ReadFile("example/normal.go")
	if err != nil {
		t.Error(err)
	}
	expect := string(b)

	generateByte, err := Generate(file, structName, pkgName)
	if err != nil {
		t.Error(err)
	}

	actual := string(generateByte)
	if expect != actual {
		t.Errorf("Generate expect %s but %s", expect, actual)
	}
}

func checkGetTypeName(input string, expectName string, expectPackage string, t *testing.T) *generator {
	g := newGenerator()

	var data map[string]interface{}
	if _, err := toml.DecodeReader(strings.NewReader(input), &data); err != nil {
		t.Error(err)
		return nil
	}

	key := ""
	for k := range data {
		key = k
	}

	actual, pkgName := g.getTypeName(key, data[key])
	if expectName != actual {
		t.Errorf("getTypeName should return name %s but `%s`", expectName, actual)
	}
	if pkgName != expectPackage {
		t.Errorf("getTypeName should return package name %s but `%s`", expectPackage, pkgName)
	}

	return g
}

func TestGetTypeName(t *testing.T) {
	input := `name = "Name"`
	checkGetTypeName(input, "string", "", t)

	input = `date = 1979-05-27T07:32:00Z`
	checkGetTypeName(input, "time.Time", "time", t)

	input = `num = 12`
	checkGetTypeName(input, "int64", "", t)

	input = `interface = []`
	checkGetTypeName(input, "[]interface{}", "", t)

	input = `numbers = [1, 2, 3]`
	checkGetTypeName(input, "[]int64", "", t)

	input = `dates = [2015-01-31T18:08:03Z, 2015-02-01T16:08:03Z]`
	checkGetTypeName(input, "[]time.Time", "time", t)

	input = `[group]
name = "A-RISE"
school = "Society of UTX"
`
	g := checkGetTypeName(input, "Group", "", t)
	_, ok := g.tomlData["Group"]
	if !ok {
		t.Errorf("getTypeName should save sub struct but not saved %v", g.tomlData)
	}

	v, ok := g.tomlParsed["Group"]
	if !ok || v {
		t.Errorf("getTypeName should save sub struct but not saved %v", g.tomlParsed)
	}
}
<commit_msg>add error test<commit_after>package gotoml

import (
	"bytes"
	"io/ioutil"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/ota42y/gotoml/example"
)

func isValidToml(t *testing.T, b bool) {
	if !b {
		t.Error("invalid toml")
	}
}

func TestExampleData(t *testing.T) {
	// example/normal.go is valid toml struct?
	var n example.Normal

	file, err := os.Open("example/normal.toml")
	if err != nil {
		t.Error(err)
	}
	// read toml file
	if _, err := toml.DecodeReader(file, &n); err != nil {
		t.Error(err)
	}

	isValidToml(t, n.Name == "Dream Sensation")

	d, err := time.Parse("2006-01-02T15:04:05Z", "2015-01-31T18:08:03Z")
	if err != nil {
		t.Error(err)
	}
	isValidToml(t, n.Date == d)

	dates := make([]time.Time, 2)
	dates[0], err = time.Parse("2006-01-02T15:04:05Z", "2015-01-31T18:08:03Z")
	if err != nil {
		t.Error(err)
	}
	dates[1], err = time.Parse("2006-01-02T15:04:05Z", "2015-02-01T16:08:03Z")
	if err != nil {
		t.Error(err)
	}

	isValidToml(t, n.Dates[0] == dates[0])
	isValidToml(t, n.Dates[1] == dates[1])

	isValidToml(t, n.Num == 64)

	isValidToml(t, n.Numbers[0] == 1)
	isValidToml(t, n.Numbers[1] == 2)
	isValidToml(t, n.Numbers[2] == 3)

	isValidToml(t, n.Group.Name == "μ's")
	isValidToml(t, n.Group.School == "National Society of Otonoki")

	isValidToml(t, n.Site.Name == "Saitama Super Arena")
	isValidToml(t, n.Site.Num == 36500)
}

func TestGenerateHead(t *testing.T) {
	pkgName := "main"

	w := new(bytes.Buffer)
	g := newGenerator()
	g.generateHead(w, pkgName)

	expect := `package main

`
	actual := string(w.Bytes())

	if expect != actual {
		t.Errorf("generateHead expect %s but %s", expect, actual)
	}

	g = newGenerator()
	g.usingPackages[""] = true
	g.usingPackages["time"] = true
	w = new(bytes.Buffer)
	g.generateHead(w, pkgName)

	expect = `package main
import (
"time"
)

`
	actual = string(w.Bytes())

	if expect != actual {
		t.Errorf("generateHead expect %s but %s", expect, actual)
	}

}

func TestGenerateError(t *testing.T) {
	invalidData := "invalid"
	input := strings.NewReader(invalidData)

	b, err := Generate(input, "", "")
	if err == nil {
		t.Error("if invalid toml data, should return error, but nil")
	}

	if b != nil {
		t.Errorf("if invalid toml data, should return nil, but %v", b)
	}
}

func TestGenerate(t *testing.T) {
	structName := "Normal"
	pkgName := "example"

	file, err := os.Open("example/normal.toml")
	if err != nil {
		t.Error(err)
	}

	b, err := ioutil.ReadFile("example/normal.go")
	if err != nil {
		t.Error(err)
	}
	expect := string(b)

	generateByte, err := Generate(file, structName, pkgName)
	if err != nil {
		t.Error(err)
	}

	actual := string(generateByte)
	if expect != actual {
		t.Errorf("Generate expect %s but %s", expect, actual)
	}
}

func checkGetTypeName(input string, expectName string, expectPackage string, t *testing.T) *generator {
	g := newGenerator()

	var data map[string]interface{}
	if _, err := toml.DecodeReader(strings.NewReader(input), &data); err != nil {
		t.Error(err)
		return nil
	}

	key := ""
	for k := range data {
		key = k
	}

	actual, pkgName := g.getTypeName(key, data[key])
	if expectName != actual {
		t.Errorf("getTypeName should return name %s but `%s`", expectName, actual)
	}
	if pkgName != expectPackage {
		t.Errorf("getTypeName should return package name %s but `%s`", expectPackage, pkgName)
	}

	return g
}

func TestGetTypeName(t *testing.T) {
	input := `name = "Name"`
	checkGetTypeName(input, "string", "", t)

	input = `date = 1979-05-27T07:32:00Z`
	checkGetTypeName(input, "time.Time", "time", t)

	input = `num = 12`
	checkGetTypeName(input, "int64", "", t)

	input = `interface = []`
	checkGetTypeName(input, "[]interface{}", "", t)

	input = `numbers = [1, 2, 3]`
	checkGetTypeName(input, "[]int64", "", t)

	input = `dates = [2015-01-31T18:08:03Z, 2015-02-01T16:08:03Z]`
	checkGetTypeName(input, "[]time.Time", "time", t)

	input = `[group]
name = "A-RISE"
school = "Society of UTX"
`
	g := checkGetTypeName(input, "Group", "", t)
	_, ok := g.tomlData["Group"]
	if !ok {
		t.Errorf("getTypeName should save sub struct but not saved %v", g.tomlData)
	}

	v, ok := g.tomlParsed["Group"]
	if !ok || v {
		t.Errorf("getTypeName should save sub struct but not saved %v", g.tomlParsed)
	}
}
<|endoftext|>


<commit_before>package workerpool

import (
	"sync"
	"testing"
	"time"
)

const max = 20

func TestExample(t *testing.T) {
	wp := New(2)
	requests := []string{"alpha", "beta", "gamma", "delta", "epsilon"}

	rspChan := make(chan string, len(requests))
	for _, r := range requests {
		r := r
		wp.Submit(func() {
			rspChan <- r
		})
	}

	wp.StopWait()

	close(rspChan)
	rspSet := map[string]struct{}{}
	for rsp := range rspChan {
		rspSet[rsp] = struct{}{}
	}
	if len(rspSet) < len(requests) {
		t.Fatal("Did not handle all requests")
	}
	for _, req := range requests {
		if _, ok := rspSet[req]; !ok {
			t.Fatal("Missing expected values:", req)
		}
	}
}

func TestMaxWorkers(t *testing.T) {
	t.Parallel()

	wp := New(0)
	if wp.maxWorkers != 1 {
		t.Fatal("should have created one worker")
	}

	wp = New(max)
	defer wp.Stop()

	started := make(chan struct{}, max)
	sync := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for all queued tasks to be dispatched to workers.
	timeout := time.After(5 * time.Second)
	if wp.waitingQueue.Len() != wp.WaitingQueueSize() {
		t.Fatal("Working Queue size returned should not be 0")
		panic("WRONG")
	}
	for startCount := 0; startCount < max; {
		select {
		case <-started:
			startCount++
		case <-timeout:
			t.Fatal("timed out waiting for workers to start")
		}
	}

	// Release workers.
	close(sync)
}

func TestReuseWorkers(t *testing.T) {
	t.Parallel()

	wp := New(5)
	defer wp.Stop()

	sync := make(chan struct{})

	// Cause worker to be created, and available for reuse before next task.
	for i := 0; i < 10; i++ {
		wp.Submit(func() { <-sync })
		sync <- struct{}{}
		time.Sleep(100 * time.Millisecond)
	}

	// If the same worker was always reused, then only one worker would have
	// been created and there should only be one ready.
	if countReady(wp) > 1 {
		t.Fatal("Worker not reused")
	}
}

func TestWorkerTimeout(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	sync := make(chan struct{})
	started := make(chan struct{}, max)
	// Cause workers to be created.  Workers wait on channel, keeping them busy
	// and causing the worker pool to create more.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for tasks to start.
	for i := 0; i < max; i++ {
		<-started
	}

	if anyReady(wp) {
		t.Fatal("number of ready workers should be zero")
	}
	// Release workers.
	close(sync)

	if countReady(wp) != max {
		t.Fatal("Expected", max, "ready workers")
	}

	// Check that a worker timed out.
	time.Sleep((idleTimeoutSec + 1) * time.Second)
	if countReady(wp) != max-1 {
		t.Fatal("First worker did not timeout")
	}

	// Check that another worker timed out.
	time.Sleep((idleTimeoutSec + 1) * time.Second)
	if countReady(wp) != max-2 {
		t.Fatal("Second worker did not timeout")
	}
}

func TestStop(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	started := make(chan struct{}, max)
	sync := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for all queued tasks to be dispatched to workers.
	timeout := time.After(5 * time.Second)
	for startCount := 0; startCount < max; {
		select {
		case <-started:
			startCount++
		case <-timeout:
			t.Fatal("timed out waiting for workers to start")
		}
	}

	// Release workers.
	close(sync)

	if wp.Stopped() {
		t.Error("pool should not be stopped")
	}

	wp.Stop()
	if anyReady(wp) {
		t.Error("should have zero workers after stop")
	}

	if !wp.Stopped() {
		t.Error("pool should be stopped")
	}

	// Start workers, and have them all wait on a channel before completing.
	wp = New(5)
	sync = make(chan struct{})
	finished := make(chan struct{}, max)
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			<-sync
			finished <- struct{}{}
		})
	}

	// Call Stop() and see that only the already running tasks were completed.
	go func() {
		time.Sleep(10000 * time.Millisecond)
		close(sync)
	}()
	wp.Stop()
	var count int
Count:
	for count < max {
		select {
		case <-finished:
			count++
		default:
			break Count
		}
	}
	if count > 5 {
		t.Error("Should not have completed any queued tasks, did", count)
	}

	// Check that calling Stop() againg is OK.
	wp.Stop()
}

func TestStopWait(t *testing.T) {
	t.Parallel()

	// Start workers, and have them all wait on a channel before completing.
	wp := New(5)
	sync := make(chan struct{})
	finished := make(chan struct{}, max)
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			<-sync
			finished <- struct{}{}
		})
	}

	// Call StopWait() and see that all tasks were completed.
	go func() {
		time.Sleep(10 * time.Millisecond)
		close(sync)
	}()
	wp.StopWait()
	for count := 0; count < max; count++ {
		select {
		case <-finished:
		default:
			t.Error("Should have completed all queued tasks")
		}
	}

	if anyReady(wp) {
		t.Error("should have zero workers after stopwait")
	}

	if !wp.Stopped() {
		t.Error("pool should be stopped")
	}

	// Make sure that calling StopWait() with no queued tasks is OK.
	wp = New(5)
	wp.StopWait()

	if anyReady(wp) {
		t.Error("should have zero workers after stopwait")
	}

	// Check that calling StopWait() againg is OK.
	wp.StopWait()
}

func TestSubmitWait(t *testing.T) {
	wp := New(1)
	defer wp.Stop()

	// Check that these are noop.
	wp.Submit(nil)
	wp.SubmitWait(nil)

	done1 := make(chan struct{})
	wp.Submit(func() {
		time.Sleep(100 * time.Millisecond)
		close(done1)
	})
	select {
	case <-done1:
		t.Fatal("Submit did not return immediately")
	default:
	}

	done2 := make(chan struct{})
	wp.SubmitWait(func() {
		time.Sleep(100 * time.Millisecond)
		close(done2)
	})
	select {
	case <-done2:
	default:
		t.Fatal("SubmitWait did not wait for function to execute")
	}
}

func TestOverflow(t *testing.T) {
	wp := New(2)
	releaseChan := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < 64; i++ {
		wp.Submit(func() { <-releaseChan })
	}

	// Start a goroutine to free the workers after calling stop.  This way
	// the dispatcher can exit, then when this goroutine runs, the workerpool
	// can exit.
	go func() {
		<-time.After(time.Millisecond)
		close(releaseChan)
	}()
	wp.Stop()

	// Now that the worker pool has exited, it is safe to inspect its waiting
	// queue without causing a race.
	qlen := wp.waitingQueue.Len()
	if qlen != 62 {
		t.Fatal("Expected 62 tasks in waiting queue, have", qlen)
	}
}

func TestStopRace(t *testing.T) {
	wp := New(20)
	releaseChan := make(chan struct{})
	workRelChan := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < 20; i++ {
		wp.Submit(func() { <-workRelChan })
	}

	time.Sleep(5 * time.Second)
	for i := 0; i < 64; i++ {
		go func() {
			<-releaseChan
			wp.Stop()
		}()
	}

	close(workRelChan)
	close(releaseChan)
}

// Run this test with race detector to test that using WaitingQueueSize has no
// race condition
func TestWaitingQueueSizeRace(t *testing.T) {
	const (
		goroutines = 10
		tasks      = 20
		workers    = 5
	)
	wp := New(workers)
	maxChan := make(chan int)
	for g := 0; g < goroutines; g++ {
		go func() {
			max := 0
			// Submit 100 tasks, checking waiting queue size each time.  Report
			// the maximum queue size seen.
			for i := 0; i < tasks; i++ {
				wp.Submit(func() {
					time.Sleep(time.Microsecond)
				})
				waiting := wp.WaitingQueueSize()
				if waiting > max {
					max = waiting
				}
			}
			maxChan <- max
		}()
	}

	// Find maximum queuesize seen by any thread.
	maxMax := 0
	for g := 0; g < goroutines; g++ {
		max := <-maxChan
		if max > maxMax {
			maxMax = max
		}
	}
	if maxMax == 0 {
		t.Error("expected to see waiting queue size > 0")
	}
	if maxMax >= goroutines*tasks {
		t.Error("should not have seen all tasks on waiting queue")
	}
}

func anyReady(w *WorkerPool) bool {
	select {
	case wkCh := <-w.readyWorkers:
		w.readyWorkers <- wkCh
		return true
	default:
	}
	return false
}

func countReady(w *WorkerPool) int {
	// Try to pull max workers off of ready queue.
	timeout := time.After(5 * time.Second)
	readyTmp := make(chan chan func(), max)
	var readyCount int
	for i := 0; i < max; i++ {
		select {
		case wkCh := <-w.readyWorkers:
			readyTmp <- wkCh
			readyCount++
		case <-timeout:
			readyCount = i
			i = max
		}
	}

	// Restore ready workers.
	close(readyTmp)
	go func() {
		for r := range readyTmp {
			w.readyWorkers <- r
		}
	}()
	return readyCount
}

/*

Run benchmarking with: go test -bench '.'

*/

func BenchmarkEnqueue(b *testing.B) {
	wp := New(1)
	defer wp.Stop()
	releaseChan := make(chan struct{})

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() { <-releaseChan })
	}
	close(releaseChan)
}

func BenchmarkEnqueue2(b *testing.B) {
	wp := New(2)
	defer wp.Stop()
	releaseChan := make(chan struct{})

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		for i := 0; i < 64; i++ {
			wp.Submit(func() { <-releaseChan })
		}
		for i := 0; i < 64; i++ {
			releaseChan <- struct{}{}
		}
	}
	close(releaseChan)
}

func BenchmarkExecute1Worker(b *testing.B) {
	wp := New(1)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute2Worker(b *testing.B) {
	wp := New(2)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute4Workers(b *testing.B) {
	wp := New(4)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute16Workers(b *testing.B) {
	wp := New(16)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute64Workers(b *testing.B) {
	wp := New(64)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}
<commit_msg>Change t.Error to t.Fatal in test file<commit_after>package workerpool

import (
	"sync"
	"testing"
	"time"
)

const max = 20

func TestExample(t *testing.T) {
	wp := New(2)
	requests := []string{"alpha", "beta", "gamma", "delta", "epsilon"}

	rspChan := make(chan string, len(requests))
	for _, r := range requests {
		r := r
		wp.Submit(func() {
			rspChan <- r
		})
	}

	wp.StopWait()

	close(rspChan)
	rspSet := map[string]struct{}{}
	for rsp := range rspChan {
		rspSet[rsp] = struct{}{}
	}
	if len(rspSet) < len(requests) {
		t.Fatal("Did not handle all requests")
	}
	for _, req := range requests {
		if _, ok := rspSet[req]; !ok {
			t.Fatal("Missing expected values:", req)
		}
	}
}

func TestMaxWorkers(t *testing.T) {
	t.Parallel()

	wp := New(0)
	if wp.maxWorkers != 1 {
		t.Fatal("should have created one worker")
	}

	wp = New(max)
	defer wp.Stop()

	started := make(chan struct{}, max)
	sync := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for all queued tasks to be dispatched to workers.
	timeout := time.After(5 * time.Second)
	if wp.waitingQueue.Len() != wp.WaitingQueueSize() {
		t.Fatal("Working Queue size returned should not be 0")
		panic("WRONG")
	}
	for startCount := 0; startCount < max; {
		select {
		case <-started:
			startCount++
		case <-timeout:
			t.Fatal("timed out waiting for workers to start")
		}
	}

	// Release workers.
	close(sync)
}

func TestReuseWorkers(t *testing.T) {
	t.Parallel()

	wp := New(5)
	defer wp.Stop()

	sync := make(chan struct{})

	// Cause worker to be created, and available for reuse before next task.
	for i := 0; i < 10; i++ {
		wp.Submit(func() { <-sync })
		sync <- struct{}{}
		time.Sleep(100 * time.Millisecond)
	}

	// If the same worker was always reused, then only one worker would have
	// been created and there should only be one ready.
	if countReady(wp) > 1 {
		t.Fatal("Worker not reused")
	}
}

func TestWorkerTimeout(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	sync := make(chan struct{})
	started := make(chan struct{}, max)
	// Cause workers to be created.  Workers wait on channel, keeping them busy
	// and causing the worker pool to create more.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for tasks to start.
	for i := 0; i < max; i++ {
		<-started
	}

	if anyReady(wp) {
		t.Fatal("number of ready workers should be zero")
	}
	// Release workers.
	close(sync)

	if countReady(wp) != max {
		t.Fatal("Expected", max, "ready workers")
	}

	// Check that a worker timed out.
	time.Sleep((idleTimeoutSec + 1) * time.Second)
	if countReady(wp) != max-1 {
		t.Fatal("First worker did not timeout")
	}

	// Check that another worker timed out.
	time.Sleep((idleTimeoutSec + 1) * time.Second)
	if countReady(wp) != max-2 {
		t.Fatal("Second worker did not timeout")
	}
}

func TestStop(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	started := make(chan struct{}, max)
	sync := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for all queued tasks to be dispatched to workers.
	timeout := time.After(5 * time.Second)
	for startCount := 0; startCount < max; {
		select {
		case <-started:
			startCount++
		case <-timeout:
			t.Fatal("timed out waiting for workers to start")
		}
	}

	// Release workers.
	close(sync)

	if wp.Stopped() {
		t.Fatal("pool should not be stopped")
	}

	wp.Stop()
	if anyReady(wp) {
		t.Fatal("should have zero workers after stop")
	}

	if !wp.Stopped() {
		t.Fatal("pool should be stopped")
	}

	// Start workers, and have them all wait on a channel before completing.
	wp = New(5)
	sync = make(chan struct{})
	finished := make(chan struct{}, max)
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			<-sync
			finished <- struct{}{}
		})
	}

	// Call Stop() and see that only the already running tasks were completed.
	go func() {
		time.Sleep(10000 * time.Millisecond)
		close(sync)
	}()
	wp.Stop()
	var count int
Count:
	for count < max {
		select {
		case <-finished:
			count++
		default:
			break Count
		}
	}
	if count > 5 {
		t.Fatal("Should not have completed any queued tasks, did", count)
	}

	// Check that calling Stop() again is OK.
	wp.Stop()
}

func TestStopWait(t *testing.T) {
	t.Parallel()

	// Start workers, and have them all wait on a channel before completing.
	wp := New(5)
	sync := make(chan struct{})
	finished := make(chan struct{}, max)
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			<-sync
			finished <- struct{}{}
		})
	}

	// Call StopWait() and see that all tasks were completed.
	go func() {
		time.Sleep(10 * time.Millisecond)
		close(sync)
	}()
	wp.StopWait()
	for count := 0; count < max; count++ {
		select {
		case <-finished:
		default:
			t.Fatal("Should have completed all queued tasks")
		}
	}

	if anyReady(wp) {
		t.Fatal("should have zero workers after stopwait")
	}

	if !wp.Stopped() {
		t.Fatal("pool should be stopped")
	}

	// Make sure that calling StopWait() with no queued tasks is OK.
	wp = New(5)
	wp.StopWait()

	if anyReady(wp) {
		t.Fatal("should have zero workers after stopwait")
	}

	// Check that calling StopWait() again is OK.
	wp.StopWait()
}

func TestSubmitWait(t *testing.T) {
	wp := New(1)
	defer wp.Stop()

	// Check that these are noop.
	wp.Submit(nil)
	wp.SubmitWait(nil)

	done1 := make(chan struct{})
	wp.Submit(func() {
		time.Sleep(100 * time.Millisecond)
		close(done1)
	})
	select {
	case <-done1:
		t.Fatal("Submit did not return immediately")
	default:
	}

	done2 := make(chan struct{})
	wp.SubmitWait(func() {
		time.Sleep(100 * time.Millisecond)
		close(done2)
	})
	select {
	case <-done2:
	default:
		t.Fatal("SubmitWait did not wait for function to execute")
	}
}

func TestOverflow(t *testing.T) {
	wp := New(2)
	releaseChan := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < 64; i++ {
		wp.Submit(func() { <-releaseChan })
	}

	// Start a goroutine to free the workers after calling stop.  This way
	// the dispatcher can exit, then when this goroutine runs, the workerpool
	// can exit.
	go func() {
		<-time.After(time.Millisecond)
		close(releaseChan)
	}()
	wp.Stop()

	// Now that the worker pool has exited, it is safe to inspect its waiting
	// queue without causing a race.
	qlen := wp.waitingQueue.Len()
	if qlen != 62 {
		t.Fatal("Expected 62 tasks in waiting queue, have", qlen)
	}
}

func TestStopRace(t *testing.T) {
	wp := New(20)
	releaseChan := make(chan struct{})
	workRelChan := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < 20; i++ {
		wp.Submit(func() { <-workRelChan })
	}

	time.Sleep(5 * time.Second)
	for i := 0; i < 64; i++ {
		go func() {
			<-releaseChan
			wp.Stop()
		}()
	}

	close(workRelChan)
	close(releaseChan)
}

// Run this test with race detector to test that using WaitingQueueSize has no
// race condition
func TestWaitingQueueSizeRace(t *testing.T) {
	const (
		goroutines = 10
		tasks      = 20
		workers    = 5
	)
	wp := New(workers)
	maxChan := make(chan int)
	for g := 0; g < goroutines; g++ {
		go func() {
			max := 0
			// Submit 100 tasks, checking waiting queue size each time.  Report
			// the maximum queue size seen.
			for i := 0; i < tasks; i++ {
				wp.Submit(func() {
					time.Sleep(time.Microsecond)
				})
				waiting := wp.WaitingQueueSize()
				if waiting > max {
					max = waiting
				}
			}
			maxChan <- max
		}()
	}

	// Find maximum queuesize seen by any thread.
	maxMax := 0
	for g := 0; g < goroutines; g++ {
		max := <-maxChan
		if max > maxMax {
			maxMax = max
		}
	}
	if maxMax == 0 {
		t.Error("expected to see waiting queue size > 0")
	}
	if maxMax >= goroutines*tasks {
		t.Error("should not have seen all tasks on waiting queue")
	}
}

func anyReady(w *WorkerPool) bool {
	select {
	case wkCh := <-w.readyWorkers:
		w.readyWorkers <- wkCh
		return true
	default:
	}
	return false
}

func countReady(w *WorkerPool) int {
	// Try to pull max workers off of ready queue.
	timeout := time.After(5 * time.Second)
	readyTmp := make(chan chan func(), max)
	var readyCount int
	for i := 0; i < max; i++ {
		select {
		case wkCh := <-w.readyWorkers:
			readyTmp <- wkCh
			readyCount++
		case <-timeout:
			readyCount = i
			i = max
		}
	}

	// Restore ready workers.
	close(readyTmp)
	go func() {
		for r := range readyTmp {
			w.readyWorkers <- r
		}
	}()
	return readyCount
}

/*

Run benchmarking with: go test -bench '.'

*/

func BenchmarkEnqueue(b *testing.B) {
	wp := New(1)
	defer wp.Stop()
	releaseChan := make(chan struct{})

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() { <-releaseChan })
	}
	close(releaseChan)
}

func BenchmarkEnqueue2(b *testing.B) {
	wp := New(2)
	defer wp.Stop()
	releaseChan := make(chan struct{})

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		for i := 0; i < 64; i++ {
			wp.Submit(func() { <-releaseChan })
		}
		for i := 0; i < 64; i++ {
			releaseChan <- struct{}{}
		}
	}
	close(releaseChan)
}

func BenchmarkExecute1Worker(b *testing.B) {
	wp := New(1)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute2Worker(b *testing.B) {
	wp := New(2)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute4Workers(b *testing.B) {
	wp := New(4)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute16Workers(b *testing.B) {
	wp := New(16)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute64Workers(b *testing.B) {
	wp := New(64)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}
<|endoftext|>


<commit_before>package controllers

import (
	"archive/tar"
	"errors"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/cloud/storage"

	"github.com/coduno/api/model"
	"github.com/coduno/api/test"
	"github.com/coduno/api/util"
	"github.com/coduno/api/util/passenger"
	"github.com/gorilla/mux"

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

func init() {
	router.Handle("/results/{resultKey}/tasks/{taskKey}/submissions", ContextHandlerFunc(PostSubmission))
	router.Handle("/results/{resultKey}/finalSubmissions/{index}", ContextHandlerFunc(FinalSubmission))
}

// PostSubmission creates a new submission.
func PostSubmission(ctx context.Context, w http.ResponseWriter, r *http.Request) (status int, err error) {
	if r.Method != "POST" {
		return http.StatusMethodNotAllowed, nil
	}

	p, ok := passenger.FromContext(ctx)
	if !ok {
		return http.StatusUnauthorized, nil
	}

	mediaType, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil {
		return http.StatusBadRequest, err
	}
	if !strings.HasPrefix(mediaType, "multipart/") {
		return http.StatusUnsupportedMediaType, nil
	}

	resultKey, err := datastore.DecodeKey(mux.Vars(r)["resultKey"])
	if err != nil {
		return http.StatusNotFound, err
	}

	if !util.HasParent(p.User, resultKey) {
		return http.StatusBadRequest, errors.New("cannot submit answer for other users")
	}

	taskKey, err := datastore.DecodeKey(mux.Vars(r)["taskKey"])
	if err != nil {
		return http.StatusNotFound, err
	}

	form, err := multipart.NewReader(r.Body, params["boundary"]).ReadForm(16 << 20)
	if err != nil {
		return http.StatusBadRequest, err
	}

	var task model.Task
	if err = datastore.Get(ctx, taskKey, &task); err != nil {
		return http.StatusNotFound, err
	}

	// Furthermore, the name of the GCS object is derived from the of the
	// encapsulating Submission. To avoid race conditions, allocate an ID.
	low, _, err := datastore.AllocateIDs(ctx, model.SubmissionKind, resultKey, 1)
	if err != nil {
		return http.StatusInternalServerError, err
	}

	submissionKey := datastore.NewKey(ctx, model.SubmissionKind, "", low, resultKey)
	storedCode := model.StoredObject{
		Bucket: util.SubmissionBucket(),
		Name:   nameObject(submissionKey) + "Code/",
	}
	submission := model.Submission{
		Task:     taskKey,
		Time:     time.Now(),
		Language: detectLanguage(form),
		Code:     storedCode,
	}

	if _, err = datastore.Put(ctx, submissionKey, &submission); err != nil {
		return http.StatusInternalServerError, err
	}

	pr, pw := io.Pipe()
	go maketar(pw, form)

	var tests model.Tests
	testKeys, err := model.NewQueryForTest().
		Ancestor(taskKey).
		GetAll(ctx, &tests)
	if err != nil {
		return http.StatusInternalServerError, err
	}

	for i, t := range tests {
		go func(i int, t model.Test) {
			// TODO(victorbalan, flowlo): Error handling
			if err := test.Tester(t.Tester).Call(ctx, *t.Key(testKeys[i]), *submission.Key(submissionKey), pr); err != nil {
				log.Warningf(ctx, "%s", err)
			}
		}(i, t)
	}

	if err := upload(ctx, storedCode.Bucket, storedCode.Name, form); err != nil {
		return http.StatusInternalServerError, err
	}

	// TODO(flowlo): Return something meaningful.
	return http.StatusOK, nil
}

type multiError []error

func (e multiError) Error() string {
	s := make([]string, 0, len(e))
	for i := range e {
		s = append(s, e[i].Error())
	}
	return strings.Join(s, "\n")
}

func upload(ctx context.Context, bucket, base string, form *multipart.Form) multiError {
	if appengine.IsDevAppServer() {
		return nil
	}

	errc := make(chan error)
	var errs []error

	for _, fhs := range form.File {
		go func(fhs []*multipart.FileHeader) {
			if len(fhs) > 1 {
				errc <- errors.New("cannot read split files")
				return
			}

			fh := fhs[0]
			f, err := fh.Open()
			if err != nil {
				errc <- err
				return
			}

			name := base + fh.Filename
			wc := storage.NewWriter(ctx, bucket, name)

			if _, err := io.Copy(wc, f); err != nil {
				errc <- err
				return
			}

			if err := wc.Close(); err != nil {
				errc <- err
				return
			}

			if err := f.Close(); err != nil {
				errc <- err
				return
			}

			errc <- nil
		}(fhs)
	}

	for range form.File {
		err := <-errc
		if err != nil {
			errs = append(errs, err)
		}
	}

	close(errc)
	return errs
}

func maketar(wc io.WriteCloser, form *multipart.Form) error {
	defer wc.Close()

	tarw := tar.NewWriter(wc)
	defer tarw.Close()

	sizeFunc := func(s io.Seeker) int64 {
		size, err := s.Seek(0, os.SEEK_END)
		if err != nil {
			return -1
		}
		if _, err = s.Seek(0, os.SEEK_SET); err != nil {
			return -1
		}
		return size
	}

	for _, fhs := range form.File {
		if len(fhs) > 1 {
			return errors.New("cannot read split files")
		}
		fh := fhs[0]
		f, err := fh.Open()
		if err != nil {
			return err
		}
		size := sizeFunc(f)
		if size < 0 {
			return errors.New("seeker can't seek")
		}
		tarw.WriteHeader(&tar.Header{
			Name: fh.Filename,
			Mode: 0600,
			Size: size,
		})
		io.Copy(tarw, f)
		f.Close()
	}
	return nil
}

func detectLanguage(form *multipart.Form) string {
	l := ""
	m := map[string]int{
		"py":   0,
		"java": 0,
		"c":    0,
		"cpp":  0,
	}
	max := 0
	for name := range form.File {
		dot := strings.LastIndex(name, ".") + 1
		if dot == 0 || dot == len(name) {
			continue
		}
		ext := name[dot:]
		cnt, ok := m[ext]
		if !ok {
			continue
		}
		cnt++
		if cnt > max {
			l = ext
		}
		m[ext] = cnt
	}
	return l
}

// FinalSubmission makes the last submission final.
func FinalSubmission(ctx context.Context, w http.ResponseWriter, r *http.Request) (status int, err error) {
	if r.Method != "POST" {
		return http.StatusMethodNotAllowed, nil
	}

	p, ok := passenger.FromContext(ctx)
	if !ok {
		return http.StatusUnauthorized, nil
	}

	var resultKey *datastore.Key
	if resultKey, err = datastore.DecodeKey(mux.Vars(r)["resultKey"]); err != nil {
		return http.StatusInternalServerError, err
	}

	if !util.HasParent(p.User, resultKey) {
		return http.StatusBadRequest, errors.New("cannot submit answer for other users")
	}

	var index int
	if index, err = strconv.Atoi(mux.Vars(r)["index"]); err != nil {
		return http.StatusInternalServerError, err
	}

	if len(r.URL.Query()["submissionKey"]) == 0 {
		return http.StatusOK, nil
	}
	var submissionKey *datastore.Key
	if submissionKey, err = datastore.DecodeKey(r.URL.Query()["submissionKey"][0]); err != nil {
		return http.StatusInternalServerError, err
	}

	var result model.Result
	if err = datastore.Get(ctx, resultKey, &result); err != nil {
		return http.StatusInternalServerError, err
	}

	result.FinalSubmissions[index] = submissionKey

	if _, err = result.Put(ctx, resultKey); err != nil {
		return http.StatusInternalServerError, err
	}
	w.Write([]byte("OK"))
	return
}

func nameObject(key *datastore.Key) string {
	name := ""
	for key != nil {
		id := key.StringID()
		if id == "" {
			id = strconv.FormatInt(key.IntID(), 10)
		}
		name = "/" + key.Kind() + "/" + id + name
		key = key.Parent()
	}
	// NOTE: The name of a GCS object must not be prefixed "/",
	// this will give you a major headache when working with
	// gcsfuse. Just don't.
	return name[1:]
}

func defaultObjectAttrs(disposition string) storage.ObjectAttrs {
	// TODO(flowlo): Establish ACLs?
	return storage.ObjectAttrs{
		ContentType:        "text/plain", // TODO(flowlo): Content types per language?
		ContentLanguage:    "",           // TODO(flowlo): Does it make sense to set this?
		ContentEncoding:    "utf-8",
		CacheControl:       "max-age=31556926", // Aggressive caching for one year
		ContentDisposition: "attachment; filename=\"" + disposition + "\"",
	}
}
<commit_msg>submission: Fix multipart file handling<commit_after>package controllers

import (
	"archive/tar"
	"errors"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/cloud/storage"

	"github.com/coduno/api/model"
	"github.com/coduno/api/test"
	"github.com/coduno/api/util"
	"github.com/coduno/api/util/passenger"
	"github.com/gorilla/mux"

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

func init() {
	router.Handle("/results/{resultKey}/tasks/{taskKey}/submissions", ContextHandlerFunc(PostSubmission))
	router.Handle("/results/{resultKey}/finalSubmissions/{index}", ContextHandlerFunc(FinalSubmission))
}

// PostSubmission creates a new submission.
func PostSubmission(ctx context.Context, w http.ResponseWriter, r *http.Request) (status int, err error) {
	if r.Method != "POST" {
		return http.StatusMethodNotAllowed, nil
	}

	p, ok := passenger.FromContext(ctx)
	if !ok {
		return http.StatusUnauthorized, nil
	}

	mediaType, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil {
		return http.StatusBadRequest, err
	}
	if !strings.HasPrefix(mediaType, "multipart/") {
		return http.StatusUnsupportedMediaType, nil
	}

	resultKey, err := datastore.DecodeKey(mux.Vars(r)["resultKey"])
	if err != nil {
		return http.StatusNotFound, err
	}

	if !util.HasParent(p.User, resultKey) {
		return http.StatusBadRequest, errors.New("cannot submit answer for other users")
	}

	taskKey, err := datastore.DecodeKey(mux.Vars(r)["taskKey"])
	if err != nil {
		return http.StatusNotFound, err
	}

	form, err := multipart.NewReader(r.Body, params["boundary"]).ReadForm(16 << 20)
	if err != nil {
		return http.StatusBadRequest, err
	}

	files, ok := form.File["files"]
	if !ok {
		return http.StatusBadRequest, errors.New("missing files")
	}

	var task model.Task
	if err = datastore.Get(ctx, taskKey, &task); err != nil {
		return http.StatusNotFound, err
	}

	// Furthermore, the name of the GCS object is derived from the of the
	// encapsulating Submission. To avoid race conditions, allocate an ID.
	low, _, err := datastore.AllocateIDs(ctx, model.SubmissionKind, resultKey, 1)
	if err != nil {
		return http.StatusInternalServerError, err
	}

	submissionKey := datastore.NewKey(ctx, model.SubmissionKind, "", low, resultKey)
	storedCode := model.StoredObject{
		Bucket: util.SubmissionBucket(),
		Name:   nameObject(submissionKey) + "Code/",
	}
	submission := model.Submission{
		Task:     taskKey,
		Time:     time.Now(),
		Language: detectLanguage(files),
		Code:     storedCode,
	}

	if _, err = datastore.Put(ctx, submissionKey, &submission); err != nil {
		return http.StatusInternalServerError, err
	}

	pr, pw := io.Pipe()
	go maketar(pw, files)

	var tests model.Tests
	testKeys, err := model.NewQueryForTest().
		Ancestor(taskKey).
		GetAll(ctx, &tests)
	if err != nil {
		return http.StatusInternalServerError, err
	}

	for i, t := range tests {
		go func(i int, t model.Test) {
			// TODO(victorbalan, flowlo): Error handling
			if err := test.Tester(t.Tester).Call(ctx, *t.Key(testKeys[i]), *submission.Key(submissionKey), pr); err != nil {
				log.Warningf(ctx, "%s", err)
			}
		}(i, t)
	}

	if err := upload(ctx, storedCode.Bucket, storedCode.Name, files); err != nil {
		return http.StatusInternalServerError, err
	}

	// TODO(flowlo): Return something meaningful.
	return http.StatusOK, nil
}

type multiError []error

func (e multiError) Error() string {
	s := make([]string, 0, len(e))
	for i := range e {
		s = append(s, e[i].Error())
	}
	return strings.Join(s, "\n")
}

func upload(ctx context.Context, bucket, base string, files []*multipart.FileHeader) multiError {
	if appengine.IsDevAppServer() {
		return nil
	}

	errc := make(chan error)
	var errs []error

	for _, fh := range files {
		go func(fh *multipart.FileHeader) {
			f, err := fh.Open()
			if err != nil {
				errc <- err
				return
			}

			name := base + fh.Filename
			wc := storage.NewWriter(ctx, bucket, name)

			if _, err := io.Copy(wc, f); err != nil {
				errc <- err
				return
			}

			if err := wc.Close(); err != nil {
				errc <- err
				return
			}

			if err := f.Close(); err != nil {
				errc <- err
				return
			}

			errc <- nil
		}(fh)
	}

	for range files {
		err := <-errc
		if err != nil {
			errs = append(errs, err)
		}
	}

	close(errc)
	return errs
}

func maketar(wc io.WriteCloser, files []*multipart.FileHeader) error {
	defer wc.Close()

	tarw := tar.NewWriter(wc)
	defer tarw.Close()

	sizeFunc := func(s io.Seeker) int64 {
		size, err := s.Seek(0, os.SEEK_END)
		if err != nil {
			return -1
		}
		if _, err = s.Seek(0, os.SEEK_SET); err != nil {
			return -1
		}
		return size
	}

	for _, fh := range files {
		f, err := fh.Open()
		if err != nil {
			return err
		}
		size := sizeFunc(f)
		if size < 0 {
			return errors.New("seeker can't seek")
		}
		tarw.WriteHeader(&tar.Header{
			Name: fh.Filename,
			Mode: 0600,
			Size: size,
		})
		io.Copy(tarw, f)
		f.Close()
	}
	return nil
}

func detectLanguage(files []*multipart.FileHeader) string {
	l := ""
	m := map[string]int{
		"py":   0,
		"java": 0,
		"c":    0,
		"cpp":  0,
	}
	max := 0

	for _, fh := range files {
		name := fh.Filename
		dot := strings.LastIndex(name, ".") + 1
		if dot == 0 || dot == len(name) {
			continue
		}
		ext := name[dot:]
		cnt, ok := m[ext]
		if !ok {
			continue
		}
		cnt++
		if cnt > max {
			l = ext
		}
		m[ext] = cnt
	}
	return l
}

// FinalSubmission makes the last submission final.
func FinalSubmission(ctx context.Context, w http.ResponseWriter, r *http.Request) (status int, err error) {
	if r.Method != "POST" {
		return http.StatusMethodNotAllowed, nil
	}

	p, ok := passenger.FromContext(ctx)
	if !ok {
		return http.StatusUnauthorized, nil
	}

	var resultKey *datastore.Key
	if resultKey, err = datastore.DecodeKey(mux.Vars(r)["resultKey"]); err != nil {
		return http.StatusInternalServerError, err
	}

	if !util.HasParent(p.User, resultKey) {
		return http.StatusBadRequest, errors.New("cannot submit answer for other users")
	}

	var index int
	if index, err = strconv.Atoi(mux.Vars(r)["index"]); err != nil {
		return http.StatusInternalServerError, err
	}

	if len(r.URL.Query()["submissionKey"]) == 0 {
		return http.StatusOK, nil
	}
	var submissionKey *datastore.Key
	if submissionKey, err = datastore.DecodeKey(r.URL.Query()["submissionKey"][0]); err != nil {
		return http.StatusInternalServerError, err
	}

	var result model.Result
	if err = datastore.Get(ctx, resultKey, &result); err != nil {
		return http.StatusInternalServerError, err
	}

	result.FinalSubmissions[index] = submissionKey

	if _, err = result.Put(ctx, resultKey); err != nil {
		return http.StatusInternalServerError, err
	}
	w.Write([]byte("OK"))
	return
}

func nameObject(key *datastore.Key) string {
	name := ""
	for key != nil {
		id := key.StringID()
		if id == "" {
			id = strconv.FormatInt(key.IntID(), 10)
		}
		name = "/" + key.Kind() + "/" + id + name
		key = key.Parent()
	}
	// NOTE: The name of a GCS object must not be prefixed "/",
	// this will give you a major headache when working with
	// gcsfuse. Just don't.
	return name[1:]
}

func defaultObjectAttrs(disposition string) storage.ObjectAttrs {
	// TODO(flowlo): Establish ACLs?
	return storage.ObjectAttrs{
		ContentType:        "text/plain", // TODO(flowlo): Content types per language?
		ContentLanguage:    "",           // TODO(flowlo): Does it make sense to set this?
		ContentEncoding:    "utf-8",
		CacheControl:       "max-age=31556926", // Aggressive caching for one year
		ContentDisposition: "attachment; filename=\"" + disposition + "\"",
	}
}
<|endoftext|>


<commit_before>package geocodio_test

import (
	"testing"

	"github.com/stevepartridge/geocodio"
)

const (
	// AddressTestOneFull       = "42370 Bob Hope Dr, Rancho Mirage, CA 92270"
	// AddressTestOneWithoutZip = "42370 Bob Hope Dr, Rancho Mirage, CA"
	// AddressTestOneNumber     = "42370"
	// AddressTestOneStreet     = "Bob Hope Dr"
	// AddressTestOneCity       = "Rancho Mirage"
	// AddressTestOneState      = "CA"
	// AddressTestOneLatitude   = 33.738645
	// AddressTestOneLongitude  = -116.407141

	AddressTestOneFull       = "1109 N Highland St, Arlington, VA 22201"
	AddressTestOneWithoutZip = "1109 N Highland St, Arlington, VA"
	AddressTestOneNumber     = "1109"
	AddressTestOneStreet     = "N Highland St"
	AddressTestOneCity       = "Arlington"
	AddressTestOneState      = "VA"
	AddressTestOneLatitude   = 38.886672
	AddressTestOneLongitude  = -77.094735

	AddressTestTwoFull       = "100 Legends Way, Boston, MA 02114"
	AddressTestTwoWithoutZip = "100 Legends Way, Boston, MA"
	AddressTestTwoNumber     = "100"
	AddressTestTwoStreet     = "Legends Way"
	AddressTestTwoCity       = "Boston"
	AddressTestTwoState      = "MA"
	AddressTestTwoLatitude   = 42.36629
	AddressTestTwoLongitude  = -71.0622
)

func TestGeocodioWithApiKey(t *testing.T) {
	_, err := geocodio.New()
	if err != nil {
		t.Error("Failed with API KEY set.", err)
	}
}

func TestGeocodioWithoutApiKey(t *testing.T) {
	_, err := geocodio.New("")
	if err == nil {
		t.Error("Did not through error when omitting API KEY")
	}
}

// 1109+N+Highland+St%2c+Arlington+VA
<commit_msg>Add test for deprecated method usage<commit_after>package geocodio_test

import (
	"os"
	"testing"

	"github.com/stevepartridge/geocodio"
)

const (
	// AddressTestOneFull       = "42370 Bob Hope Dr, Rancho Mirage, CA 92270"
	// AddressTestOneWithoutZip = "42370 Bob Hope Dr, Rancho Mirage, CA"
	// AddressTestOneNumber     = "42370"
	// AddressTestOneStreet     = "Bob Hope Dr"
	// AddressTestOneCity       = "Rancho Mirage"
	// AddressTestOneState      = "CA"
	// AddressTestOneLatitude   = 33.738645
	// AddressTestOneLongitude  = -116.407141

	AddressTestOneFull       = "1109 N Highland St, Arlington, VA 22201"
	AddressTestOneWithoutZip = "1109 N Highland St, Arlington, VA"
	AddressTestOneNumber     = "1109"
	AddressTestOneStreet     = "N Highland St"
	AddressTestOneCity       = "Arlington"
	AddressTestOneState      = "VA"
	AddressTestOneLatitude   = 38.886672
	AddressTestOneLongitude  = -77.094735

	AddressTestTwoFull       = "100 Legends Way, Boston, MA 02114"
	AddressTestTwoWithoutZip = "100 Legends Way, Boston, MA"
	AddressTestTwoNumber     = "100"
	AddressTestTwoStreet     = "Legends Way"
	AddressTestTwoCity       = "Boston"
	AddressTestTwoState      = "MA"
	AddressTestTwoLatitude   = 42.36629
	AddressTestTwoLongitude  = -71.0622
)

func TestGeocodioWithApiKey(t *testing.T) {
	_, err := geocodio.New()
	if err != nil {
		t.Error("Failed with API KEY set.", err)
	}
}

func TestGeocodioWithoutApiKey(t *testing.T) {
	_, err := geocodio.New("")
	if err == nil {
		t.Error("Did not through error when omitting API KEY")
	}
}

func TestGeocodioDeprecatedWithApiKey(t *testing.T) {
	_, err := geocodio.NewGeocodio(os.Getenv("API_KEY"))
	if err != nil {
		t.Error("Failed with API KEY set.", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"go/token"

	"github.com/davyxu/gosproto/meta"
)

const goCodeTemplate = `// Generated by github.com/davyxu/gosproto/sprotogen
// DO NOT EDIT!

package {{.PackageName}}

import (
	"reflect"
	"github.com/davyxu/gosproto"
	"github.com/davyxu/goobjfmt"
	{{if .CellnetReg}}"github.com/davyxu/cellnet/codec/sproto"{{end}}
)

{{range $a, $enumobj := .Enums}}
type {{.Name}} int32
const (	{{range .StFields}}
	{{$enumobj.Name}}_{{.Name}} {{$enumobj.Name}} = {{.TagNumber}} {{end}}
)

var {{$enumobj.Name}}_ValueByName = map[string]int32{ {{range .StFields}}
	"{{.Name}}": {{.TagNumber}}, {{end}}
}

var {{$enumobj.Name}}_NameByValue = map[int32]string{ {{range .StFields}}
	{{.TagNumber}}: "{{.Name}}" , {{end}}
}

func (self {{$enumobj.Name}}) String() string {
	return sproto.EnumName({{$enumobj.Name}}_NameByValue, int32(self))
}
{{end}}

{{range $a, $stobj := .Structs}}
type {{.Name}} struct{
	{{range .StFields}}
		{{.GoFieldName}} {{.GoTypeName}} {{.GoTags}}
	{{end}}
}

func (self *{{.Name}}) String() string { return goobjfmt.CompactTextString(self) }

{{range .StFields}}{{if .IsExtendType}}
func (self *{{$stobj.Name}}) {{.GoExtendFieldGetterName}}() {{.GoExtendFieldGetterType}} {
	{{.GoExtendFieldGetter}}
}
func (self *{{$stobj.Name}}) Set{{.GoExtendFieldGetterName}}(v {{.GoExtendFieldGetterType}})  {
	{{.GoExtendFieldSetter}}
}
{{end}} {{end}}
{{end}}

var SProtoStructs = []reflect.Type{
{{range .Structs}}
	reflect.TypeOf((*{{.Name}})(nil)).Elem(), // {{.MsgID}} {{end}}
}

{{if .CellnetReg}}
func init() {
	sprotocodec.AutoRegisterMessageMeta(SProtoStructs)
}
{{end}}

`

func (self *fieldModel) GoExtendFieldGetterName() string {
	pname := publicFieldName(self.Name)

	if token.Lookup(pname).IsKeyword() {
		return pname + "_"
	}

	return pname
}

func (self *fieldModel) GoExtendFieldGetter() string {

	switch self.Type {
	case meta.FieldType_Float32,
		meta.FieldType_Float64:
		return fmt.Sprintf("return %s(self.%s) * %f", self.GoExtendFieldGetterType(), self.GoFieldName(), 1.0/float32(self.ExtendTypePrecision()))
	}

	return "unknown extend type:" + self.Type.String()
}

func (self *fieldModel) GoExtendFieldSetter() string {

	switch self.Type {
	case meta.FieldType_Float32:
		return fmt.Sprintf("self.%s = int32(v* %d)", self.GoFieldName(), self.ExtendTypePrecision())
	case meta.FieldType_Float64:
		return fmt.Sprintf("self.%s = int64(v* %d)", self.GoFieldName(), self.ExtendTypePrecision())
	}

	return "unknown extend type:" + self.Type.String()
}

func (self *fieldModel) GoExtendFieldGetterType() string {
	var b bytes.Buffer
	if self.Repeatd {
		b.WriteString("[]")
	}

	// 字段类型映射go的类型
	switch self.Type {
	case meta.FieldType_Float32:
		b.WriteString("float32")
	case meta.FieldType_Float64:
		b.WriteString("float64")
	default:
		b.WriteString("unknown extend type:" + self.Type.String())
	}

	return b.String()
}

func (self *fieldModel) GoFieldName() string {

	var pname string

	// 扩展类型不能直接访问
	if self.IsExtendType() {
		pname = "Extend_" + self.Name
	} else {
		pname = publicFieldName(self.Name)
	}

	// 碰到关键字在尾部加_
	if token.Lookup(pname).IsKeyword() {
		return pname + "_"
	}

	return pname
}

func (self *fieldModel) GoTypeName() string {

	var b bytes.Buffer
	if self.Repeatd {
		b.WriteString("[]")
	}

	if self.Type == meta.FieldType_Struct {
		b.WriteString("*")
	}

	// 字段类型映射go的类型
	switch self.Type {
	case meta.FieldType_Integer:
		b.WriteString("int")
	case meta.FieldType_Bool:
		b.WriteString("bool")
	case meta.FieldType_Struct,
		meta.FieldType_Enum:
		b.WriteString(self.Complex.Name)
	case meta.FieldType_Float32:
		b.WriteString("int32")
	case meta.FieldType_Float64:
		b.WriteString("int64")
	default:
		b.WriteString(self.Type.String())
	}

	return b.String()
}

func (self *fieldModel) GoTags() string {

	var b bytes.Buffer

	b.WriteString("`sproto:\"")

	// 整形类型对解码层都视为整形
	switch self.Type {
	case meta.FieldType_Int32,
		meta.FieldType_Int64,
		meta.FieldType_UInt32,
		meta.FieldType_UInt64,
		meta.FieldType_Float32,
		meta.FieldType_Float64,
		meta.FieldType_Enum:
		b.WriteString("integer")
	default:
		b.WriteString(self.Kind())
	}

	b.WriteString(",")

	b.WriteString(fmt.Sprintf("%d", self.TagNumber()))
	b.WriteString(",")

	if self.Repeatd {
		b.WriteString("array,")
	}

	b.WriteString(fmt.Sprintf("name=%s", self.GoFieldName()))

	b.WriteString("\"`")

	return b.String()
}

func gen_go(fileD *meta.FileDescriptor, packageName, filename string, cellnetReg bool) {

	fm := &fileModel{
		FileDescriptor: fileD,
		PackageName:    packageName,
		CellnetReg:     cellnetReg,
	}

	addData(fm, fileD)

	generateCode("sp->go", goCodeTemplate, filename, fm, &generateOption{
		formatGoCode: true,
	})

}
<commit_msg>fixed: golang enum mapper var name<commit_after>package main

import (
	"bytes"
	"fmt"
	"go/token"

	"github.com/davyxu/gosproto/meta"
)

const goCodeTemplate = `// Generated by github.com/davyxu/gosproto/sprotogen
// DO NOT EDIT!

package {{.PackageName}}

import (
	"reflect"
	"github.com/davyxu/gosproto"
	"github.com/davyxu/goobjfmt"
	{{if .CellnetReg}}"github.com/davyxu/cellnet/codec/sproto"{{end}}
)

{{range $a, $enumobj := .Enums}}
type {{.Name}} int32
const (	{{range .StFields}}
	{{$enumobj.Name}}_{{.Name}} {{$enumobj.Name}} = {{.TagNumber}} {{end}}
)

var {{$enumobj.Name}}MapperValueByName = map[string]int32{ {{range .StFields}}
	"{{.Name}}": {{.TagNumber}}, {{end}}
}

var {{$enumobj.Name}}MapperNameByValue = map[int32]string{ {{range .StFields}}
	{{.TagNumber}}: "{{.Name}}" , {{end}}
}

func (self {{$enumobj.Name}}) String() string {
	return sproto.EnumName({{$enumobj.Name}}MapperNameByValue, int32(self))
}
{{end}}

{{range $a, $stobj := .Structs}}
type {{.Name}} struct{
	{{range .StFields}}
		{{.GoFieldName}} {{.GoTypeName}} {{.GoTags}}
	{{end}}
}

func (self *{{.Name}}) String() string { return goobjfmt.CompactTextString(self) }

{{range .StFields}}{{if .IsExtendType}}
func (self *{{$stobj.Name}}) {{.GoExtendFieldGetterName}}() {{.GoExtendFieldGetterType}} {
	{{.GoExtendFieldGetter}}
}
func (self *{{$stobj.Name}}) Set{{.GoExtendFieldGetterName}}(v {{.GoExtendFieldGetterType}})  {
	{{.GoExtendFieldSetter}}
}
{{end}} {{end}}
{{end}}

var SProtoStructs = []reflect.Type{
{{range .Structs}}
	reflect.TypeOf((*{{.Name}})(nil)).Elem(), // {{.MsgID}} {{end}}
}

{{if .CellnetReg}}
func init() {
	sprotocodec.AutoRegisterMessageMeta(SProtoStructs)
}
{{end}}

`

func (self *fieldModel) GoExtendFieldGetterName() string {
	pname := publicFieldName(self.Name)

	if token.Lookup(pname).IsKeyword() {
		return pname + "_"
	}

	return pname
}

func (self *fieldModel) GoExtendFieldGetter() string {

	switch self.Type {
	case meta.FieldType_Float32,
		meta.FieldType_Float64:
		return fmt.Sprintf("return %s(self.%s) * %f", self.GoExtendFieldGetterType(), self.GoFieldName(), 1.0/float32(self.ExtendTypePrecision()))
	}

	return "unknown extend type:" + self.Type.String()
}

func (self *fieldModel) GoExtendFieldSetter() string {

	switch self.Type {
	case meta.FieldType_Float32:
		return fmt.Sprintf("self.%s = int32(v* %d)", self.GoFieldName(), self.ExtendTypePrecision())
	case meta.FieldType_Float64:
		return fmt.Sprintf("self.%s = int64(v* %d)", self.GoFieldName(), self.ExtendTypePrecision())
	}

	return "unknown extend type:" + self.Type.String()
}

func (self *fieldModel) GoExtendFieldGetterType() string {
	var b bytes.Buffer
	if self.Repeatd {
		b.WriteString("[]")
	}

	// 字段类型映射go的类型
	switch self.Type {
	case meta.FieldType_Float32:
		b.WriteString("float32")
	case meta.FieldType_Float64:
		b.WriteString("float64")
	default:
		b.WriteString("unknown extend type:" + self.Type.String())
	}

	return b.String()
}

func (self *fieldModel) GoFieldName() string {

	var pname string

	// 扩展类型不能直接访问
	if self.IsExtendType() {
		pname = "Extend_" + self.Name
	} else {
		pname = publicFieldName(self.Name)
	}

	// 碰到关键字在尾部加_
	if token.Lookup(pname).IsKeyword() {
		return pname + "_"
	}

	return pname
}

func (self *fieldModel) GoTypeName() string {

	var b bytes.Buffer
	if self.Repeatd {
		b.WriteString("[]")
	}

	if self.Type == meta.FieldType_Struct {
		b.WriteString("*")
	}

	// 字段类型映射go的类型
	switch self.Type {
	case meta.FieldType_Integer:
		b.WriteString("int")
	case meta.FieldType_Bool:
		b.WriteString("bool")
	case meta.FieldType_Struct,
		meta.FieldType_Enum:
		b.WriteString(self.Complex.Name)
	case meta.FieldType_Float32:
		b.WriteString("int32")
	case meta.FieldType_Float64:
		b.WriteString("int64")
	default:
		b.WriteString(self.Type.String())
	}

	return b.String()
}

func (self *fieldModel) GoTags() string {

	var b bytes.Buffer

	b.WriteString("`sproto:\"")

	// 整形类型对解码层都视为整形
	switch self.Type {
	case meta.FieldType_Int32,
		meta.FieldType_Int64,
		meta.FieldType_UInt32,
		meta.FieldType_UInt64,
		meta.FieldType_Float32,
		meta.FieldType_Float64,
		meta.FieldType_Enum:
		b.WriteString("integer")
	default:
		b.WriteString(self.Kind())
	}

	b.WriteString(",")

	b.WriteString(fmt.Sprintf("%d", self.TagNumber()))
	b.WriteString(",")

	if self.Repeatd {
		b.WriteString("array,")
	}

	b.WriteString(fmt.Sprintf("name=%s", self.GoFieldName()))

	b.WriteString("\"`")

	return b.String()
}

func gen_go(fileD *meta.FileDescriptor, packageName, filename string, cellnetReg bool) {

	fm := &fileModel{
		FileDescriptor: fileD,
		PackageName:    packageName,
		CellnetReg:     cellnetReg,
	}

	addData(fm, fileD)

	generateCode("sp->go", goCodeTemplate, filename, fm, &generateOption{
		formatGoCode: true,
	})

}
<|endoftext|>


<commit_before>// Copyright © 2017 The vt-go authors. All Rights Reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package vt

import (
	"bytes"
	"encoding/json"
	"fmt"
	"time"
)

// ObjectDescriptor is a pair (ID, type) describing a VirusTotal API object.
type ObjectDescriptor struct {
	ID                string                 `json:"id,omitempty"`
	Type              string                 `json:"type,omitempty"`
	ContextAttributes map[string]interface{} `json:"context_attributes,omitempty"`
}

// Object represents a VirusTotal API object.
type Object struct {
	ID                string                   `json:"id,omitempty"`
	Type              string                   `json:"type,omitempty"`
	Attributes        map[string]interface{}   `json:"attributes,omitempty"`
	ContextAttributes map[string]interface{}   `json:"context_attributes,omitempty"`
	Relationships     map[string]*Relationship `json:"relationships,omitempty"`
	Links             Links                    `json:"links,omitempty"`
}

// Links contains links related to an API object.
type Links struct {
	Self string `json:"self"`
	Next string `json:"next"`
}

// Relationship contains information about a related API object.
type Relationship struct {
	Data  json.RawMessage `json:"data,omitempty"`
	Links Links           `json:"links,omitempty"`

	// IsOneToOne is true if this is a one-to-one relationshio and False if
	// otherwise. If true RelatedObjects contains one object at most.
	IsOneToOne     bool
	RelatedObjects []ObjectDescriptor
}

// NewObject creates a new object.
func NewObject() *Object {
	return &Object{Attributes: make(map[string]interface{})}
}

// UnmarshalJSON unmarshals a VirusTotal API object from data.
func (obj *Object) UnmarshalJSON(data []byte) error {

	type object Object

	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.UseNumber()

	o := object{}
	if err := decoder.Decode(&o); err != nil {
		return err
	}

	obj.Type = o.Type
	obj.ID = o.ID
	obj.Attributes = o.Attributes
	obj.ContextAttributes = o.ContextAttributes
	obj.Relationships = o.Relationships

	for _, v := range obj.Relationships {
		// Try unmarshalling as an array first, if it fails this is a one-to-one
		// relationship, so we should try unmarshalling a single object descriptor.
		if err := json.Unmarshal(v.Data, &v.RelatedObjects); err != nil {
			v.IsOneToOne = true
			var o ObjectDescriptor
			if err = json.Unmarshal(v.Data, &o); err != nil {
				return err
			}
			v.RelatedObjects = append(v.RelatedObjects, o)
		}
	}

	return nil
}

func (obj *Object) getAttributeNumber(name string) (n json.Number, err error) {
	if attrValue, attrExists := obj.Attributes[name]; attrExists {
		n, isNumber := attrValue.(json.Number)
		if !isNumber {
			err = fmt.Errorf("attribute \"%s\" is not a number", name)
		}
		return n, err
	}
	return n, fmt.Errorf("attribute \"%s\" does not exists", name)
}

func (obj *Object) getContextAttributeNumber(name string) (n json.Number, err error) {
	if attrValue, attrExists := obj.ContextAttributes[name]; attrExists {
		n, isNumber := attrValue.(json.Number)
		if !isNumber {
			err = fmt.Errorf("context attribute \"%s\" is not a number", name)
		}
		return n, err
	}
	return n, fmt.Errorf("context attribute \"%s\" does not exists", name)
}

// GetInt64 returns an attribute as an int64. It returns the attribute's
// value or an error if the attribute doesn't exist or is not a number.
func (obj *Object) GetInt64(name string) (int64, error) {
	n, err := obj.getAttributeNumber(name)
	if err == nil {
		return n.Int64()
	}
	return 0, err
}

// MustGetInt64 is like GetInt64, but it panic in case of error.
func (obj *Object) MustGetInt64(name string) int64 {
	result, err := obj.GetInt64(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetFloat64 returns an attribute as a float64. It returns the attribute's
// value or an error if the attribute doesn't exist or is not a number.
func (obj *Object) GetFloat64(name string) (float64, error) {
	n, err := obj.getAttributeNumber(name)
	if err == nil {
		return n.Float64()
	}
	return 0, err
}

// MustGetFloat64 is like GetFloat64, but it panic in case of error.
func (obj *Object) MustGetFloat64(name string) float64 {
	result, err := obj.GetFloat64(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetString returns an attribute as a string. It returns the attribute's
// valueor an error if the attribute doesn't exist or is not a string.
func (obj *Object) GetString(name string) (s string, err error) {
	if attrValue, attrExists := obj.Attributes[name]; attrExists {
		s, isString := attrValue.(string)
		if !isString {
			err = fmt.Errorf("attribute \"%s\" is not a string", name)
		}
		return s, err
	}
	return "", fmt.Errorf("attribute \"%s\" does not exists", name)
}

// MustGetString is like GetString, but it panic in case of error.
func (obj *Object) MustGetString(name string) string {
	result, err := obj.GetString(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetTime returns an attribute as a time. It returns the attribute's
// value or an error if the attribute doesn't exist or is not a time.
func (obj *Object) GetTime(name string) (t time.Time, err error) {
	n, err := obj.getAttributeNumber(name)
	if err == nil {
		i, err := n.Int64()
		return time.Unix(i, 0), err
	}
	return time.Unix(0, 0), err
}

// MustGetTime is like GetTime, but it panic in case of error.
func (obj *Object) MustGetTime(name string) time.Time {
	result, err := obj.GetTime(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetBool returns an attribute as a boolean. It returns the attribute's
// value or an error if the attribute doesn't exist or is not a boolean.
func (obj *Object) GetBool(name string) (b bool, err error) {
	if attrValue, attrExists := obj.Attributes[name]; attrExists {
		b, isBool := attrValue.(bool)
		if !isBool {
			err = fmt.Errorf("context attribute \"%s\" is not a bool", name)
		}
		return b, err
	}
	return false, fmt.Errorf("context attribute \"%s\" does not exists", name)
}

// MustGetBool is like GetTime, but it panic in case of error.
func (obj *Object) MustGetBool(name string) bool {
	result, err := obj.GetBool(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetContextInt64 returns a context attribute as an int64. It returns the
// attribute's value or an error if the attribute doesn't exist or is not a
// number.
func (obj *Object) GetContextInt64(name string) (int64, error) {
	n, err := obj.getContextAttributeNumber(name)
	if err == nil {
		return n.Int64()
	}
	return 0, err
}

// GetContextFloat64 returns a context attribute as an float64. It returns the
// attribute's value or an error if the attribute doesn't exist or is not a
// number.
func (obj *Object) GetContextFloat64(name string) (float64, error) {
	n, err := obj.getContextAttributeNumber(name)
	if err == nil {
		return n.Float64()
	}
	return 0, err
}

// GetContextString returns a context attribute as a string. It returns the
// attribute's value or an error if the attribute doesn't exist or is not a
// string.
func (obj *Object) GetContextString(name string) (s string, err error) {
	if attrValue, attrExists := obj.ContextAttributes[name]; attrExists {
		s, isString := attrValue.(string)
		if !isString {
			err = fmt.Errorf("context attribute \"%s\" is not a string", name)
		}
		return s, err
	}
	return "", fmt.Errorf("context attribute \"%s\" does not exists", name)
}

// GetContextBool returns a context attribute as a bool. It returns the
// attribute's value or an error if the attribute doesn't exist or is not a
// bool.
func (obj *Object) GetContextBool(name string) (b bool, err error) {
	if attrValue, attrExists := obj.ContextAttributes[name]; attrExists {
		b, isBool := attrValue.(bool)
		if !isBool {
			err = fmt.Errorf("context attribute \"%s\" is not a bool", name)
		}
		return b, err
	}
	return false, fmt.Errorf("context attribute \"%s\" does not exists", name)
}
<commit_msg>Unmarshal relationships correctly when it has a null value.<commit_after>// Copyright © 2017 The vt-go authors. All Rights Reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package vt

import (
	"bytes"
	"encoding/json"
	"fmt"
	"time"
)

// ObjectDescriptor is a pair (ID, type) describing a VirusTotal API object.
type ObjectDescriptor struct {
	ID                string                 `json:"id,omitempty"`
	Type              string                 `json:"type,omitempty"`
	ContextAttributes map[string]interface{} `json:"context_attributes,omitempty"`
}

// Object represents a VirusTotal API object.
type Object struct {
	ID                string                   `json:"id,omitempty"`
	Type              string                   `json:"type,omitempty"`
	Attributes        map[string]interface{}   `json:"attributes,omitempty"`
	ContextAttributes map[string]interface{}   `json:"context_attributes,omitempty"`
	Relationships     map[string]*Relationship `json:"relationships,omitempty"`
	Links             Links                    `json:"links,omitempty"`
}

// Links contains links related to an API object.
type Links struct {
	Self string `json:"self"`
	Next string `json:"next"`
}

// Relationship contains information about a related API object.
type Relationship struct {
	Data  json.RawMessage `json:"data,omitempty"`
	Links Links           `json:"links,omitempty"`

	// IsOneToOne is true if this is a one-to-one relationshio and False if
	// otherwise. If true RelatedObjects contains one object at most.
	IsOneToOne     bool
	RelatedObjects []ObjectDescriptor
}

// NewObject creates a new object.
func NewObject() *Object {
	return &Object{Attributes: make(map[string]interface{})}
}

// UnmarshalJSON unmarshals a VirusTotal API object from data.
func (obj *Object) UnmarshalJSON(data []byte) error {

	type object Object

	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.UseNumber()

	o := object{}
	if err := decoder.Decode(&o); err != nil {
		return err
	}

	obj.Type = o.Type
	obj.ID = o.ID
	obj.Attributes = o.Attributes
	obj.ContextAttributes = o.ContextAttributes
	obj.Relationships = o.Relationships

	for _, v := range obj.Relationships {
		var o ObjectDescriptor
		// Try unmarshalling as an ObjectDescriptor first, if it fails this is
		// a one-to-many relationship, so we try unmarshalling as an array of
		// ObjectDescriptor.
		if err := json.Unmarshal(v.Data, &o); err == nil {
			v.IsOneToOne = true
			// If the value is null the ObjectDescriptor will have an empty
			// ID and Type.
			if o.ID == "" && o.Type == "" {
				v.RelatedObjects = nil
			} else {
				v.RelatedObjects = append(v.RelatedObjects, o)
			}
		} else {
			if err := json.Unmarshal(v.Data, &v.RelatedObjects); err != nil {
				return err
			}
		}
	}

	return nil
}

func (obj *Object) getAttributeNumber(name string) (n json.Number, err error) {
	if attrValue, attrExists := obj.Attributes[name]; attrExists {
		n, isNumber := attrValue.(json.Number)
		if !isNumber {
			err = fmt.Errorf("attribute \"%s\" is not a number", name)
		}
		return n, err
	}
	return n, fmt.Errorf("attribute \"%s\" does not exists", name)
}

func (obj *Object) getContextAttributeNumber(name string) (n json.Number, err error) {
	if attrValue, attrExists := obj.ContextAttributes[name]; attrExists {
		n, isNumber := attrValue.(json.Number)
		if !isNumber {
			err = fmt.Errorf("context attribute \"%s\" is not a number", name)
		}
		return n, err
	}
	return n, fmt.Errorf("context attribute \"%s\" does not exists", name)
}

// GetInt64 returns an attribute as an int64. It returns the attribute's
// value or an error if the attribute doesn't exist or is not a number.
func (obj *Object) GetInt64(name string) (int64, error) {
	n, err := obj.getAttributeNumber(name)
	if err == nil {
		return n.Int64()
	}
	return 0, err
}

// MustGetInt64 is like GetInt64, but it panic in case of error.
func (obj *Object) MustGetInt64(name string) int64 {
	result, err := obj.GetInt64(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetFloat64 returns an attribute as a float64. It returns the attribute's
// value or an error if the attribute doesn't exist or is not a number.
func (obj *Object) GetFloat64(name string) (float64, error) {
	n, err := obj.getAttributeNumber(name)
	if err == nil {
		return n.Float64()
	}
	return 0, err
}

// MustGetFloat64 is like GetFloat64, but it panic in case of error.
func (obj *Object) MustGetFloat64(name string) float64 {
	result, err := obj.GetFloat64(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetString returns an attribute as a string. It returns the attribute's
// valueor an error if the attribute doesn't exist or is not a string.
func (obj *Object) GetString(name string) (s string, err error) {
	if attrValue, attrExists := obj.Attributes[name]; attrExists {
		s, isString := attrValue.(string)
		if !isString {
			err = fmt.Errorf("attribute \"%s\" is not a string", name)
		}
		return s, err
	}
	return "", fmt.Errorf("attribute \"%s\" does not exists", name)
}

// MustGetString is like GetString, but it panic in case of error.
func (obj *Object) MustGetString(name string) string {
	result, err := obj.GetString(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetTime returns an attribute as a time. It returns the attribute's
// value or an error if the attribute doesn't exist or is not a time.
func (obj *Object) GetTime(name string) (t time.Time, err error) {
	n, err := obj.getAttributeNumber(name)
	if err == nil {
		i, err := n.Int64()
		return time.Unix(i, 0), err
	}
	return time.Unix(0, 0), err
}

// MustGetTime is like GetTime, but it panic in case of error.
func (obj *Object) MustGetTime(name string) time.Time {
	result, err := obj.GetTime(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetBool returns an attribute as a boolean. It returns the attribute's
// value or an error if the attribute doesn't exist or is not a boolean.
func (obj *Object) GetBool(name string) (b bool, err error) {
	if attrValue, attrExists := obj.Attributes[name]; attrExists {
		b, isBool := attrValue.(bool)
		if !isBool {
			err = fmt.Errorf("context attribute \"%s\" is not a bool", name)
		}
		return b, err
	}
	return false, fmt.Errorf("context attribute \"%s\" does not exists", name)
}

// MustGetBool is like GetTime, but it panic in case of error.
func (obj *Object) MustGetBool(name string) bool {
	result, err := obj.GetBool(name)
	if err != nil {
		panic(err)
	}
	return result
}

// GetContextInt64 returns a context attribute as an int64. It returns the
// attribute's value or an error if the attribute doesn't exist or is not a
// number.
func (obj *Object) GetContextInt64(name string) (int64, error) {
	n, err := obj.getContextAttributeNumber(name)
	if err == nil {
		return n.Int64()
	}
	return 0, err
}

// GetContextFloat64 returns a context attribute as an float64. It returns the
// attribute's value or an error if the attribute doesn't exist or is not a
// number.
func (obj *Object) GetContextFloat64(name string) (float64, error) {
	n, err := obj.getContextAttributeNumber(name)
	if err == nil {
		return n.Float64()
	}
	return 0, err
}

// GetContextString returns a context attribute as a string. It returns the
// attribute's value or an error if the attribute doesn't exist or is not a
// string.
func (obj *Object) GetContextString(name string) (s string, err error) {
	if attrValue, attrExists := obj.ContextAttributes[name]; attrExists {
		s, isString := attrValue.(string)
		if !isString {
			err = fmt.Errorf("context attribute \"%s\" is not a string", name)
		}
		return s, err
	}
	return "", fmt.Errorf("context attribute \"%s\" does not exists", name)
}

// GetContextBool returns a context attribute as a bool. It returns the
// attribute's value or an error if the attribute doesn't exist or is not a
// bool.
func (obj *Object) GetContextBool(name string) (b bool, err error) {
	if attrValue, attrExists := obj.ContextAttributes[name]; attrExists {
		b, isBool := attrValue.(bool)
		if !isBool {
			err = fmt.Errorf("context attribute \"%s\" is not a bool", name)
		}
		return b, err
	}
	return false, fmt.Errorf("context attribute \"%s\" does not exists", name)
}
<|endoftext|>


<commit_before><commit_msg>Speed up duplicate checking using a hash<commit_after><|endoftext|>


<commit_before><commit_msg>Remove traces of go2idl<commit_after><|endoftext|>


<commit_before><commit_msg>Format mailer `From` field to include name<commit_after><|endoftext|>


<commit_before>package sparta

import (
	"errors"
	"fmt"
	"github.com/Sirupsen/logrus"
	"reflect"
	"strings"
)

// Common IAM Policy Actions for Lambda push-source configuration management.
// The configuration is handled by CustomResources inserted into the generated
// CloudFormation template.
var PushSourceConfigurationActions = map[string][]string{
	"s3.amazonaws.com": {"s3:GetBucketNotificationConfiguration",
		"s3:PutBucketNotificationConfiguration"},
	"sns.amazonaws.com": {"sns:ConfirmSubscription",
		"sns:GetTopicAttributes",
		"sns:Subscribe",
		"sns:Unsubscribe"},
}

func ensureConfiguratorLambdaResource(awsPrincipalName string, sourceArn string, resources ArbitraryJSONObject, logger *logrus.Logger) (string, error) {
	// AWS service basename
	awsServiceName := strings.ToUpper(strings.SplitN(awsPrincipalName, ".", 2)[0])

	//////////////////////////////////////////////////////////////////////////////
	// IAM Role definition
	// TODO - Check sourceArn for equivalence
	iamResourceName, err := ensureIAMRoleResource(awsPrincipalName, sourceArn, resources, logger)
	if nil != err {
		return "", err
	}

	iamRoleRef := ArbitraryJSONObject{
		"Fn::GetAtt": []string{iamResourceName, "Arn"},
	}
	// Custom handler resource for this service type
	subscriberHandlerName := fmt.Sprintf("%sSubscriber", awsServiceName)
	_, exists := resources[subscriberHandlerName]
	if !exists {
		logger.Info("Creating Subscription Lambda Resource for AWS service: ", awsServiceName)

		//////////////////////////////////////////////////////////////////////////////
		// Custom Resource Lambda Handler
		// NOTE: This path depends on `go generate` already having processed the provision
		// directory with the https://github.com/tdewolff/minify/tree/master/cmd/minify contents
		scriptHandlerPath := fmt.Sprintf("/resources/provision/%s.min.js", strings.ToLower(awsServiceName))
		logger.Debug("Lambda Source: ", scriptHandlerPath)

		customResourceHandlerDef := ArbitraryJSONObject{
			"Type": "AWS::Lambda::Function",
			"Properties": ArbitraryJSONObject{
				"Code": ArbitraryJSONObject{
					"ZipFile": FSMustString(false, scriptHandlerPath),
				},
				"Role":    iamRoleRef,
				"Handler": "index.handler",
				"Runtime": "nodejs",
				"Timeout": "30",
			},
		}
		resources[subscriberHandlerName] = customResourceHandlerDef
	}
	return subscriberHandlerName, nil
}

func ensureIAMRoleResource(awsPrincipalName string, sourceArn string, resources ArbitraryJSONObject, logger *logrus.Logger) (string, error) {
	principalActions, exists := PushSourceConfigurationActions[awsPrincipalName]
	if !exists {
		return "", errors.New("Unsupported principal for IAM role creation: " + awsPrincipalName)
	}

	// First determine if there is one provisioned...
	var iamRoleResourceNames []string
	for eachName, eachResource := range resources {
		logger.Debug("Checking IAM Policy equality: ", eachName)
		if eachResource.(ArbitraryJSONObject)["Type"] == "AWS::IAM::Role" {
			properties := eachResource.(ArbitraryJSONObject)["Properties"]
			policies := properties.(ArbitraryJSONObject)["Policies"]
			for _, eachPolicyEntry := range policies.([]ArbitraryJSONObject) {
				policyDocument := eachPolicyEntry["PolicyDocument"]
				statements := policyDocument.(ArbitraryJSONObject)["Statement"]
				for _, eachStatement := range statements.([]ArbitraryJSONObject) {
					if eachStatement["Resource"] == sourceArn &&
						reflect.DeepEqual(eachStatement["Action"], principalActions) {
						iamRoleResourceNames = append(iamRoleResourceNames, eachName)
					}
				}
			}
		}
	}
	logger.WithFields(logrus.Fields{
		"MatchingIAMRoleNames": iamRoleResourceNames,
		"PrincipalActions":     principalActions,
		"Principal":            awsPrincipalName,
	}).Debug("Ensuring IAM Role results")

	if len(iamRoleResourceNames) > 1 {
		return "", errors.New("More than 1 IAM Role found for entry: " + awsPrincipalName)
	} else if len(iamRoleResourceNames) == 1 {
		logger.Debug("Using prexisting IAM Role: " + iamRoleResourceNames[0])
		return iamRoleResourceNames[0], nil
	} else {
		// Provision a new one and add it...
		newIAMRoleResourceName := CloudFormationResourceName("IAMRole")
		logger.Debug("Inserting new IAM Role: ", newIAMRoleResourceName)

		statements := CommonIAMStatements
		logger.Info("IAMRole Actions: ", principalActions)

		statements = append(statements, ArbitraryJSONObject{
			"Effect":   "Allow",
			"Action":   principalActions,
			"Resource": sourceArn,
		})

		iamPolicy := ArbitraryJSONObject{"Type": "AWS::IAM::Role",
			"Properties": ArbitraryJSONObject{
				"AssumeRolePolicyDocument": AssumePolicyDocument,
				"Policies": []ArbitraryJSONObject{
					{
						"PolicyName": "configurator",
						"PolicyDocument": ArbitraryJSONObject{
							"Version":   "2012-10-17",
							"Statement": statements,
						},
					},
				},
			},
		}
		resources[newIAMRoleResourceName] = iamPolicy
		return newIAMRoleResourceName, nil
	}
}
<commit_msg>Fix IAM role creation to use proper crews & improve resource name<commit_after>package sparta

import (
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/Sirupsen/logrus"
	"strings"
)

const salt = "213EA743-A98F-499D-8FEF-B87015FE13E7"

// Common IAM Policy Actions for Lambda push-source configuration management.
// The configuration is handled by CustomResources inserted into the generated
// CloudFormation template.
var PushSourceConfigurationActions = map[string][]string{
	"s3.amazonaws.com": {"s3:GetBucketLocation",
		"s3:GetBucketNotification",
		"s3:PutBucketNotification",
		"s3:GetBucketNotificationConfiguration",
		"s3:PutBucketNotificationConfiguration"},
	"sns.amazonaws.com": {"sns:ConfirmSubscription",
		"sns:GetTopicAttributes",
		"sns:Subscribe",
		"sns:Unsubscribe"},
}

func ensureConfiguratorLambdaResource(awsPrincipalName string, sourceArn string, resources ArbitraryJSONObject, logger *logrus.Logger) (string, error) {
	// AWS service basename
	awsServiceName := strings.ToUpper(strings.SplitN(awsPrincipalName, ".", 2)[0])

	//////////////////////////////////////////////////////////////////////////////
	// IAM Role definition
	// TODO - Check sourceArn for equivalence
	iamResourceName, err := ensureIAMRoleResource(awsServiceName, awsPrincipalName, sourceArn, resources, logger)
	if nil != err {
		return "", err
	}

	iamRoleRef := ArbitraryJSONObject{
		"Fn::GetAtt": []string{iamResourceName, "Arn"},
	}
	// Custom handler resource for this service type
	subscriberHandlerName := fmt.Sprintf("%sSubscriber", awsServiceName)
	_, exists := resources[subscriberHandlerName]
	if !exists {
		logger.Info("Creating Subscription Lambda Resource for AWS service: ", awsServiceName)

		//////////////////////////////////////////////////////////////////////////////
		// Custom Resource Lambda Handler
		// NOTE: This path depends on `go generate` already having processed the provision
		// directory with the https://github.com/tdewolff/minify/tree/master/cmd/minify contents
		scriptHandlerPath := fmt.Sprintf("/resources/provision/%s.min.js", strings.ToLower(awsServiceName))
		logger.Debug("Lambda Source: ", scriptHandlerPath)

		customResourceHandlerDef := ArbitraryJSONObject{
			"Type": "AWS::Lambda::Function",
			"Properties": ArbitraryJSONObject{
				"Code": ArbitraryJSONObject{
					"ZipFile": FSMustString(false, scriptHandlerPath),
				},
				"Role":    iamRoleRef,
				"Handler": "index.handler",
				"Runtime": "nodejs",
				"Timeout": "30",
			},
		}
		resources[subscriberHandlerName] = customResourceHandlerDef
	}
	return subscriberHandlerName, nil
}

func ensureIAMRoleResource(awsServiceName string, awsPrincipalName string, sourceArn string, resources ArbitraryJSONObject, logger *logrus.Logger) (string, error) {
	principalActions, exists := PushSourceConfigurationActions[awsPrincipalName]
	if !exists {
		return "", errors.New("Unsupported principal for IAM role creation: " + awsPrincipalName)
	}

	hash := sha1.New()
	hash.Write([]byte(fmt.Sprintf("%s%s", awsPrincipalName, salt)))
	roleName := fmt.Sprintf("%sConfigIAMRole%s", awsServiceName, hex.EncodeToString(hash.Sum(nil)))

	logger.WithFields(logrus.Fields{
		"PrincipalActions": principalActions,
		"Principal":        awsPrincipalName,
	}).Debug("Ensuring IAM Role results")

	_, exists = resources[roleName]

	// If it exists, make sure these permissions are enabled on it...
	if exists {
		logger.Debug("Using prexisting IAM Role: " + roleName)
		return roleName, nil
	} else {
		// Provision a new one and add it...
		newIAMRoleResourceName := CloudFormationResourceName("IAMRole")
		logger.Debug("Inserting new IAM Role: ", newIAMRoleResourceName)

		statements := CommonIAMStatements
		logger.Info("IAMRole Actions: ", principalActions)

		statements = append(statements, ArbitraryJSONObject{
			"Effect":   "Allow",
			"Action":   principalActions,
			"Resource": sourceArn,
		})

		iamPolicy := ArbitraryJSONObject{"Type": "AWS::IAM::Role",
			"Properties": ArbitraryJSONObject{
				"AssumeRolePolicyDocument": AssumePolicyDocument,
				"Policies": []ArbitraryJSONObject{
					{
						"PolicyName": fmt.Sprintf("%sConfigurator%s", awsServiceName, CloudFormationResourceName("")),
						"PolicyDocument": ArbitraryJSONObject{
							"Version":   "2012-10-17",
							"Statement": statements,
						},
					},
				},
			},
		}
		resources[roleName] = iamPolicy
		return roleName, nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 DSR Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package x509

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"strings"

	"github.com/zigbee-alliance/distributed-compliance-ledger/x/pki/types"
)

type X509Certificate struct {
	Issuer         string
	SerialNumber   string
	Subject        string
	SubjectKeyID   string
	AuthorityKeyID string
	Certificate    *x509.Certificate
}

func DecodeX509Certificate(pemCertificate string) (*X509Certificate, error) {
	block, _ := pem.Decode([]byte(pemCertificate))
	if block == nil {
		return nil, types.NewErrInvalidCertificate("Could not decode pem certificate")
	}

	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, types.NewErrInvalidCertificate(fmt.Sprintf("Could not parse certificate: %v", err.Error()))
	}

	certificate := X509Certificate{
		Issuer:         cert.Issuer.String(),
		SerialNumber:   cert.SerialNumber.String(),
		Subject:        cert.Subject.String(),
		SubjectKeyID:   BytesToHex(cert.SubjectKeyId),
		AuthorityKeyID: BytesToHex(cert.AuthorityKeyId),
		Certificate:    cert,
	}

	return &certificate, nil
}

func BytesToHex(bytes []byte) string {
	if bytes == nil {
		return ""
	}

	bytesHex := make([]string, len(bytes))
	for i, b := range bytes {
		bytesHex[i] = fmt.Sprintf("%X", b)
	}

	return strings.Join(bytesHex, ":")
}

func (c X509Certificate) Verify(parent *X509Certificate) error {
	roots := x509.NewCertPool()
	roots.AddCert(parent.Certificate)

	opts := x509.VerifyOptions{Roots: roots}

	if _, err := c.Certificate.Verify(opts); err != nil {
		return types.NewErrInvalidCertificate(fmt.Sprintf("Certificate verification failed. Error: %v", err))
	}

	return nil
}

func (c X509Certificate) IsSelfSigned() bool {
	if len(c.AuthorityKeyID) > 0 {
		return c.Issuer == c.Subject && c.AuthorityKeyID == c.SubjectKeyID
	} else {
		return c.Issuer == c.Subject
	}
}
<commit_msg>More human friendly way to display Matter OIDs in the certificates Fixes #270<commit_after>// Copyright 2020 DSR Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package x509

import (
	"bytes"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"strings"

	"github.com/zigbee-alliance/distributed-compliance-ledger/x/pki/types"
)

type X509Certificate struct {
	Issuer         string
	SerialNumber   string
	Subject        string
	SubjectKeyID   string
	AuthorityKeyID string
	Certificate    *x509.Certificate
}

func DecodeX509Certificate(pemCertificate string) (*X509Certificate, error) {
	block, _ := pem.Decode([]byte(pemCertificate))
	if block == nil {
		return nil, types.NewErrInvalidCertificate("Could not decode pem certificate")
	}

	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, types.NewErrInvalidCertificate(fmt.Sprintf("Could not parse certificate: %v", err.Error()))
	}

	certificate := X509Certificate{
		Issuer:         TranslateToVidOrPid([]byte(cert.Issuer.String())),
		SerialNumber:   cert.SerialNumber.String(),
		Subject:        TranslateToVidOrPid([]byte(cert.Subject.String())),
		SubjectKeyID:   BytesToHex(cert.SubjectKeyId),
		AuthorityKeyID: BytesToHex(cert.AuthorityKeyId),
		Certificate:    cert,
	}

	return &certificate, nil
}

func TranslateToVidOrPid(subject []byte) string {
	oldVID := []byte("1.3.6.1.4.1.37244.2.1")
	oldPID := []byte("1.3.6.1.4.1.37244.2.2")

	var vid, vidValue, pid, pidValue []byte

	if i := bytes.Index(subject, oldVID); i >= 0 {
		// get vid value from subject
		vidValue = subject[i+len(oldVID)+2 : len(string(subject))]
		// get last 8 numbers from vidValue
		vidValue = vidValue[len(vidValue)-8:]

		decoded, _ := hex.DecodeString(string(vidValue))

		vid = []byte(("vid=0x") + string(decoded))

		subject = subject[0:i]
		subject = []byte(string(subject) + string(vid))
	} else if i := bytes.Index(subject, oldPID); i >= 0 {
		// get pid value from subject
		pidValue = subject[i+len(oldPID)+2 : len(string(subject))]
		// get last 8 number from pidValue
		pidValue = pidValue[len(pidValue)-8:]

		decoded, _ := hex.DecodeString(string(pidValue))

		pid = []byte(("pid=0x") + string(decoded))

		subject = subject[0:i]
		subject = []byte(string(subject) + string(pid))
	}

	return string(subject)
}

func BytesToHex(bytes []byte) string {
	if bytes == nil {
		return ""
	}

	bytesHex := make([]string, len(bytes))
	for i, b := range bytes {
		bytesHex[i] = fmt.Sprintf("%X", b)
	}

	return strings.Join(bytesHex, ":")
}

func (c X509Certificate) Verify(parent *X509Certificate) error {
	roots := x509.NewCertPool()
	roots.AddCert(parent.Certificate)

	opts := x509.VerifyOptions{Roots: roots}

	if _, err := c.Certificate.Verify(opts); err != nil {
		return types.NewErrInvalidCertificate(fmt.Sprintf("Certificate verification failed. Error: %v", err))
	}

	return nil
}

func (c X509Certificate) IsSelfSigned() bool {
	if len(c.AuthorityKeyID) > 0 {
		return c.Issuer == c.Subject && c.AuthorityKeyID == c.SubjectKeyID
	} else {
		return c.Issuer == c.Subject
	}
}
<|endoftext|>