<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"
)

func countFile() string {
	return os.ExpandEnv("${HOME}/.traktor-charts.count")
}
func requiresUpdate(count int) bool {
	shouldUpdate := false

	oldCount, fileErr := ioutil.ReadFile(countFile())
	if fileErr != nil {
		shouldUpdate = true
	} else {
		oldFileCount, _ := strconv.Atoi(string(oldCount))
		fmt.Printf("Found %d old entries\n", oldFileCount)
		if count != oldFileCount {
			shouldUpdate = true
		}
	}
	return shouldUpdate
}

func main() {
	fmt.Println("NI directories:", traktorDir(""))

	historyPaths, _ := traktorHistoryPaths(traktorDir(""))
	archiveFiles, _ := traktorArchiveFiles(historyPaths)

	db, err := initializeDB("traktor-charts.db")
	if err != true {
		fmt.Println("Error initializing db", err)
	}

	fileCount := 0
	for _, fileName := range archiveFiles {
		entries, _ := traktorParseFile(fileName)
		for _, entry := range entries.TraktorXMLEntryList {
			insertEntry(db, entries, entry)
		}
		fileCount++
	}
	fmt.Println("Found", fileCount, "archive files")

	jsonBytes := getExportData(db)
	db.Close()

	if requiresUpdate(fileCount) {
		if httpPostResults(jsonBytes) {
			_ = ioutil.WriteFile(countFile(), []byte(strconv.Itoa(fileCount)), 0600)
		} else {
			fmt.Println("Invalid credentials to post.")
			os.Exit(3)
		}
	} else {
		fmt.Println("No new traktor archive files found")
		os.Exit(1)
	}

}

func httpPostResults(traktorBody []byte) bool {
	url := "https://djcharts.io/api/import"
	fmt.Println("URL:>", url)

	token, _ := ioutil.ReadFile(os.ExpandEnv("${HOME}/.traktor-charts"))
	basicAuthToken := strings.TrimSuffix(string(token), "\n")

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(traktorBody))
	req.Header.Set("Content-Type", "application/json")
	req.SetBasicAuth("X", basicAuthToken)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	fmt.Println("Response Status:", resp.Status)
	body, _ := ioutil.ReadAll(resp.Body)
	fmt.Println("Response Body:", string(body))
	return resp.Status == "201 Created"
}
<commit_msg>add a simple pulse task to know if the app needs to be configured or not<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"
)

func countFile() string {
	return os.ExpandEnv("${HOME}/.traktor-charts.count")
}

func requiresUpdate(count int) bool {
	shouldUpdate := false

	oldCount, fileErr := ioutil.ReadFile(countFile())
	if fileErr != nil {
		shouldUpdate = true
	} else {
		oldFileCount, _ := strconv.Atoi(string(oldCount))
		fmt.Printf("Found %d old entries\n", oldFileCount)
		if count != oldFileCount {
			shouldUpdate = true
		}
	}
	return shouldUpdate
}

func main() {
	fmt.Println("NI directories:", traktorDir(""))

	historyPaths, _ := traktorHistoryPaths(traktorDir(""))
	archiveFiles, _ := traktorArchiveFiles(historyPaths)

	db, err := initializeDB("traktor-charts.db")
	if err != true {
		fmt.Println("Error initializing db", err)
	}

	fileCount := 0
	for _, fileName := range archiveFiles {
		entries, _ := traktorParseFile(fileName)
		for _, entry := range entries.TraktorXMLEntryList {
			insertEntry(db, entries, entry)
		}
		fileCount++
	}
	fmt.Println("Found", fileCount, "archive files")

	jsonBytes := getExportData(db)
	db.Close()

  if httpPulse() {
    if requiresUpdate(fileCount) {
      if httpPostResults(jsonBytes) {
        _ = ioutil.WriteFile(countFile(), []byte(strconv.Itoa(fileCount)), 0600)
      } else {
        fmt.Println("Invalid credentials to post.")
        os.Exit(3)
      }
    } else {
      fmt.Println("No new traktor archive files found")
      os.Exit(1)
    }
  } else {
    fmt.Println("Unable to pulse. Your app is likely misconfigured")
    os.Exit(3)
  }

}

func httpPostResults(traktorBody []byte) bool {
	url := "https://djcharts.io/api/import"
	fmt.Println("URL:>", url)

	token, _ := ioutil.ReadFile(os.ExpandEnv("${HOME}/.traktor-charts"))
	basicAuthToken := strings.TrimSuffix(string(token), "\n")

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(traktorBody))
	req.Header.Set("Content-Type", "application/json")
	req.SetBasicAuth("X", basicAuthToken)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	fmt.Println("Response Status:", resp.Status)
	body, _ := ioutil.ReadAll(resp.Body)
	fmt.Println("Response Body:", string(body))
	return resp.Status == "201 Created"
}

func httpPulse() bool {
	url := "https://djcharts.io/api/pulse"
	fmt.Println("URL:>", url)

	token, _ := ioutil.ReadFile(os.ExpandEnv("${HOME}/.traktor-charts"))
	basicAuthToken := strings.TrimSuffix(string(token), "\n")

	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("Content-Type", "application/json")
	req.SetBasicAuth("X", basicAuthToken)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	fmt.Println("Response Status:", resp.Status)
	body, _ := ioutil.ReadAll(resp.Body)
	fmt.Println("Response Body:", string(body))
	return resp.Status == "200 OK"
}
<|endoftext|>


<commit_before>package trollmode

import (
	"math"
	"math/rand"
	"socialapi/models"
	"socialapi/rest"
	"socialapi/workers/common/runner"
	"socialapi/workers/common/tests"
	"strconv"
	"testing"
	"time"

	"github.com/koding/bongo"
	. "github.com/smartystreets/goconvey/convey"
	"labix.org/v2/mgo/bson"
)

func TestMarkedAsTroll(t *testing.T) {
	r := runner.New("TrollMode-Test")
	err := r.Init()
	if err != nil {
		panic(err)
	}

	defer r.Close()
	// disable logs
	// r.Log.SetLevel(logging.CRITICAL)

	Convey("given a controller", t, func() {

		// cretae admin user
		adminUser := models.NewAccount()
		adminUser.OldId = bson.NewObjectId().Hex()
		adminUser, err = rest.CreateAccount(adminUser)
		tests.ResultedWithNoErrorCheck(adminUser, err)

		// create troll user
		trollUser := models.NewAccount()
		trollUser.OldId = bson.NewObjectId().Hex()
		trollUser, err := rest.CreateAccount(trollUser)
		tests.ResultedWithNoErrorCheck(trollUser, err)
		trollUser.IsTroll = true

		// create normal user
		normalUser := models.NewAccount()
		normalUser.OldId = bson.NewObjectId().Hex()
		normalUser, err = rest.CreateAccount(normalUser)
		tests.ResultedWithNoErrorCheck(normalUser, err)

		// create groupName
		rand.Seed(time.Now().UnixNano())
		groupName := "testgroup" + strconv.FormatInt(rand.Int63(), 10)
		groupChannel, err := rest.CreateChannelByGroupNameAndType(
			adminUser.Id,
			groupName,
			models.Channel_TYPE_GROUP,
		)

		controller := NewController(r.Log)

		Convey("err should be nil", func() {
			So(err, ShouldBeNil)
		})

		Convey("controller should be set", func() {
			So(controller, ShouldNotBeNil)
		})

		Convey("should return nil when given nil account", func() {
			So(controller.MarkedAsTroll(nil), ShouldBeNil)
		})

		Convey("should return nil when account id given 0", func() {
			So(controller.MarkedAsTroll(models.NewAccount()), ShouldBeNil)
		})

		Convey("non existing account should not give error", func() {
			a := models.NewAccount()
			a.Id = math.MaxInt64
			So(controller.MarkedAsTroll(a), ShouldBeNil)
		})

		Convey("non existing account should not give error", func() {
			a := models.NewAccount()
			a.Id = math.MaxInt64
			So(controller.MarkedAsTroll(a), ShouldBeNil)
		})

		/////////////////////////  marking all content ////////////////////////
		// mark channel
		Convey("private channels of a troll should be marked as exempt", func() {
			// fetch from api, because we need to test system from there
			privatemessageChannelId1, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId1, ShouldBeGreaterThan, 0)

			privatemessageChannelId2, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId2, ShouldBeGreaterThan, 0)

			So(controller.markChannels(trollUser), ShouldBeNil)

			// fetch channel from db
			c1 := models.NewChannel()
			err = c1.ById(privatemessageChannelId1)
			So(err, ShouldBeNil)
			So(c1.Id, ShouldEqual, privatemessageChannelId1)
			// check here
			So(c1.MetaBits.IsTroll(), ShouldBeTrue)

			// fetch channel from db
			c2 := models.NewChannel()
			err = c2.ById(privatemessageChannelId2)
			So(err, ShouldBeNil)
			So(c2.Id, ShouldEqual, privatemessageChannelId2)

			// check here
			So(c2.MetaBits.IsTroll(), ShouldBeTrue)
		})

		// mark channel
		Convey("public channels of a troll should not be marked as exempt", nil)

		// mark channel_participant
		Convey("participations of a troll should not be marked as exempt", func() {
			// fetch from api, because we need to test system from there
			privatemessageChannelId1, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId1, ShouldBeGreaterThan, 0)

			privatemessageChannelId2, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId2, ShouldBeGreaterThan, 0)

			So(controller.markParticipations(trollUser), ShouldBeNil)

			var participations []models.ChannelParticipant

			query := &bongo.Query{
				Selector: map[string]interface{}{
					"account_id": trollUser.Id,
				},
			}

			err = models.NewChannelParticipant().Some(&participations, query)
			So(err, ShouldBeNil)
			for _, participation := range participations {
				So(participation.MetaBits.IsTroll(), ShouldBeTrue)
			}
		})

		// mark channel_message_list
		Convey("massages that are in all channels that are created by troll, should be marked as exempt", func() {
			// post, err := rest.CreatePost(groupChannel.Id, trollUser.Id)
			// tests.ResultedWithNoErrorCheck(post, err)

			// post, err = rest.CreatePost(groupChannel.Id, trollUser.Id)
			// tests.ResultedWithNoErrorCheck(post, err)
		})

		// mark channel_message
		Convey("messages of a troll should be marked as exempt", func() {
			post1, err := rest.CreatePost(groupChannel.Id, trollUser.Id)
			tests.ResultedWithNoErrorCheck(post1, err)

			post2, err := rest.CreatePost(groupChannel.Id, trollUser.Id)
			tests.ResultedWithNoErrorCheck(post2, err)

			So(controller.markMessages(trollUser), ShouldBeNil)

			cm := models.NewChannelMessage()
			q := &bongo.Query{
				Selector: map[string]interface{}{
					"account_id": trollUser.Id,
				},
			}

			var messages []models.ChannelMessage
			err = cm.Some(&messages, q)
			So(err, ShouldBeNil)

			for _, message := range messages {
				So(message.MetaBits.IsTroll(), ShouldBeTrue)
			}
		})

		// mark interactions
		Convey("interactions of a troll should be marked as exempt", func() {
			post1, err := rest.CreatePost(groupChannel.Id, trollUser.Id)
			tests.ResultedWithNoErrorCheck(post1, err)

			err = rest.AddInteraction("like", post1.Id, trollUser.Id)
			So(err, ShouldBeNil)

			So(controller.markInteractions(trollUser), ShouldBeNil)

			cm := models.NewInteraction()
			q := &bongo.Query{
				Selector: map[string]interface{}{
					"account_id": trollUser.Id,
				},
			}

			var interactions []models.Interaction
			err = cm.Some(&interactions, q)
			So(err, ShouldBeNil)

			for _, interaction := range interactions {
				So(interaction.MetaBits.IsTroll(), ShouldBeTrue)
				So(interaction.MetaBits.IsTroll(), ShouldBeTrue)
			}
		})
		})

		// update channel data while creating
		Convey("when a troll creates a channel, meta_bits should be set", func() {
			privatemessageChannelId1, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId1, ShouldBeGreaterThan, 0)

			// fetch channel from db
			c1 := models.NewChannel()
			err = c1.ById(privatemessageChannelId1)
			So(err, ShouldBeNil)
			So(c1.Id, ShouldEqual, privatemessageChannelId1)

			So(c1.MetaBits.IsTroll(), ShouldBeTrue)
		})

	})
}

func createPrivateMessageChannel(accountId int64, groupName string) (int64, error) {
	// create first private channel
	cmc, err := rest.SendPrivateMessage(
		accountId,
		"this is a body for private message @sinan",
		groupName,
	)
	if err != nil {
		return 0, err
	}

	return cmc.Channel.Id, nil
}
<commit_msg>social: troll: test updating channel_participant data while creating<commit_after>package trollmode

import (
	"math"
	"math/rand"
	"socialapi/models"
	"socialapi/rest"
	"socialapi/workers/common/runner"
	"socialapi/workers/common/tests"
	"strconv"
	"testing"
	"time"

	"github.com/koding/bongo"
	. "github.com/smartystreets/goconvey/convey"
	"labix.org/v2/mgo/bson"
)

func TestMarkedAsTroll(t *testing.T) {
	r := runner.New("TrollMode-Test")
	err := r.Init()
	if err != nil {
		panic(err)
	}

	defer r.Close()
	// disable logs
	// r.Log.SetLevel(logging.CRITICAL)

	Convey("given a controller", t, func() {

		// cretae admin user
		adminUser := models.NewAccount()
		adminUser.OldId = bson.NewObjectId().Hex()
		adminUser, err = rest.CreateAccount(adminUser)
		tests.ResultedWithNoErrorCheck(adminUser, err)

		// create troll user
		trollUser := models.NewAccount()
		trollUser.OldId = bson.NewObjectId().Hex()
		trollUser, err := rest.CreateAccount(trollUser)
		tests.ResultedWithNoErrorCheck(trollUser, err)
		trollUser.IsTroll = true

		// create normal user
		normalUser := models.NewAccount()
		normalUser.OldId = bson.NewObjectId().Hex()
		normalUser, err = rest.CreateAccount(normalUser)
		tests.ResultedWithNoErrorCheck(normalUser, err)

		// create groupName
		rand.Seed(time.Now().UnixNano())
		groupName := "testgroup" + strconv.FormatInt(rand.Int63(), 10)
		groupChannel, err := rest.CreateChannelByGroupNameAndType(
			adminUser.Id,
			groupName,
			models.Channel_TYPE_GROUP,
		)

		controller := NewController(r.Log)

		Convey("err should be nil", func() {
			So(err, ShouldBeNil)
		})

		Convey("controller should be set", func() {
			So(controller, ShouldNotBeNil)
		})

		Convey("should return nil when given nil account", func() {
			So(controller.MarkedAsTroll(nil), ShouldBeNil)
		})

		Convey("should return nil when account id given 0", func() {
			So(controller.MarkedAsTroll(models.NewAccount()), ShouldBeNil)
		})

		Convey("non existing account should not give error", func() {
			a := models.NewAccount()
			a.Id = math.MaxInt64
			So(controller.MarkedAsTroll(a), ShouldBeNil)
		})

		Convey("non existing account should not give error", func() {
			a := models.NewAccount()
			a.Id = math.MaxInt64
			So(controller.MarkedAsTroll(a), ShouldBeNil)
		})

		/////////////////////////  marking all content ////////////////////////
		// mark channel
		Convey("private channels of a troll should be marked as exempt", func() {
			// fetch from api, because we need to test system from there
			privatemessageChannelId1, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId1, ShouldBeGreaterThan, 0)

			privatemessageChannelId2, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId2, ShouldBeGreaterThan, 0)

			So(controller.markChannels(trollUser), ShouldBeNil)

			// fetch channel from db
			c1 := models.NewChannel()
			err = c1.ById(privatemessageChannelId1)
			So(err, ShouldBeNil)
			So(c1.Id, ShouldEqual, privatemessageChannelId1)
			// check here
			So(c1.MetaBits.IsTroll(), ShouldBeTrue)

			// fetch channel from db
			c2 := models.NewChannel()
			err = c2.ById(privatemessageChannelId2)
			So(err, ShouldBeNil)
			So(c2.Id, ShouldEqual, privatemessageChannelId2)

			// check here
			So(c2.MetaBits.IsTroll(), ShouldBeTrue)
		})

		// mark channel
		Convey("public channels of a troll should not be marked as exempt", nil)

		// mark channel_participant
		Convey("participations of a troll should not be marked as exempt", func() {
			// fetch from api, because we need to test system from there
			privatemessageChannelId1, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId1, ShouldBeGreaterThan, 0)

			privatemessageChannelId2, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId2, ShouldBeGreaterThan, 0)

			So(controller.markParticipations(trollUser), ShouldBeNil)

			var participations []models.ChannelParticipant

			query := &bongo.Query{
				Selector: map[string]interface{}{
					"account_id": trollUser.Id,
				},
			}

			err = models.NewChannelParticipant().Some(&participations, query)
			So(err, ShouldBeNil)
			for _, participation := range participations {
				So(participation.MetaBits.IsTroll(), ShouldBeTrue)
			}
		})

		// mark channel_message_list
		Convey("massages that are in all channels that are created by troll, should be marked as exempt", func() {
			// post, err := rest.CreatePost(groupChannel.Id, trollUser.Id)
			// tests.ResultedWithNoErrorCheck(post, err)

			// post, err = rest.CreatePost(groupChannel.Id, trollUser.Id)
			// tests.ResultedWithNoErrorCheck(post, err)
		})

		// mark channel_message
		Convey("messages of a troll should be marked as exempt", func() {
			post1, err := rest.CreatePost(groupChannel.Id, trollUser.Id)
			tests.ResultedWithNoErrorCheck(post1, err)

			post2, err := rest.CreatePost(groupChannel.Id, trollUser.Id)
			tests.ResultedWithNoErrorCheck(post2, err)

			So(controller.markMessages(trollUser), ShouldBeNil)

			cm := models.NewChannelMessage()
			q := &bongo.Query{
				Selector: map[string]interface{}{
					"account_id": trollUser.Id,
				},
			}

			var messages []models.ChannelMessage
			err = cm.Some(&messages, q)
			So(err, ShouldBeNil)

			for _, message := range messages {
				So(message.MetaBits.IsTroll(), ShouldBeTrue)
			}
		})

		// mark interactions
		Convey("interactions of a troll should be marked as exempt", func() {
			post1, err := rest.CreatePost(groupChannel.Id, trollUser.Id)
			tests.ResultedWithNoErrorCheck(post1, err)

			err = rest.AddInteraction("like", post1.Id, trollUser.Id)
			So(err, ShouldBeNil)

			So(controller.markInteractions(trollUser), ShouldBeNil)

			cm := models.NewInteraction()
			q := &bongo.Query{
				Selector: map[string]interface{}{
					"account_id": trollUser.Id,
				},
			}

			var interactions []models.Interaction
			err = cm.Some(&interactions, q)
			So(err, ShouldBeNil)

			for _, interaction := range interactions {
				So(interaction.MetaBits.IsTroll(), ShouldBeTrue)
				So(interaction.MetaBits.IsTroll(), ShouldBeTrue)
			}
		})
		})

		// update channel data while creating
		Convey("when a troll creates a channel, meta_bits should be set", func() {
			privatemessageChannelId1, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId1, ShouldBeGreaterThan, 0)

			// fetch channel from db
			c1 := models.NewChannel()
			err = c1.ById(privatemessageChannelId1)
			So(err, ShouldBeNil)
			So(c1.Id, ShouldEqual, privatemessageChannelId1)

			So(c1.MetaBits.IsTroll(), ShouldBeTrue)
		})

		// update channel_participant data while creating
		Convey("when a troll is added to a channel as participant, meta_bits should be set", func() {
			privatemessageChannelId, err := createPrivateMessageChannel(trollUser.Id, groupName)
			So(err, ShouldBeNil)
			So(privatemessageChannelId, ShouldBeGreaterThan, 0)

			// fetch channel from db
			cp := models.NewChannelParticipant()
			cp.AccountId = trollUser.Id
			cp.ChannelId = privatemessageChannelId

			So(cp.FetchParticipant(), ShouldBeNil)
			So(cp.AccountId, ShouldEqual, trollUser.Id)

			So(cp.MetaBits.IsTroll(), ShouldBeTrue)
		})
	})
}

func createPrivateMessageChannel(accountId int64, groupName string) (int64, error) {
	// create first private channel
	cmc, err := rest.SendPrivateMessage(
		accountId,
		"this is a body for private message @sinan",
		groupName,
	)
	if err != nil {
		return 0, err
	}

	return cmc.Channel.Id, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io"
	"os"
	"regexp"
	"strings"
	"testing"
)

func run(expr string) string {
	old := os.Stdout
	r, w, _ := os.Pipe()
	os.Stdout = w

	execute_expression(expr)

	outc := make(chan string)
	go func() {
		var buf bytes.Buffer
		io.Copy(&buf, r)
		outc <- buf.String()
	}()

	w.Close()
	os.Stdout = old
	out := <-outc
	return strings.TrimSpace(unpad(out))
}

func expect(t *testing.T, output string, target string) {
	if target != output {
		t.Error("\nexpected:\n" + target + "\n\ngot:\n" + output)
	}
}

func unpad(result string) string {
	re := regexp.MustCompile("[ ]*\\n")
	return re.ReplaceAllString(result, "\n")
}

func TestEndswith(t *testing.T) {
	out := run("name in 'testdata/sub1' where name endswith '.txt'")
	expect(t, out, "test.txt\ntest2.txt")
}

func TestEndswithIgnorecase(t *testing.T) {
	out := run("name in 'testdata/sub1' where name ignorecase endswith '.TxT'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestStartswith(t *testing.T) {
	out := run("name in 'testdata/sub1' where name startswith 'test'")
	expect(t, out, "test.txt\ntest2.txt")
}

func TestStartswithIgnorecase(t *testing.T) {
	out := run("name in 'testdata/sub1' where name ignorecase startswith 'Test'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestIsfile(t *testing.T) {
	out := run("name in 'testdata/sub2' where isfile")
	expect(t, out, "a\nb\nc\nd")
}

func TestIsdir(t *testing.T) {
	out := run("name in 'testdata/sub2' where isdir")
	expect(t, out, "dir1/\ndir2/")
}

func TestNameEq(t *testing.T) {
	out := run("name in 'testdata/sub1' where name = 'test.txt'")
	expect(t, out, "test.txt")
}

func TestNameNeq(t *testing.T) {
	out := run("name in 'testdata/sub1' where name != 'test.txt'")
	expect(t, out, "Test3.Txt\ntest2.txt")
}

func TestNameContains(t *testing.T) {
	out := run("name in 'testdata/sub1' where name contains '2'")
	expect(t, out, "test2.txt")
}

func TestNameContainsIgnorecase(t *testing.T) {
	out := run("name in 'testdata/sub1' where name ignorecase contains 'Test'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestContentContains(t *testing.T) {
	out := run("name in 'testdata/sub1' where content contains 'some'")
	expect(t, out, "test.txt\ntest2.txt")
}

func TestContentContainsIgnorecase(t *testing.T) {
	out := run("name in 'testdata/sub1' where content ignorecase contains 'some'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestContentContainsIgnorecaseWithCaps(t *testing.T) {
	out := run("name in 'testdata/sub1' where content ignorecase contains 'SOME'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestContentEndswith(t *testing.T) {
	out := run("name in 'testdata/sub1' where content endswith '2'")
	expect(t, out, "test2.txt")
}

func TestContentEndswithWithContains(t *testing.T) {
	out := run("name in 'testdata/sub1' where content contains 'data' and content endswith '2'")
	expect(t, out, "test2.txt")
}

func TestContentStartsContainsAndEnds(t *testing.T) {
	out := run("name in 'testdata/sub1' where content startswith 'some' and content contains 'data' and content endswith '2'")
	expect(t, out, "test2.txt")
}

func TestPathExtraction(t *testing.T) {
	out := run("path in 'testdata/sub1' where name = 'test.txt'")
	if out != "testdata/sub1/test.txt" && out != "testdata\\sub1\\test.txt" {
		t.Error("\nexpected:\n" + "testdata/sub1/test.txt\n\ngot:\n" + out)
	}
}

func TestSizeExtraction(t *testing.T) {
	out := run("size in 'testdata/sub1' where name = 'test.txt'")
	expect(t, out, "10")
}

func TestCompoundSelect(t *testing.T) {
	out := run("name, size in 'testdata/sub1' where name = 'test.txt'")
	expect(t, out, "10          test.txt")
}

func TestOrExpression(t *testing.T) {
	out := run("name in 'testdata/sub2' where isfile or isdir")
	expect(t, out, "a\nb\ndir1/\nc\ndir2/\nd")
}

func TestAndExpression(t *testing.T) {
	out := run("name in 'testdata/sub1' where name startswith 'test.' and name endswith '.txt'")
	expect(t, out, "test.txt")
}

func TestNotExpression(t *testing.T) {
	out := run("name in 'testdata/sub1' where not name startswith 'test'")
	expect(t, out, "Test3.Txt")
}

func TestCompoundExpression(t *testing.T) {
	out := run("name in 'testdata/sub1' where name startswith 'T' or (name startswith 't' and not name contains '2')")
	expect(t, out, "Test3.Txt\ntest.txt")
}

func TestMultipleRoots(t *testing.T) {
	out := run("name in 'testdata/sub1', 'testdata/sub3' where name endswith '.txt'")
	expect(t, out, "test.txt\ntest2.txt\ntest4.txt")
}

func TestNameMatches(t *testing.T) {
	out := run("name in 'testdata/sub1' where name matches '[tT]est[0-9]'")
	expect(t, out, "Test3.Txt\ntest2.txt")
}

func TestPathMatches(t *testing.T) {
	out := run("name in 'testdata/sub1' where path matches '[tT]est[0-9]'")
	expect(t, out, "Test3.Txt\ntest2.txt")
}

func TestContentMatches(t *testing.T) {
	out := run("name in 'testdata/sub1' where content matches '[tT]ext'")
	expect(t, out, "Test3.Txt\ntest.txt")
}
<commit_msg>Initialize maxRoutines in tests.<commit_after>package main

import (
	"bytes"
	"io"
	"os"
	"regexp"
	"strings"
	"testing"
	"runtime"
)

func run(expr string) string {
	old := os.Stdout
	r, w, _ := os.Pipe()
	os.Stdout = w

	maxRoutines = runtime.NumCPU()
	execute_expression(expr)

	outc := make(chan string)
	go func() {
		var buf bytes.Buffer
		io.Copy(&buf, r)
		outc <- buf.String()
	}()

	w.Close()
	os.Stdout = old
	out := <-outc
	return strings.TrimSpace(unpad(out))
}

func expect(t *testing.T, output string, target string) {
	if target != output {
		t.Error("\nexpected:\n" + target + "\n\ngot:\n" + output)
	}
}

func unpad(result string) string {
	re := regexp.MustCompile("[ ]*\\n")
	return re.ReplaceAllString(result, "\n")
}

func TestEndswith(t *testing.T) {
	out := run("name in 'testdata/sub1' where name endswith '.txt'")
	expect(t, out, "test.txt\ntest2.txt")
}

func TestEndswithIgnorecase(t *testing.T) {
	out := run("name in 'testdata/sub1' where name ignorecase endswith '.TxT'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestStartswith(t *testing.T) {
	out := run("name in 'testdata/sub1' where name startswith 'test'")
	expect(t, out, "test.txt\ntest2.txt")
}

func TestStartswithIgnorecase(t *testing.T) {
	out := run("name in 'testdata/sub1' where name ignorecase startswith 'Test'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestIsfile(t *testing.T) {
	out := run("name in 'testdata/sub2' where isfile")
	expect(t, out, "a\nb\nc\nd")
}

func TestIsdir(t *testing.T) {
	out := run("name in 'testdata/sub2' where isdir")
	expect(t, out, "dir1/\ndir2/")
}

func TestNameEq(t *testing.T) {
	out := run("name in 'testdata/sub1' where name = 'test.txt'")
	expect(t, out, "test.txt")
}

func TestNameNeq(t *testing.T) {
	out := run("name in 'testdata/sub1' where name != 'test.txt'")
	expect(t, out, "Test3.Txt\ntest2.txt")
}

func TestNameContains(t *testing.T) {
	out := run("name in 'testdata/sub1' where name contains '2'")
	expect(t, out, "test2.txt")
}

func TestNameContainsIgnorecase(t *testing.T) {
	out := run("name in 'testdata/sub1' where name ignorecase contains 'Test'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestContentContains(t *testing.T) {
	out := run("name in 'testdata/sub1' where content contains 'some'")
	expect(t, out, "test.txt\ntest2.txt")
}

func TestContentContainsIgnorecase(t *testing.T) {
	out := run("name in 'testdata/sub1' where content ignorecase contains 'some'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestContentContainsIgnorecaseWithCaps(t *testing.T) {
	out := run("name in 'testdata/sub1' where content ignorecase contains 'SOME'")
	expect(t, out, "Test3.Txt\ntest.txt\ntest2.txt")
}

func TestContentEndswith(t *testing.T) {
	out := run("name in 'testdata/sub1' where content endswith '2'")
	expect(t, out, "test2.txt")
}

func TestContentEndswithWithContains(t *testing.T) {
	out := run("name in 'testdata/sub1' where content contains 'data' and content endswith '2'")
	expect(t, out, "test2.txt")
}

func TestContentStartsContainsAndEnds(t *testing.T) {
	out := run("name in 'testdata/sub1' where content startswith 'some' and content contains 'data' and content endswith '2'")
	expect(t, out, "test2.txt")
}

func TestPathExtraction(t *testing.T) {
	out := run("path in 'testdata/sub1' where name = 'test.txt'")
	if out != "testdata/sub1/test.txt" && out != "testdata\\sub1\\test.txt" {
		t.Error("\nexpected:\n" + "testdata/sub1/test.txt\n\ngot:\n" + out)
	}
}

func TestSizeExtraction(t *testing.T) {
	out := run("size in 'testdata/sub1' where name = 'test.txt'")
	expect(t, out, "10")
}

func TestCompoundSelect(t *testing.T) {
	out := run("name, size in 'testdata/sub1' where name = 'test.txt'")
	expect(t, out, "10          test.txt")
}

func TestOrExpression(t *testing.T) {
	out := run("name in 'testdata/sub2' where isfile or isdir")
	expect(t, out, "a\nb\ndir1/\nc\ndir2/\nd")
}

func TestAndExpression(t *testing.T) {
	out := run("name in 'testdata/sub1' where name startswith 'test.' and name endswith '.txt'")
	expect(t, out, "test.txt")
}

func TestNotExpression(t *testing.T) {
	out := run("name in 'testdata/sub1' where not name startswith 'test'")
	expect(t, out, "Test3.Txt")
}

func TestCompoundExpression(t *testing.T) {
	out := run("name in 'testdata/sub1' where name startswith 'T' or (name startswith 't' and not name contains '2')")
	expect(t, out, "Test3.Txt\ntest.txt")
}

func TestMultipleRoots(t *testing.T) {
	out := run("name in 'testdata/sub1', 'testdata/sub3' where name endswith '.txt'")
	expect(t, out, "test.txt\ntest2.txt\ntest4.txt")
}

func TestNameMatches(t *testing.T) {
	out := run("name in 'testdata/sub1' where name matches '[tT]est[0-9]'")
	expect(t, out, "Test3.Txt\ntest2.txt")
}

func TestPathMatches(t *testing.T) {
	out := run("name in 'testdata/sub1' where path matches '[tT]est[0-9]'")
	expect(t, out, "Test3.Txt\ntest2.txt")
}

func TestContentMatches(t *testing.T) {
	out := run("name in 'testdata/sub1' where content matches '[tT]ext'")
	expect(t, out, "Test3.Txt\ntest.txt")
}
<|endoftext|>


<commit_before>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 ogletest_test

import (
	"errors"
	"fmt"
	"io/ioutil"
	"path"
	"os"
	"os/exec"
	"strings"
	"testing"
)

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

// getCaseNames looks for integration test cases as files in the
// integration_test_cases directory.
func getCaseNames() ([]string, error) {
	// Open the test cases directory.
	dir, err := os.Open("integration_test_cases")
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Opening dir: %v", err))
	}

	// Get a list of the names in the directory.
	names, err := dir.Readdirnames(0)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Readdirnames: %v", err))
	}

	// Filter the names.
	result := make([]string, len(names))
	resultLen := 0
	for _, name := range names {
		// Skip golden files.
		if strings.HasPrefix(name, "golden.") {
			continue
		}

		// Check for the right format.
		if !strings.HasSuffix(name, "_test.go") {
			return nil, errors.New(fmt.Sprintf("Unexpected file: %s", name))
		}

		// Store the name minus the extension.
		result[resultLen] = name[len(name) - 2:]
	}

	return result, nil
}

func writeStringToFileOrDie(contents []byte, path string) {
	if err := ioutil.WriteFile(path, contents, 0600); err != nil {
		panic("iotuil.WriteFile: " + err.Error())
	}
}

func readFileOrDie(path string) []byte {
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		panic("iotuil.ReadFile: " + err.Error())
	}

	return contents
}

// runTestCase runs the case with the supplied name (e.g. "passing_test"), and
// returns its output and exit code.
func runTestCase(name string) ([]byte, int, error) {
	// Create a temporary directory for the test files.
	tempDir, err := ioutil.TempDir("", "ogletest_integration_test")
	if err != nil {
		return nil, 0, errors.New("ioutil.TempDir: " + err.Error())
	}

	// Create a makefile within the directory.
	makefileContents := "include $(GOROOT)/src/Make.inc\n" +
		"TARG = github.com/jacobsa/ogletest/foobar\n" +
		"GOFILES = " + name + ".go\n" +
		"include $(GOROOT)/src/Make.pkg\n"

	writeStringToFileOrDie([]byte(makefileContents), path.Join(tempDir, "Makefile"))

	// Create the test source file.
	sourceFile := name + ".go"
	testContents := readFileOrDie(path.Join("integration_test_cases", sourceFile))
	writeStringToFileOrDie(testContents, path.Join(tempDir, sourceFile))

	// Invoke gotest.
	cmd := exec.Command("gotest")
	cmd.Dir = tempDir
	output, err := cmd.Output()

	// Did the process exist with zero code?
	if err == nil {
		return output, 0, nil
	}

	// Make sure the process actually exited.
	exitError, ok := err.(*exec.ExitError)
	if !ok || !exitError.Exited() {
		return nil, 0, errors.New("exec.Command.Output: " + err.Error())
	}

	return output, exitError.ExitStatus(), nil
}

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func TestGoldenFiles(t *testing.T) {
	// We expect there to be at least one case.
	caseNames, err := getCaseNames()
	if err != nil || len(caseNames) == 0 {
		t.Fatalf("Error getting cases: %v", err)
	}
}
<commit_msg>Implemented part of TestGoldenFiles, for #2.<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 ogletest_test

import (
	"errors"
	"fmt"
	"io/ioutil"
	"path"
	"os"
	"os/exec"
	"strings"
	"testing"
)

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

// getCaseNames looks for integration test cases as files in the
// integration_test_cases directory.
func getCaseNames() ([]string, error) {
	// Open the test cases directory.
	dir, err := os.Open("integration_test_cases")
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Opening dir: %v", err))
	}

	// Get a list of the names in the directory.
	names, err := dir.Readdirnames(0)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Readdirnames: %v", err))
	}

	// Filter the names.
	result := make([]string, len(names))
	resultLen := 0
	for _, name := range names {
		// Skip golden files.
		if strings.HasPrefix(name, "golden.") {
			continue
		}

		// Check for the right format.
		if !strings.HasSuffix(name, "_test.go") {
			return nil, errors.New(fmt.Sprintf("Unexpected file: %s", name))
		}

		// Store the name minus the extension.
		result[resultLen] = name[len(name) - 2:]
	}

	return result, nil
}

func writeStringToFileOrDie(contents []byte, path string) {
	if err := ioutil.WriteFile(path, contents, 0600); err != nil {
		panic("iotuil.WriteFile: " + err.Error())
	}
}

func readFileOrDie(path string) []byte {
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		panic("iotuil.ReadFile: " + err.Error())
	}

	return contents
}

// runTestCase runs the case with the supplied name (e.g. "passing_test"), and
// returns its output and exit code.
func runTestCase(name string) ([]byte, int, error) {
	// Create a temporary directory for the test files.
	tempDir, err := ioutil.TempDir("", "ogletest_integration_test")
	if err != nil {
		return nil, 0, errors.New("ioutil.TempDir: " + err.Error())
	}

	// Create a makefile within the directory.
	makefileContents := "include $(GOROOT)/src/Make.inc\n" +
		"TARG = github.com/jacobsa/ogletest/foobar\n" +
		"GOFILES = " + name + ".go\n" +
		"include $(GOROOT)/src/Make.pkg\n"

	writeStringToFileOrDie([]byte(makefileContents), path.Join(tempDir, "Makefile"))

	// Create the test source file.
	sourceFile := name + ".go"
	testContents := readFileOrDie(path.Join("integration_test_cases", sourceFile))
	writeStringToFileOrDie(testContents, path.Join(tempDir, sourceFile))

	// Invoke gotest.
	cmd := exec.Command("gotest")
	cmd.Dir = tempDir
	output, err := cmd.Output()

	// Did the process exist with zero code?
	if err == nil {
		return output, 0, nil
	}

	// Make sure the process actually exited.
	exitError, ok := err.(*exec.ExitError)
	if !ok || !exitError.Exited() {
		return nil, 0, errors.New("exec.Command.Output: " + err.Error())
	}

	return output, exitError.ExitStatus(), nil
}

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func TestGoldenFiles(t *testing.T) {
	// We expect there to be at least one case.
	caseNames, err := getCaseNames()
	if err != nil || len(caseNames) == 0 {
		t.Fatalf("Error getting cases: %v", err)
	}

	// Run each test case.
	for _, caseName := range caseNames {
		t.Logf("Running test case: %s", caseName)

		// Run the test case.
		_, exitCode, err := runTestCase(caseName)
		if err != nil {
			t.Fatalf("Running test case %s: %v", caseName, err)
		}

		// Check the status code. We assume all test cases fail except for passing_test.
		shouldPass := caseName == "passint_test"
		didPass := exitCode == 0
		if shouldPass != didPass {
			t.Errorf("Bad exit code for test case %s: %d", caseName, exitCode)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Andreas Klauer <Andreas.Klauer@metamorpher.de>
// License: MIT

// Package nbd uses the Linux NBD layer to emulate a block device in user space
package nbd

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"sync"
	"syscall"
	"time"
)

const (
	// Defined in <linux/fs.h>:
	BLKROSET = 4701

	// Defined in <linux/nbd.h>:
	ioctlSetSock       = 43776
	ioctlSetBlockSize  = 43777
	ioctlSetSize       = 43778
	ioctlDoIt          = 43779
	ioctlClearSock     = 43780
	ioctlClearQueue    = 43781
	ioctlSetSizeBlocks = 43783
	ioctlDisconnect    = 43784
	ioctlSetFlags      = 43786
)

const (
	cmdRead  = 0
	cmdWrite = 1
	cmdDisc  = 2
	cmdFlush = 3
	cmdTrim  = 4
)

const (
	flagHasFlags  = (1 << 0) // nbd-server supports flags
	flagSendFlush = (1 << 2) // can flush writeback cache
	flagSendTrim  = (1 << 5) // Send TRIM (discard)
	// flagReadOnly   = (1 << 1) // device is read-only
	// flagSendFUA    = (1 << 3) // Send FUA (Force Unit Access)
	// flagRotational = (1 << 4) // Use elevator algorithm - rotational media
)

const (
	magicRequest = 0x25609513
	magicReply   = 0x67446698
	// Do *not* use magics: 0x12560953 0x96744668.
)

const (
	errIO = 5
)

// ioctl() helper function
func ioctl(a1, a2, a3 uintptr) (err error) {
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, a1, a2, a3)
	if errno != 0 {
		err = errno
	}
	return err
}

// Device interface is a subset of os.File.
type Device interface {
	ReadAt(b []byte, off int64) (n int, err error)
	WriteAt(b []byte, off int64) (n int, err error)
	Sync() error
	Trim(off, len int64) error
	Size() uint64
	Close() error
}

type NBD struct {
	device    Device
	size      int64
	blocksize int64
	nbd       *os.File
	socket    int
	setsocket int
	closer    chan error
}

func Create(device Device, size int64, blocksize int64) *NBD {
	if size >= 0 {
		return &NBD{
			device:    device,
			size:      size,
			blocksize: blocksize,
			nbd:       nil,
			socket:    0,
		}
	}
	return nil
}

// return true if connected
func (nbd *NBD) IsConnected() bool {
	return nbd.nbd != nil && nbd.socket > 0
}

func (nbd *NBD) Size() int64 {
	return nbd.size
}

func (nbd *NBD) SetSize(size int64) error {
	if err := ioctl(nbd.nbd.Fd(), ioctlSetSize, uintptr(size)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_SIZE",
			Err:  err,
		}
	}
	return nil
}

func (nbd *NBD) SetBlockSize(blocksize int64) error {
	if err := ioctl(nbd.nbd.Fd(), ioctlSetBlockSize, uintptr(blocksize)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_BLKSIZE",
			Err:  err,
		}
	}
	return nil
}

func FindDevice() (string, error) {
	// FIXME: Oh god... fixme.
	// find free nbd device
	for i := 0; ; i++ {
		dev := fmt.Sprintf("/dev/nbd%d", i)
		if _, err := os.Stat(dev); os.IsNotExist(err) {
			break // no more devices
		}
		if _, err := os.Stat(fmt.Sprintf("/sys/block/nbd%d/pid", i)); !os.IsNotExist(err) {
			continue // busy
		}
		return dev, nil
	}
	return "", errors.New("no devices available")
}

func (nbd *NBD) OpenDevice(dev string) (string, error) {
	f, err := os.Open(dev)
	if err != nil {
		return "", err
	}
	nbd.nbd = f

	// possible candidate
	ioctl(f.Fd(), BLKROSET, 0) // I'm really sorry about this
	pair, err := syscall.Socketpair(syscall.SOCK_STREAM, syscall.AF_UNIX, 0)
	if err != nil {
		return "", err
	}
	if err := ioctl(f.Fd(), ioctlSetSock, uintptr(pair[0])); err != nil {
		return "", err
	}

	nbd.setsocket = pair[0] // FIXME: We shouldn't hold on to this.
	nbd.socket = pair[1]
	return dev, nil
}

func (nbd *NBD) Serve() error {
	blksized := true
	if err := nbd.SetSize(nbd.size); err != nil {
		return err // already set by nbd.Size()
	}
	if err := nbd.SetBlockSize(nbd.blocksize); err != nil {
		// This is a hack around the changes made to the kernel in 4.6
		// (particularly commit 37091fdd831f28a6509008542174ed324dd645bc)
		// -- because the size of the device is cached at 0, the blocksize can't change
		// until connected. So we'll do a workaround on newer kernels, but man, it'd be
		// nice to fix this. There needs to be a little better logic kernel-side around changing size
		// even when disconnected. Changing it only when connected is fine -- but keep my intent.
		blksized = false
	}
	if err := ioctl(nbd.nbd.Fd(), ioctlSetFlags, uintptr(flagSendFlush|flagSendTrim)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_FLAGS",
			Err:  err,
		}
	}

	c := &serverConn{
		rw: os.NewFile(uintptr(nbd.socket), "<nbd socket>"),
	}
	// TODO(barakmich): Scale up NBD by handling multiple requests.
	// Requires thread-safety across the block.BlockFile/torus.File
	//n := runtime.GOMAXPROCS(0) - 1
	n := 1

	wg := new(sync.WaitGroup)
	wg.Add(n)
	for i := 0; i < n; i++ {
		go c.serveLoop(nbd.device, wg)
	}
	if !blksized {
		// Back to the hack.
		go func(nbd *NBD) {
			// Hopefully we'll be connected in 500 millis.
			// If not, we'll proceed with the standard blocksize of 1K.
			time.Sleep(time.Microsecond * 500)
			err := nbd.SetBlockSize(nbd.blocksize)
			if err != nil {
				log.Printf("Couldn't upgrade blocksize: %s", err)
			}
		}(nbd)
	}
	if err := waitDisconnect(nbd.nbd); err != nil {
		return err
	}
	wg.Wait()
	return nil
}

func (nbd *NBD) Close() error {
	// FIXME: Now I understand what these did, and it terrifies me.
	// syscall.Write(nbd.setsocket, make([]byte, 28))
	// syscall.Close(nbd.setsocket)
	nbd.Disconnect()
	syscall.Close(nbd.socket)
	return nbd.nbd.Close()
}

func (nbd *NBD) Disconnect() error {
	if nbd.nbd == nil {
		return nil
	}
	if err := ioctl(nbd.nbd.Fd(), ioctlDisconnect, 0); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_DISCONNECT",
			Err:  err,
		}
	}
	return nil
}

func waitDisconnect(f *os.File) error {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	// NBD_DO_IT does not return until disconnect
	if err := ioctl(f.Fd(), ioctlDoIt, 0); err != nil {
		return &os.PathError{
			Path: f.Name(),
			Op:   "ioctl NBD_DO_IT",
			Err:  err,
		}
	}
	if err := ioctl(f.Fd(), ioctlClearQueue, 0); err != nil {
		return &os.PathError{
			Path: f.Name(),
			Op:   "ioctl NBD_CLEAR_QUEUE",
			Err:  err,
		}
	}
	if err := ioctl(f.Fd(), ioctlClearSock, 0); err != nil {
		return &os.PathError{
			Path: f.Name(),
			Op:   "ioctl NBD_CLEAR_SOCK",
			Err:  err,
		}
	}
	return nil
}

type serverConn struct {
	mu sync.Mutex
	rw io.ReadWriteCloser
}

func (c *serverConn) serveLoop(dev Device, wg *sync.WaitGroup) error {
	if wg != nil {
		defer wg.Done()
	}

	hdr := new(reqHeader)
	buf := make([]byte, 16) // FIXME: Maybe this should be hdr[:16] instead?
	for {
		c.mu.Lock()

		if _, err := io.ReadFull(c.rw, hdr[:]); err != nil {
			// FIXME: Are there any valid short reads we need to handle?
			c.mu.Unlock()
			return err
		}

		if magic := hdr.magic(); magic != magicRequest {
			c.mu.Unlock()
			return fmt.Errorf("nbd: invalid magic: 0x%x", magic)
		}

		cmd, _ := hdr.command()
		if cmd == cmdWrite {
			buf = hdr.resize(buf)
			if _, err := io.ReadFull(c.rw, buf[16:]); err != nil {
				return err
			}
		}
		c.mu.Unlock()

		switch cmd {
		case cmdRead:
			buf = hdr.resize(buf)
			if _, err := dev.ReadAt(buf[16:], hdr.offset()); err != nil {
				hdr.putReplyHeader(buf, errIO)
			} else {
				hdr.putReplyHeader(buf, 0)
			}
		case cmdWrite:
			if _, err := dev.WriteAt(buf[16:], hdr.offset()); err != nil {
				hdr.putReplyHeader(buf, errIO)
			} else {
				hdr.putReplyHeader(buf, 0)
			}
			buf = buf[:16]
		case cmdTrim:
			if err := dev.Trim(hdr.offset(), int64(hdr.length())); err != nil {
				log.Printf("nbd: trim error: %s", err)
			}
			fallthrough
		case cmdFlush:
			if err := dev.Sync(); err != nil {
				log.Printf("nbd: sync error: %s", err)
			}
			hdr.putReplyHeader(buf, 0)
			buf = buf[:16]
		case cmdDisc:
			// FIXME: We're actually supposed to wait for outstanding requests to finish.
			if err := dev.Sync(); err != nil {
				log.Printf("nbd: sync error: %s", err)
			}
			return c.rw.Close()
		default:
			return errors.New("nbd: invalid command")
		}

		// FIXME: Are we sure this whole write is going to be atomic?
		if _, err := c.rw.Write(buf); err != nil {
			return err
		}
	}
}

type reqHeader [28]byte

func (h *reqHeader) command() (cmd, flags uint16) {
	return binary.BigEndian.Uint16(h[6:8]), binary.BigEndian.Uint16(h[4:6])
}

func (h *reqHeader) handle() uint64 {
	return binary.BigEndian.Uint64(h[8:16])
}

func (h *reqHeader) length() uint32 {
	return binary.BigEndian.Uint32(h[24:28])
}

func (h *reqHeader) magic() uint32 {
	return binary.BigEndian.Uint32(h[0:4])
}

func (h *reqHeader) offset() int64 {
	const maxInt64 = (^uint64(0) >> 1) + 1
	if u := binary.BigEndian.Uint64(h[16:24]); u <= maxInt64 {
		return int64(u)
	}
	panic("nbd: offset overflow")
}

func (h *reqHeader) resize(buf []byte) []byte {
	if n := int(h.length()) + 16; n > cap(buf) {
		return make([]byte, n)
	} else {
		return buf[:n]
	}
}

func (h *reqHeader) putReplyHeader(dst []byte, err uint32) {
	binary.BigEndian.PutUint32(dst[0:4], magicReply)
	binary.BigEndian.PutUint32(dst[4:8], err)
	copy(dst[8:16], h[8:16])
}
<commit_msg>nbd: don't return error during clean up process<commit_after>// Copyright (C) 2014 Andreas Klauer <Andreas.Klauer@metamorpher.de>
// License: MIT

// Package nbd uses the Linux NBD layer to emulate a block device in user space
package nbd

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"sync"
	"syscall"
	"time"
)

const (
	// Defined in <linux/fs.h>:
	BLKROSET = 4701

	// Defined in <linux/nbd.h>:
	ioctlSetSock       = 43776
	ioctlSetBlockSize  = 43777
	ioctlSetSize       = 43778
	ioctlDoIt          = 43779
	ioctlClearSock     = 43780
	ioctlClearQueue    = 43781
	ioctlSetSizeBlocks = 43783
	ioctlDisconnect    = 43784
	ioctlSetFlags      = 43786
)

const (
	cmdRead  = 0
	cmdWrite = 1
	cmdDisc  = 2
	cmdFlush = 3
	cmdTrim  = 4
)

const (
	flagHasFlags  = (1 << 0) // nbd-server supports flags
	flagSendFlush = (1 << 2) // can flush writeback cache
	flagSendTrim  = (1 << 5) // Send TRIM (discard)
	// flagReadOnly   = (1 << 1) // device is read-only
	// flagSendFUA    = (1 << 3) // Send FUA (Force Unit Access)
	// flagRotational = (1 << 4) // Use elevator algorithm - rotational media
)

const (
	magicRequest = 0x25609513
	magicReply   = 0x67446698
	// Do *not* use magics: 0x12560953 0x96744668.
)

const (
	errIO = 5
)

// ioctl() helper function
func ioctl(a1, a2, a3 uintptr) (err error) {
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, a1, a2, a3)
	if errno != 0 {
		err = errno
	}
	return err
}

// Device interface is a subset of os.File.
type Device interface {
	ReadAt(b []byte, off int64) (n int, err error)
	WriteAt(b []byte, off int64) (n int, err error)
	Sync() error
	Trim(off, len int64) error
	Size() uint64
	Close() error
}

type NBD struct {
	device    Device
	size      int64
	blocksize int64
	nbd       *os.File
	socket    int
	setsocket int
	closer    chan error
}

func Create(device Device, size int64, blocksize int64) *NBD {
	if size >= 0 {
		return &NBD{
			device:    device,
			size:      size,
			blocksize: blocksize,
			nbd:       nil,
			socket:    0,
		}
	}
	return nil
}

// return true if connected
func (nbd *NBD) IsConnected() bool {
	return nbd.nbd != nil && nbd.socket > 0
}

func (nbd *NBD) Size() int64 {
	return nbd.size
}

func (nbd *NBD) SetSize(size int64) error {
	if err := ioctl(nbd.nbd.Fd(), ioctlSetSize, uintptr(size)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_SIZE",
			Err:  err,
		}
	}
	return nil
}

func (nbd *NBD) SetBlockSize(blocksize int64) error {
	if err := ioctl(nbd.nbd.Fd(), ioctlSetBlockSize, uintptr(blocksize)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_BLKSIZE",
			Err:  err,
		}
	}
	return nil
}

func FindDevice() (string, error) {
	// FIXME: Oh god... fixme.
	// find free nbd device
	for i := 0; ; i++ {
		dev := fmt.Sprintf("/dev/nbd%d", i)
		if _, err := os.Stat(dev); os.IsNotExist(err) {
			break // no more devices
		}
		if _, err := os.Stat(fmt.Sprintf("/sys/block/nbd%d/pid", i)); !os.IsNotExist(err) {
			continue // busy
		}
		return dev, nil
	}
	return "", errors.New("no devices available")
}

func (nbd *NBD) OpenDevice(dev string) (string, error) {
	f, err := os.Open(dev)
	if err != nil {
		return "", err
	}
	nbd.nbd = f

	// possible candidate
	ioctl(f.Fd(), BLKROSET, 0) // I'm really sorry about this
	pair, err := syscall.Socketpair(syscall.SOCK_STREAM, syscall.AF_UNIX, 0)
	if err != nil {
		return "", err
	}
	if err := ioctl(f.Fd(), ioctlSetSock, uintptr(pair[0])); err != nil {
		return "", err
	}

	nbd.setsocket = pair[0] // FIXME: We shouldn't hold on to this.
	nbd.socket = pair[1]
	return dev, nil
}

func (nbd *NBD) Serve() error {
	blksized := true
	if err := nbd.SetSize(nbd.size); err != nil {
		return err // already set by nbd.Size()
	}
	if err := nbd.SetBlockSize(nbd.blocksize); err != nil {
		// This is a hack around the changes made to the kernel in 4.6
		// (particularly commit 37091fdd831f28a6509008542174ed324dd645bc)
		// -- because the size of the device is cached at 0, the blocksize can't change
		// until connected. So we'll do a workaround on newer kernels, but man, it'd be
		// nice to fix this. There needs to be a little better logic kernel-side around changing size
		// even when disconnected. Changing it only when connected is fine -- but keep my intent.
		blksized = false
	}
	if err := ioctl(nbd.nbd.Fd(), ioctlSetFlags, uintptr(flagSendFlush|flagSendTrim)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_FLAGS",
			Err:  err,
		}
	}

	c := &serverConn{
		rw: os.NewFile(uintptr(nbd.socket), "<nbd socket>"),
	}
	// TODO(barakmich): Scale up NBD by handling multiple requests.
	// Requires thread-safety across the block.BlockFile/torus.File
	//n := runtime.GOMAXPROCS(0) - 1
	n := 1

	wg := new(sync.WaitGroup)
	wg.Add(n)
	for i := 0; i < n; i++ {
		go c.serveLoop(nbd.device, wg)
	}
	if !blksized {
		// Back to the hack.
		go func(nbd *NBD) {
			// Hopefully we'll be connected in 500 millis.
			// If not, we'll proceed with the standard blocksize of 1K.
			time.Sleep(time.Microsecond * 500)
			err := nbd.SetBlockSize(nbd.blocksize)
			if err != nil {
				log.Printf("Couldn't upgrade blocksize: %s", err)
			}
		}(nbd)
	}

	waitDisconnect(nbd.nbd)

	wg.Wait()
	return nil
}

func (nbd *NBD) Close() error {
	// FIXME: Now I understand what these did, and it terrifies me.
	// syscall.Write(nbd.setsocket, make([]byte, 28))
	// syscall.Close(nbd.setsocket)
	nbd.Disconnect()
	syscall.Close(nbd.socket)
	return nbd.nbd.Close()
}

func (nbd *NBD) Disconnect() error {
	if nbd.nbd == nil {
		return nil
	}
	if err := ioctl(nbd.nbd.Fd(), ioctlDisconnect, 0); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_DISCONNECT",
			Err:  err,
		}
	}
	return nil
}

func waitDisconnect(f *os.File) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	// NBD_DO_IT does not return until disconnect
	if err := ioctl(f.Fd(), ioctlDoIt, 0); err != nil {
		clog.Errorf("error %s: ioctl returned %v", f.Name(), err)
	}

	clog.Debugf("Running disconnection to %s", f.Name())

	if err := ioctl(f.Fd(), ioctlClearQueue, 0); err != nil {
		clog.Errorf("error clear queue for %s. ioctl returned: %v", f.Name(), err)
	}
	if err := ioctl(f.Fd(), ioctlClearSock, 0); err != nil {
		clog.Errorf("error clear socket for %s. ioctl returned: %v", f.Name(), err)
	}
}

type serverConn struct {
	mu sync.Mutex
	rw io.ReadWriteCloser
}

func (c *serverConn) serveLoop(dev Device, wg *sync.WaitGroup) error {
	if wg != nil {
		defer wg.Done()
	}

	hdr := new(reqHeader)
	buf := make([]byte, 16) // FIXME: Maybe this should be hdr[:16] instead?
	for {
		c.mu.Lock()

		if _, err := io.ReadFull(c.rw, hdr[:]); err != nil {
			// FIXME: Are there any valid short reads we need to handle?
			c.mu.Unlock()
			return err
		}

		if magic := hdr.magic(); magic != magicRequest {
			c.mu.Unlock()
			return fmt.Errorf("nbd: invalid magic: 0x%x", magic)
		}

		cmd, _ := hdr.command()
		if cmd == cmdWrite {
			buf = hdr.resize(buf)
			if _, err := io.ReadFull(c.rw, buf[16:]); err != nil {
				return err
			}
		}
		c.mu.Unlock()

		switch cmd {
		case cmdRead:
			buf = hdr.resize(buf)
			if _, err := dev.ReadAt(buf[16:], hdr.offset()); err != nil {
				hdr.putReplyHeader(buf, errIO)
			} else {
				hdr.putReplyHeader(buf, 0)
			}
		case cmdWrite:
			if _, err := dev.WriteAt(buf[16:], hdr.offset()); err != nil {
				hdr.putReplyHeader(buf, errIO)
			} else {
				hdr.putReplyHeader(buf, 0)
			}
			buf = buf[:16]
		case cmdTrim:
			if err := dev.Trim(hdr.offset(), int64(hdr.length())); err != nil {
				log.Printf("nbd: trim error: %s", err)
			}
			fallthrough
		case cmdFlush:
			if err := dev.Sync(); err != nil {
				log.Printf("nbd: sync error: %s", err)
			}
			hdr.putReplyHeader(buf, 0)
			buf = buf[:16]
		case cmdDisc:
			// FIXME: We're actually supposed to wait for outstanding requests to finish.
			if err := dev.Sync(); err != nil {
				log.Printf("nbd: sync error: %s", err)
			}
			return c.rw.Close()
		default:
			return errors.New("nbd: invalid command")
		}

		// FIXME: Are we sure this whole write is going to be atomic?
		if _, err := c.rw.Write(buf); err != nil {
			return err
		}
	}
}

type reqHeader [28]byte

func (h *reqHeader) command() (cmd, flags uint16) {
	return binary.BigEndian.Uint16(h[6:8]), binary.BigEndian.Uint16(h[4:6])
}

func (h *reqHeader) handle() uint64 {
	return binary.BigEndian.Uint64(h[8:16])
}

func (h *reqHeader) length() uint32 {
	return binary.BigEndian.Uint32(h[24:28])
}

func (h *reqHeader) magic() uint32 {
	return binary.BigEndian.Uint32(h[0:4])
}

func (h *reqHeader) offset() int64 {
	const maxInt64 = (^uint64(0) >> 1) + 1
	if u := binary.BigEndian.Uint64(h[16:24]); u <= maxInt64 {
		return int64(u)
	}
	panic("nbd: offset overflow")
}

func (h *reqHeader) resize(buf []byte) []byte {
	if n := int(h.length()) + 16; n > cap(buf) {
		return make([]byte, n)
	} else {
		return buf[:n]
	}
}

func (h *reqHeader) putReplyHeader(dst []byte, err uint32) {
	binary.BigEndian.PutUint32(dst[0:4], magicReply)
	binary.BigEndian.PutUint32(dst[4:8], err)
	copy(dst[8:16], h[8:16])
}
<|endoftext|>


<commit_before>package irc

import(
	"testing"
)

func TestParseMessage(t *testing.T) {
	msg_string := ":kyle!~kyle@localhost PRIVMSG #tenyks :tenyks: messages are awesome"
	msg := ParseMessage(msg_string)
	if msg == nil {
		t.Error("Expected", Message{}, "got", msg)
	}

	if msg.Command != "PRIVMSG" {
		t.Error("Expected", Message{}, "got", msg)
	}

	if msg.Trail != "tenyks: messages are awesome" {
		t.Error("Expected", "tenyks: messages are awesome", "got", msg)
	}
}
<commit_msg>Test params for things like target<commit_after>package irc

import(
	"testing"
)

func TestParseMessage(t *testing.T) {
	msg_string := ":kyle!~kyle@localhost PRIVMSG #tenyks :tenyks: messages are awesome"
	msg := ParseMessage(msg_string)
	if msg == nil {
		t.Error("Expected", Message{}, "got", msg)
	}

	if msg.Command != "PRIVMSG" {
		t.Error("Expected", Message{}, "got", msg)
	}

	if msg.Trail != "tenyks: messages are awesome" {
		t.Error("Expected", "tenyks: messages are awesome", "got", msg)
	}

	if msg.Params[0] != "#tenyks" {
		t.Error("Expected", "#tenyks", "got", msg.Params[0])
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/go-errors/errors"
	"github.com/mdp/qrterminal"
	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/common"
	"github.com/privacybydesign/irmago/server"
	sseclient "github.com/sietseringers/go-sse"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// requestCmd represents the request command
var requestCmd = &cobra.Command{
	Use:   "request",
	Short: "Generate an IRMA session request",
	Run: func(cmd *cobra.Command, args []string) {
		request, _, err := configureRequest(cmd)
		if err != nil {
			die("", err)
		}

		flags := cmd.Flags()
		authmethod, _ := flags.GetString("authmethod")
		var output string
		if authmethod == "none" || authmethod == "token" {
			output = prettyprint(request)
		} else {
			key, _ := flags.GetString("key")
			name, _ := flags.GetString("name")
			if output, err = signRequest(request, name, authmethod, key); err != nil {
				die("Failed to sign request", err)
			}
		}

		fmt.Println(output)
	},
}

func configureJWTKey(authmethod, key string) (interface{}, jwt.SigningMethod, error) {
	var (
		err    error
		sk     interface{}
		jwtalg jwt.SigningMethod
		bts    []byte
	)
	// If the key refers to an existing file, use contents of the file as key
	if bts, err = common.ReadKey("", key); err != nil {
		bts = []byte(key)
	}
	switch authmethod {
	case "hmac":
		jwtalg = jwt.SigningMethodHS256
		if sk, err = common.Base64Decode(bts); err != nil {
			return nil, nil, err
		}
	case "rsa":
		jwtalg = jwt.SigningMethodRS256
		if sk, err = jwt.ParseRSAPrivateKeyFromPEM(bts); err != nil {
			return nil, nil, err
		}
	default:
		return nil, nil, errors.Errorf("Unsupported signing algorithm: '%s'", authmethod)
	}

	return sk, jwtalg, nil
}

func signRequest(request irma.RequestorRequest, name, authmethod, key string) (string, error) {
	sk, jwtalg, err := configureJWTKey(authmethod, key)
	if err != nil {
		return "", err
	}
	return irma.SignRequestorRequest(request, jwtalg, sk, name)
}

func configureRequest(cmd *cobra.Command) (irma.RequestorRequest, *irma.Configuration, error) {
	irmaconfigPath, err := cmd.Flags().GetString("schemes-path")
	if err != nil {
		return nil, nil, err
	}
	irmaconfig, err := irma.NewConfiguration(irmaconfigPath, irma.ConfigurationOptions{})
	if err != nil {
		return nil, nil, err
	}
	if err = irmaconfig.ParseFolder(); err != nil {
		return nil, nil, err
	}
	if len(irmaconfig.SchemeManagers) == 0 {
		if err = irmaconfig.DownloadDefaultSchemes(); err != nil {
			return nil, nil, err
		}
	}

	request, err := constructSessionRequest(cmd, irmaconfig)
	if err != nil {
		return nil, nil, err
	}

	return request, irmaconfig, nil
}

// Helper functions

func wait(initialStatus server.Status, transport *irma.HTTPTransport, statuschan chan server.Status) {
	events := make(chan *sseclient.Event)

	go func() {
		for {
			if e := <-events; e != nil && e.Type != "open" {
				status := server.Status(strings.Trim(string(e.Data), `"`))
				statuschan <- status
				if status.Finished() {
					return
				}
			}
		}
	}()

	if err := sseclient.Notify(nil, transport.Server+"statusevents", true, events); err != nil {
		fmt.Println("SSE failed, fallback to polling", err)
		close(events)
		poll(initialStatus, transport, statuschan)
		return
	}
}

// poll recursively polls the session status until a final status is received.
func poll(initialStatus server.Status, transport *irma.HTTPTransport, statuschan chan server.Status) {
	// First we wait
	<-time.NewTimer(pollInterval).C

	// Get session status
	var s string
	if err := transport.Get("status", &s); err != nil {
		_ = server.LogFatal(err)
	}
	status := server.Status(strings.Trim(s, `"`))

	// report if status changed
	if status != initialStatus {
		statuschan <- status
	}

	if status.Finished() {
		return
	}
	go poll(status, transport, statuschan)
}

func constructSessionRequest(cmd *cobra.Command, conf *irma.Configuration) (irma.RequestorRequest, error) {
	disclose, _ := cmd.Flags().GetStringArray("disclose")
	issue, _ := cmd.Flags().GetStringArray("issue")
	sign, _ := cmd.Flags().GetStringArray("sign")
	message, _ := cmd.Flags().GetString("message")
	jsonrequest, _ := cmd.Flags().GetString("request")

	if len(disclose) == 0 && len(issue) == 0 && len(sign) == 0 && message == "" {
		if jsonrequest == "" {
			return nil, errors.New("Provide either a complete session request using --request or construct one using the other flags")
		}
		request, err := server.ParseSessionRequest(jsonrequest)
		if err != nil {
			return nil, err
		}
		return request, nil
	}

	if jsonrequest != "" {
		return nil, errors.New("Provide either a complete session request using --request or construct one using the other flags")
	}

	if len(sign) != 0 {
		if len(disclose) != 0 {
			return nil, errors.New("cannot combine disclosure and signature sessions, use either --disclose or --sign")
		}
		if len(issue) != 0 {
			return nil, errors.New("cannot combine issuance and signature sessions, use either --issue or --sign")
		}
		if message == "" {
			return nil, errors.New("signature sessions require a message to be signed using --message")
		}
	}

	var request irma.RequestorRequest
	if len(disclose) != 0 && len(issue) == 0 {
		disclose, err := parseAttrs(disclose, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.ServiceProviderRequest{
			Request: irma.NewDisclosureRequest(),
		}
		request.SessionRequest().(*irma.DisclosureRequest).Disclose = disclose
	}
	if len(sign) != 0 {
		disclose, err := parseAttrs(sign, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.SignatureRequestorRequest{
			Request: irma.NewSignatureRequest(message),
		}
		request.SessionRequest().(*irma.SignatureRequest).Disclose = disclose
	}
	if len(issue) != 0 {
		creds, err := parseCredentials(issue, conf)
		if err != nil {
			return nil, err
		}
		disclose, err := parseAttrs(disclose, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.IdentityProviderRequest{
			Request: irma.NewIssuanceRequest(creds),
		}
		request.SessionRequest().(*irma.IssuanceRequest).Disclose = disclose
	}

	return request, nil
}

func parseCredentials(credentialsStr []string, conf *irma.Configuration) ([]*irma.CredentialRequest, error) {
	list := make([]*irma.CredentialRequest, 0, len(credentialsStr))

	for _, credStr := range credentialsStr {
		parts := strings.Split(credStr, "=")
		if len(parts) != 2 {
			return nil, errors.New("--issue argument must contain exactly 1 = sign")
		}
		credIdStr, attrsStr := parts[0], parts[1]
		credtype := conf.CredentialTypes[irma.NewCredentialTypeIdentifier(credIdStr)]
		if credtype == nil {
			return nil, errors.New("unknown credential type: " + credIdStr)
		}

		// we dont need values for the random blind attributes,
		// since they are randomly generated during issuance
		missingValues := credtype.RandomBlinds()
		if len(credtype.RevocationServers) > 0 {
			missingValues += 1
		}
		required := len(credtype.AttributeTypes) - missingValues
		attrsSlice := strings.Split(attrsStr, ",")
		if len(attrsSlice) != required {
			return nil, errors.Errorf("%d attributes required but %d provided for %s", required, len(attrsSlice), credIdStr)
		}

		// Construct the attribute -> value map for in he CredentialRequest
		attrs := make(map[string]string, len(attrsSlice))
		j := 0 // idx in attrSlice
		for _, typ := range credtype.AttributeTypes {
			if !typ.RandomBlind && !typ.RevocationAttribute {
				attrs[typ.ID] = attrsSlice[j]
				j += 1
			}
		}
		list = append(list, &irma.CredentialRequest{
			CredentialTypeID: irma.NewCredentialTypeIdentifier(credIdStr),
			Attributes:       attrs,
		})
	}

	return list, nil
}

func parseAttrs(attrsStr []string, conf *irma.Configuration) (irma.AttributeConDisCon, error) {
	list := make(irma.AttributeConDisCon, 0, len(attrsStr))
	for _, disjunctionStr := range attrsStr {
		disjunction := irma.AttributeDisCon{}
		attrids := strings.Split(disjunctionStr, ",")
		for _, attridStr := range attrids {
			attrid := irma.NewAttributeTypeIdentifier(attridStr)
			if conf.AttributeTypes[attrid] == nil {
				return nil, errors.New("unknown attribute: " + attridStr)
			}
			disjunction = append(disjunction, irma.AttributeCon{irma.AttributeRequest{Type: attrid}})
		}
		list = append(list, disjunction)
	}
	return list, nil
}

func startServer(port int) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", irmaServer.HandlerFunc())
	httpServer = &http.Server{Addr: ":" + strconv.Itoa(port), Handler: mux}
	go func() {
		err := httpServer.ListenAndServe()
		if err != nil && err != http.ErrServerClosed {
			die("Failed to start server", err)
		}
	}()
}

func printQr(qr *irma.Qr, noqr bool) error {
	qrBts, err := json.Marshal(qr)
	if err != nil {
		return err
	}
	if noqr {
		fmt.Println(string(qrBts))
	} else {
		qrterminal.GenerateWithConfig(string(qrBts), qrterminal.Config{
			Level:     qrterminal.L,
			Writer:    os.Stdout,
			BlackChar: qrterminal.BLACK,
			WhiteChar: qrterminal.WHITE,
		})
	}
	return nil
}

func printSessionResult(result *server.SessionResult) {
	fmt.Println("Session result:")
	fmt.Println(prettyprint(result))
}

func init() {
	RootCmd.AddCommand(requestCmd)

	flags := requestCmd.Flags()
	flags.SortFlags = false

	addRequestFlags(flags)
}

func authmethodAlias(f *pflag.FlagSet, name string) pflag.NormalizedName {
	switch name {
	case "authmethod":
		name = "auth-method"
		break
	}
	return pflag.NormalizedName(name)
}

func addRequestFlags(flags *pflag.FlagSet) {
	flags.StringP("schemes-path", "s", irma.DefaultSchemesPath(), "path to irma_configuration")
	flags.StringP("auth-method", "a", "none", "Authentication method to server (none, token, rsa, hmac)")
	flags.SetNormalizeFunc(authmethodAlias)
	flags.String("key", "", "Key to sign request with")
	flags.String("name", "", "Requestor name")
	flags.StringArray("disclose", nil, "Add an attribute disjunction (comma-separated)")
	flags.StringArray("issue", nil, "Add a credential to issue")
	flags.StringArray("sign", nil, "Add an attribute disjunction to signature session")
	flags.String("message", "", "Message to sign in signature session")
}
<commit_msg>update issue command to use correct amount of missing values<commit_after>package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/go-errors/errors"
	"github.com/mdp/qrterminal"
	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/common"
	"github.com/privacybydesign/irmago/server"
	sseclient "github.com/sietseringers/go-sse"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// requestCmd represents the request command
var requestCmd = &cobra.Command{
	Use:   "request",
	Short: "Generate an IRMA session request",
	Run: func(cmd *cobra.Command, args []string) {
		request, _, err := configureRequest(cmd)
		if err != nil {
			die("", err)
		}

		flags := cmd.Flags()
		authmethod, _ := flags.GetString("authmethod")
		var output string
		if authmethod == "none" || authmethod == "token" {
			output = prettyprint(request)
		} else {
			key, _ := flags.GetString("key")
			name, _ := flags.GetString("name")
			if output, err = signRequest(request, name, authmethod, key); err != nil {
				die("Failed to sign request", err)
			}
		}

		fmt.Println(output)
	},
}

func configureJWTKey(authmethod, key string) (interface{}, jwt.SigningMethod, error) {
	var (
		err    error
		sk     interface{}
		jwtalg jwt.SigningMethod
		bts    []byte
	)
	// If the key refers to an existing file, use contents of the file as key
	if bts, err = common.ReadKey("", key); err != nil {
		bts = []byte(key)
	}
	switch authmethod {
	case "hmac":
		jwtalg = jwt.SigningMethodHS256
		if sk, err = common.Base64Decode(bts); err != nil {
			return nil, nil, err
		}
	case "rsa":
		jwtalg = jwt.SigningMethodRS256
		if sk, err = jwt.ParseRSAPrivateKeyFromPEM(bts); err != nil {
			return nil, nil, err
		}
	default:
		return nil, nil, errors.Errorf("Unsupported signing algorithm: '%s'", authmethod)
	}

	return sk, jwtalg, nil
}

func signRequest(request irma.RequestorRequest, name, authmethod, key string) (string, error) {
	sk, jwtalg, err := configureJWTKey(authmethod, key)
	if err != nil {
		return "", err
	}
	return irma.SignRequestorRequest(request, jwtalg, sk, name)
}

func configureRequest(cmd *cobra.Command) (irma.RequestorRequest, *irma.Configuration, error) {
	irmaconfigPath, err := cmd.Flags().GetString("schemes-path")
	if err != nil {
		return nil, nil, err
	}
	irmaconfig, err := irma.NewConfiguration(irmaconfigPath, irma.ConfigurationOptions{})
	if err != nil {
		return nil, nil, err
	}
	if err = irmaconfig.ParseFolder(); err != nil {
		return nil, nil, err
	}
	if len(irmaconfig.SchemeManagers) == 0 {
		if err = irmaconfig.DownloadDefaultSchemes(); err != nil {
			return nil, nil, err
		}
	}

	request, err := constructSessionRequest(cmd, irmaconfig)
	if err != nil {
		return nil, nil, err
	}

	return request, irmaconfig, nil
}

// Helper functions

func wait(initialStatus server.Status, transport *irma.HTTPTransport, statuschan chan server.Status) {
	events := make(chan *sseclient.Event)

	go func() {
		for {
			if e := <-events; e != nil && e.Type != "open" {
				status := server.Status(strings.Trim(string(e.Data), `"`))
				statuschan <- status
				if status.Finished() {
					return
				}
			}
		}
	}()

	if err := sseclient.Notify(nil, transport.Server+"statusevents", true, events); err != nil {
		fmt.Println("SSE failed, fallback to polling", err)
		close(events)
		poll(initialStatus, transport, statuschan)
		return
	}
}

// poll recursively polls the session status until a final status is received.
func poll(initialStatus server.Status, transport *irma.HTTPTransport, statuschan chan server.Status) {
	// First we wait
	<-time.NewTimer(pollInterval).C

	// Get session status
	var s string
	if err := transport.Get("status", &s); err != nil {
		_ = server.LogFatal(err)
	}
	status := server.Status(strings.Trim(s, `"`))

	// report if status changed
	if status != initialStatus {
		statuschan <- status
	}

	if status.Finished() {
		return
	}
	go poll(status, transport, statuschan)
}

func constructSessionRequest(cmd *cobra.Command, conf *irma.Configuration) (irma.RequestorRequest, error) {
	disclose, _ := cmd.Flags().GetStringArray("disclose")
	issue, _ := cmd.Flags().GetStringArray("issue")
	sign, _ := cmd.Flags().GetStringArray("sign")
	message, _ := cmd.Flags().GetString("message")
	jsonrequest, _ := cmd.Flags().GetString("request")

	if len(disclose) == 0 && len(issue) == 0 && len(sign) == 0 && message == "" {
		if jsonrequest == "" {
			return nil, errors.New("Provide either a complete session request using --request or construct one using the other flags")
		}
		request, err := server.ParseSessionRequest(jsonrequest)
		if err != nil {
			return nil, err
		}
		return request, nil
	}

	if jsonrequest != "" {
		return nil, errors.New("Provide either a complete session request using --request or construct one using the other flags")
	}

	if len(sign) != 0 {
		if len(disclose) != 0 {
			return nil, errors.New("cannot combine disclosure and signature sessions, use either --disclose or --sign")
		}
		if len(issue) != 0 {
			return nil, errors.New("cannot combine issuance and signature sessions, use either --issue or --sign")
		}
		if message == "" {
			return nil, errors.New("signature sessions require a message to be signed using --message")
		}
	}

	var request irma.RequestorRequest
	if len(disclose) != 0 && len(issue) == 0 {
		disclose, err := parseAttrs(disclose, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.ServiceProviderRequest{
			Request: irma.NewDisclosureRequest(),
		}
		request.SessionRequest().(*irma.DisclosureRequest).Disclose = disclose
	}
	if len(sign) != 0 {
		disclose, err := parseAttrs(sign, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.SignatureRequestorRequest{
			Request: irma.NewSignatureRequest(message),
		}
		request.SessionRequest().(*irma.SignatureRequest).Disclose = disclose
	}
	if len(issue) != 0 {
		creds, err := parseCredentials(issue, conf)
		if err != nil {
			return nil, err
		}
		disclose, err := parseAttrs(disclose, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.IdentityProviderRequest{
			Request: irma.NewIssuanceRequest(creds),
		}
		request.SessionRequest().(*irma.IssuanceRequest).Disclose = disclose
	}

	return request, nil
}

func parseCredentials(credentialsStr []string, conf *irma.Configuration) ([]*irma.CredentialRequest, error) {
	list := make([]*irma.CredentialRequest, 0, len(credentialsStr))

	for _, credStr := range credentialsStr {
		parts := strings.Split(credStr, "=")
		if len(parts) != 2 {
			return nil, errors.New("--issue argument must contain exactly 1 = sign")
		}
		credIdStr, attrsStr := parts[0], parts[1]
		credtype := conf.CredentialTypes[irma.NewCredentialTypeIdentifier(credIdStr)]
		if credtype == nil {
			return nil, errors.New("unknown credential type: " + credIdStr)
		}

		// we dont need values for the random blind attributes,
		// since they are randomly generated during issuance
		missingValues := len(credtype.RandomBlinds())
		if len(credtype.RevocationServers) > 0 {
			missingValues += 1
		}
		required := len(credtype.AttributeTypes) - missingValues
		attrsSlice := strings.Split(attrsStr, ",")
		if len(attrsSlice) != required {
			return nil, errors.Errorf("%d attributes required but %d provided for %s", required, len(attrsSlice), credIdStr)
		}

		// Construct the attribute -> value map for in the CredentialRequest
		attrs := make(map[string]string, len(attrsSlice))
		j := 0 // idx in attrSlice
		for _, typ := range credtype.AttributeTypes {
			if !typ.RandomBlind && !typ.RevocationAttribute {
				attrs[typ.ID] = attrsSlice[j]
				j += 1
			}
		}
		list = append(list, &irma.CredentialRequest{
			CredentialTypeID: irma.NewCredentialTypeIdentifier(credIdStr),
			Attributes:       attrs,
		})
	}

	return list, nil
}

func parseAttrs(attrsStr []string, conf *irma.Configuration) (irma.AttributeConDisCon, error) {
	list := make(irma.AttributeConDisCon, 0, len(attrsStr))
	for _, disjunctionStr := range attrsStr {
		disjunction := irma.AttributeDisCon{}
		attrids := strings.Split(disjunctionStr, ",")
		for _, attridStr := range attrids {
			attrid := irma.NewAttributeTypeIdentifier(attridStr)
			if conf.AttributeTypes[attrid] == nil {
				return nil, errors.New("unknown attribute: " + attridStr)
			}
			disjunction = append(disjunction, irma.AttributeCon{irma.AttributeRequest{Type: attrid}})
		}
		list = append(list, disjunction)
	}
	return list, nil
}

func startServer(port int) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", irmaServer.HandlerFunc())
	httpServer = &http.Server{Addr: ":" + strconv.Itoa(port), Handler: mux}
	go func() {
		err := httpServer.ListenAndServe()
		if err != nil && err != http.ErrServerClosed {
			die("Failed to start server", err)
		}
	}()
}

func printQr(qr *irma.Qr, noqr bool) error {
	qrBts, err := json.Marshal(qr)
	if err != nil {
		return err
	}
	if noqr {
		fmt.Println(string(qrBts))
	} else {
		qrterminal.GenerateWithConfig(string(qrBts), qrterminal.Config{
			Level:     qrterminal.L,
			Writer:    os.Stdout,
			BlackChar: qrterminal.BLACK,
			WhiteChar: qrterminal.WHITE,
		})
	}
	return nil
}

func printSessionResult(result *server.SessionResult) {
	fmt.Println("Session result:")
	fmt.Println(prettyprint(result))
}

func init() {
	RootCmd.AddCommand(requestCmd)

	flags := requestCmd.Flags()
	flags.SortFlags = false

	addRequestFlags(flags)
}

func authmethodAlias(f *pflag.FlagSet, name string) pflag.NormalizedName {
	switch name {
	case "authmethod":
		name = "auth-method"
		break
	}
	return pflag.NormalizedName(name)
}

func addRequestFlags(flags *pflag.FlagSet) {
	flags.StringP("schemes-path", "s", irma.DefaultSchemesPath(), "path to irma_configuration")
	flags.StringP("auth-method", "a", "none", "Authentication method to server (none, token, rsa, hmac)")
	flags.SetNormalizeFunc(authmethodAlias)
	flags.String("key", "", "Key to sign request with")
	flags.String("name", "", "Requestor name")
	flags.StringArray("disclose", nil, "Add an attribute disjunction (comma-separated)")
	flags.StringArray("issue", nil, "Add a credential to issue")
	flags.StringArray("sign", nil, "Add an attribute disjunction to signature session")
	flags.String("message", "", "Message to sign in signature session")
}
<|endoftext|>