<commit_before>package mux

import (
	"github.com/oinume/lekcije/server/controller"
	"github.com/oinume/lekcije/server/controller/api"
	"github.com/oinume/lekcije/server/controller/middleware"
	"goji.io"
	"goji.io/pat"
)

func Create() *goji.Mux {
	mux := goji.NewMux()
	mux.UseC(middleware.AccessLogger)
	mux.UseC(middleware.SetDbToContext)
	mux.UseC(middleware.SetLoggedInUserToContext)
	mux.UseC(middleware.LoginRequiredFilter)
	mux.UseC(middleware.CORS)

	mux.HandleFunc(pat.Get("/static/*"), controller.Static)
	mux.HandleFuncC(pat.Get("/"), controller.Index)
	mux.HandleFuncC(pat.Get("/logout"), controller.Logout)
	mux.HandleFuncC(pat.Get("/oauth/google"), controller.OAuthGoogle)
	mux.HandleFuncC(pat.Get("/oauth/google/callback"), controller.OAuthGoogleCallback)
	mux.HandleFuncC(pat.Post("/me/followingTeachers/create"), controller.PostMeFollowingTeachersCreate)
	mux.HandleFuncC(pat.Post("/me/followingTeachers/delete"), controller.PostMeFollowingTeachersDelete)
	mux.HandleFuncC(pat.Get("/me/setting"), controller.GetMeSetting)
	mux.HandleFuncC(pat.Get("/me/setting/update"), controller.PostMeSettingUpdate)
	mux.HandleFuncC(pat.Get("/api/status"), api.GetStatus)
	mux.HandleFuncC(pat.Get("/api/me/followingTeachers"), api.GetMeFollowingTeachers)

	return mux
}
<commit_msg>Get -> Post<commit_after>package mux

import (
	"github.com/oinume/lekcije/server/controller"
	"github.com/oinume/lekcije/server/controller/api"
	"github.com/oinume/lekcije/server/controller/middleware"
	"goji.io"
	"goji.io/pat"
)

func Create() *goji.Mux {
	mux := goji.NewMux()
	mux.UseC(middleware.AccessLogger)
	mux.UseC(middleware.SetDbToContext)
	mux.UseC(middleware.SetLoggedInUserToContext)
	mux.UseC(middleware.LoginRequiredFilter)
	mux.UseC(middleware.CORS)

	mux.HandleFunc(pat.Get("/static/*"), controller.Static)
	mux.HandleFuncC(pat.Get("/"), controller.Index)
	mux.HandleFuncC(pat.Get("/logout"), controller.Logout)
	mux.HandleFuncC(pat.Get("/oauth/google"), controller.OAuthGoogle)
	mux.HandleFuncC(pat.Get("/oauth/google/callback"), controller.OAuthGoogleCallback)
	mux.HandleFuncC(pat.Post("/me/followingTeachers/create"), controller.PostMeFollowingTeachersCreate)
	mux.HandleFuncC(pat.Post("/me/followingTeachers/delete"), controller.PostMeFollowingTeachersDelete)
	mux.HandleFuncC(pat.Get("/me/setting"), controller.GetMeSetting)
	mux.HandleFuncC(pat.Post("/me/setting/update"), controller.PostMeSettingUpdate)
	mux.HandleFuncC(pat.Get("/api/status"), api.GetStatus)
	mux.HandleFuncC(pat.Get("/api/me/followingTeachers"), api.GetMeFollowingTeachers)

	return mux
}
<|endoftext|>


<commit_before>package server

import (
	"time"

	"gopkg.in/mgo.v2/bson"
)

type Team struct {
	Id     bson.ObjectId `json:"id,omitempty" bson:"_id,omitempty"`
	Number int           `json:"number"`
	Name   string        `json:"name"`
	Ip     string        `json:"ip"`
	Hash   string        `json:"-"`
}

type Challenge struct {
	Id          bson.ObjectId `json:"id,omitempty" bson:"_id,omitempty"`
	Group       string        `json:"group"`
	Name        string        `json:"name"`
	Description string        `json:"description"`
	Flag        string        `json:"-" bson:"-"`
	Points      int           `json:"points"`
}

type Result struct {
	Id         bson.ObjectId `json:"id,omitempty" bson:"_id,omitempty"`
	Type       string        `json:"type" bson:"type"`
	Group      string        `json:"group" bson:"group"`
	Teamname   string        `json:"teamname" bson:"teamname"`
	Teamnumber int           `json:"teamnumber" bson:"teamnumber"`
	Details    string        `json:"details" bson:"details"`
	Points     int           `json:"points" bson:"points"`
}

// Authentication Queries
func GetTeamByTeamname(teamname string) (Team, error) {
	t := Team{}

	session, teamCollection := GetSessionAndCollection("teams")
	defer session.Close()

	err := teamCollection.Find(bson.M{"name": teamname}).One(&t)
	if err != nil {
		Logger.Printf("Error finding team by Teamname %s err: %v\n", teamname, err)
		return t, err
	}
	return t, nil
}

func GetTeamById(id *bson.ObjectId) (Team, error) {
	t := Team{}

	session, teamCollection := GetSessionAndCollection("teams")
	defer session.Close()

	err := teamCollection.Find(bson.M{"_id": id}).One(&t)
	if err != nil {
		Logger.Printf("Error finding team by ID %v err: %v\n", id, err)
		return t, err
	}
	return t, nil
}

// Get Team name and ip only used for service checking
func DataGetTeamIps() ([]Team, error) {
	t := []Team{}

	session, teamCollection := GetSessionAndCollection("teams")
	defer session.Close()

	err := teamCollection.Find(nil).Select(bson.M{"_id": false, "name": true, "number": true, "ip": true}).All(&t)
	if err != nil {
		//Logger.Printf("Error finding teams: %v\n", err)
		return t, err
	}
	return t, nil
}

func DataGetTeams() []Team {
	t := []Team{}

	session, chalCollection := GetSessionAndCollection("teams")
	defer session.Close()

	err := chalCollection.Find(nil).Sort("number").Select(bson.M{"_id": 0, "number": 1, "name": 1}).All(&t)
	if err != nil {
		Logger.Printf("Could not get team info: %v\n", err)
		return t
	}
	return t
}

// Query statements
func DataGetChallenges(group string) ([]Challenge, error) {
	challenges := []Challenge{}

	session, chalCollection := GetSessionAndCollection("challenges")
	defer session.Close()

	err := chalCollection.Find(bson.M{"group": group}).Sort("description").Select(bson.M{"_id": 0, "flag": 0}).All(&challenges)
	if err != nil {
		return challenges, err
	}
	return challenges, nil
}

func DataCheckFlag(team Team, chal Challenge) (int, error) {
	session, chalCollection := GetSessionAndCollection("challenges")
	defer session.Close()
	var err error

	if len(chal.Name) > 0 {
		err = chalCollection.Find(bson.M{"flag": chal.Flag, "name": chal.Name}).Select(bson.M{"_id": 0, "flag": 0}).One(&chal)
	} else {
		err = chalCollection.Find(bson.M{"flag": chal.Flag}).Select(bson.M{"_id": 0, "flag": 0}).One(&chal)
	}
	if err != nil {
		// Wrong flag = 1
		return 1, err
	} else {
		if !HasFlag(team.Name, chal.Name) {
			// Correct flag = 0
			result := Result{
				Type:       "CTF",
				Group:      chal.Group,
				Teamname:   team.Name,
				Teamnumber: team.Number,
				Details:    chal.Name,
				Points:     chal.Points,
			}
			return 0, DataAddResult(result)
		} else {
			// Got challenge already
			return 2, nil
		}
	}
}

func HasFlag(teamname, challengeName string) bool {
	chal := Challenge{}

	session, resultCollection := GetSessionAndCollection("results")
	defer session.Close()

	// TODO: Do not need the returned document.
	// Need to find better way to check if exists
	err := resultCollection.Find(bson.M{"teamname": teamname, "details": challengeName}).One(&chal)
	if err != nil {
		// TODO: Log error
		return false
	}
	return true
}

func DataGetTotalChallenges() (map[string]int, error) {
	session, collection := GetSessionAndCollection("challenges")
	defer session.Close()
	totals := make(map[string]int)
	t := bson.M{}

	pipe := collection.Pipe([]bson.M{
		{"$group": bson.M{"_id": "$group", "total": bson.M{"$sum": 1}}},
		{"$sort": bson.M{"_id": -1}},
	})
	iter := pipe.Iter()
	for iter.Next(&t) {
		totals[t["_id"].(string)] = t["total"].(int)
	}
	if err := iter.Close(); err != nil {
		Logger.Printf("Error getting challenges: %v\n", err)
		return totals, err
	}
	return totals, nil
}

func DataGetTeamChallenges(teamname string) (map[string]int, error) {
	session, collection := GetSessionAndCollection("results")
	defer session.Close()
	acquired := make(map[string]int)
	a := bson.M{}
	pipe := collection.Pipe([]bson.M{
		{"$match": bson.M{"teamname": teamname, "type": "CTF"}},
		{"$group": bson.M{"_id": "$group", "acquired": bson.M{"$sum": 1}}},
		{"$sort": bson.M{"_id": -1}},
	})
	iter := pipe.Iter()
	for iter.Next(&a) {
		acquired[a["_id"].(string)] = a["acquired"].(int)
	}
	if err := iter.Close(); err != nil {
		Logger.Printf("Error getting challenges: %v\n", err)
		return acquired, err
	}
	return acquired, nil
}

func DataGetTeamScore(teamname string) int {
	session, collection := GetSessionAndCollection("results")
	defer session.Close()
	points := bson.M{}
	var p int
	pipe := collection.Pipe([]bson.M{
		{"$match": bson.M{"teamname": teamname}},
		{"$group": bson.M{"_id": nil, "points": bson.M{"$sum": "$points"}}},
		{"$project": bson.M{"_id": 0, "points": 1}},
	})
	err := pipe.One(&points)
	if err != nil {
		Logger.Printf("Error getting team points: %v\n", err)
	} else {
		p = points["points"].(int)
	}
	return p
}

func DataGetAllScore() []Result {
	session, collection := GetSessionAndCollection("results")
	defer session.Close()
	tmScore := []Result{}

	pipe := collection.Pipe([]bson.M{
		{"$group": bson.M{"_id": bson.M{"tname": "$teamname", "tnum": "$teamnumber"}, "points": bson.M{"$sum": "$points"}}},
		{"$project": bson.M{"_id": 0, "points": 1, "teamnumber": "$_id.tnum", "teamname": "$_id.tname"}},
		{"$sort": bson.M{"teamnumber": 1}},
	})
	err := pipe.All(&tmScore)
	if err != nil {
		Logger.Printf("Error getting all team scores: %v\n", err)
	}
	if len(tmScore) == 0 {
		teams := DataGetTeams()
		for _, t := range teams {
			tmScore = append(tmScore, Result{Teamname: t.Name, Teamnumber: t.Number, Points: 0})
		}
	}
	return tmScore
}

func DataGetServiceStatus() []Result {
	r := []Result{}

	session, results := GetSessionAndCollection("results")
	defer session.Close()

	err := results.Find(bson.M{"type": "Service"}).All(&r)
	if err != nil {
		Logger.Printf("Error getting service status: %v\n", err)
	}
	return r
}

func DataGetLastServiceResult() time.Time {
	session, results := GetSessionAndCollection("results")
	defer session.Close()
	id := bson.M{}
	pipe := results.Pipe([]bson.M{
		{"$match": bson.M{"type": "Service"}},
		{"$sort": bson.M{"_id": 1}},
		{"$group": bson.M{"_id": nil, "last": bson.M{"$last": "$_id"}}},
		{"$project": bson.M{"_id": 0, "last": 1}},
	})
	err := pipe.One(&id)
	if err != nil {
		Logger.Printf("Error getting last Service result: %v\n", err)
	}
	time := id["last"].(bson.ObjectId).Time()
	return time
}

func DataGetLastResult() time.Time {
	session, results := GetSessionAndCollection("results")
	defer session.Close()
	id := bson.M{}
	pipe := results.Pipe([]bson.M{
		{"$sort": bson.M{"_id": 1}},
		{"$group": bson.M{"_id": nil, "last": bson.M{"$last": "$_id"}}},
		{"$project": bson.M{"_id": 0, "last": 1}},
	})
	err := pipe.One(&id)
	var time time.Time
	if err != nil {
		Logger.Printf("Error getting last document: %v\n", err)
	} else {
		time = id["last"].(bson.ObjectId).Time()
	}
	return time
}

// Insert statements
func DataAddResult(result Result) error {
	session, collection := GetSessionAndCollection("results")
	defer session.Close()
	err := collection.Insert(result)
	if err != nil {
		//Logger.Printf("Error inserting %s to team %s: %v", result.Details, result.Teamname, err)
		return err
	}
	return nil
}
<commit_msg>Fix scoreboard not displaying all team score when there are no results<commit_after>package server

import (
	"time"

	"gopkg.in/mgo.v2/bson"
)

type Team struct {
	Id     bson.ObjectId `json:"id,omitempty" bson:"_id,omitempty"`
	Number int           `json:"number"`
	Name   string        `json:"name"`
	Ip     string        `json:"ip"`
	Hash   string        `json:"-"`
}

type Challenge struct {
	Id          bson.ObjectId `json:"id,omitempty" bson:"_id,omitempty"`
	Group       string        `json:"group"`
	Name        string        `json:"name"`
	Description string        `json:"description"`
	Flag        string        `json:"-" bson:"-"`
	Points      int           `json:"points"`
}

type Result struct {
	Id         bson.ObjectId `json:"id,omitempty" bson:"_id,omitempty"`
	Type       string        `json:"type" bson:"type"`
	Group      string        `json:"group" bson:"group"`
	Teamname   string        `json:"teamname" bson:"teamname"`
	Teamnumber int           `json:"teamnumber" bson:"teamnumber"`
	Details    string        `json:"details" bson:"details"`
	Points     int           `json:"points" bson:"points"`
}

// Authentication Queries
func GetTeamByTeamname(teamname string) (Team, error) {
	t := Team{}

	session, teamCollection := GetSessionAndCollection("teams")
	defer session.Close()

	err := teamCollection.Find(bson.M{"name": teamname}).One(&t)
	if err != nil {
		Logger.Printf("Error finding team by Teamname %s err: %v\n", teamname, err)
		return t, err
	}
	return t, nil
}

func GetTeamById(id *bson.ObjectId) (Team, error) {
	t := Team{}

	session, teamCollection := GetSessionAndCollection("teams")
	defer session.Close()

	err := teamCollection.Find(bson.M{"_id": id}).One(&t)
	if err != nil {
		Logger.Printf("Error finding team by ID %v err: %v\n", id, err)
		return t, err
	}
	return t, nil
}

// Get Team name and ip only used for service checking
func DataGetTeamIps() ([]Team, error) {
	t := []Team{}

	session, teamCollection := GetSessionAndCollection("teams")
	defer session.Close()

	err := teamCollection.Find(nil).Select(bson.M{"_id": false, "name": true, "number": true, "ip": true}).All(&t)
	if err != nil {
		//Logger.Printf("Error finding teams: %v\n", err)
		return t, err
	}
	return t, nil
}

func DataGetTeams() []Team {
	t := []Team{}

	session, chalCollection := GetSessionAndCollection("teams")
	defer session.Close()

	err := chalCollection.Find(nil).Sort("number").Select(bson.M{"_id": 0, "number": 1, "name": 1}).All(&t)
	if err != nil {
		Logger.Printf("Could not get team info: %v\n", err)
		return t
	}
	return t
}

// Query statements
func DataGetChallenges(group string) ([]Challenge, error) {
	challenges := []Challenge{}

	session, chalCollection := GetSessionAndCollection("challenges")
	defer session.Close()

	err := chalCollection.Find(bson.M{"group": group}).Sort("description").Select(bson.M{"_id": 0, "flag": 0}).All(&challenges)
	if err != nil {
		return challenges, err
	}
	return challenges, nil
}

func DataCheckFlag(team Team, chal Challenge) (int, error) {
	session, chalCollection := GetSessionAndCollection("challenges")
	defer session.Close()
	var err error

	if len(chal.Name) > 0 {
		err = chalCollection.Find(bson.M{"flag": chal.Flag, "name": chal.Name}).Select(bson.M{"_id": 0, "flag": 0}).One(&chal)
	} else {
		err = chalCollection.Find(bson.M{"flag": chal.Flag}).Select(bson.M{"_id": 0, "flag": 0}).One(&chal)
	}
	if err != nil {
		// Wrong flag = 1
		return 1, err
	} else {
		if !HasFlag(team.Name, chal.Name) {
			// Correct flag = 0
			result := Result{
				Type:       "CTF",
				Group:      chal.Group,
				Teamname:   team.Name,
				Teamnumber: team.Number,
				Details:    chal.Name,
				Points:     chal.Points,
			}
			return 0, DataAddResult(result)
		} else {
			// Got challenge already
			return 2, nil
		}
	}
}

func HasFlag(teamname, challengeName string) bool {
	chal := Challenge{}

	session, resultCollection := GetSessionAndCollection("results")
	defer session.Close()

	// TODO: Do not need the returned document.
	// Need to find better way to check if exists
	err := resultCollection.Find(bson.M{"teamname": teamname, "details": challengeName}).One(&chal)
	if err != nil {
		// TODO: Log error
		return false
	}
	return true
}

func DataGetTotalChallenges() (map[string]int, error) {
	session, collection := GetSessionAndCollection("challenges")
	defer session.Close()
	totals := make(map[string]int)
	t := bson.M{}
	pipe := collection.Pipe([]bson.M{
		{"$group": bson.M{"_id": "$group", "total": bson.M{"$sum": 1}}},
		{"$sort": bson.M{"_id": -1}},
	})
	iter := pipe.Iter()
	for iter.Next(&t) {
		totals[t["_id"].(string)] = t["total"].(int)
	}
	if err := iter.Close(); err != nil {
		Logger.Printf("Error getting challenges: %v\n", err)
		return totals, err
	}
	return totals, nil
}

func DataGetTeamChallenges(teamname string) (map[string]int, error) {
	session, collection := GetSessionAndCollection("results")
	defer session.Close()
	acquired := make(map[string]int)
	a := bson.M{}
	pipe := collection.Pipe([]bson.M{
		{"$match": bson.M{"teamname": teamname, "type": "CTF"}},
		{"$group": bson.M{"_id": "$group", "acquired": bson.M{"$sum": 1}}},
		{"$sort": bson.M{"_id": -1}},
	})
	iter := pipe.Iter()
	for iter.Next(&a) {
		acquired[a["_id"].(string)] = a["acquired"].(int)
	}
	if err := iter.Close(); err != nil {
		Logger.Printf("Error getting challenges: %v\n", err)
		return acquired, err
	}
	return acquired, nil
}

func DataGetTeamScore(teamname string) int {
	session, collection := GetSessionAndCollection("results")
	defer session.Close()
	points := bson.M{}
	var p int
	pipe := collection.Pipe([]bson.M{
		{"$match": bson.M{"teamname": teamname}},
		{"$group": bson.M{"_id": nil, "points": bson.M{"$sum": "$points"}}},
		{"$project": bson.M{"_id": 0, "points": 1}},
	})
	err := pipe.One(&points)
	if err != nil {
		Logger.Printf("Error getting team points: %v\n", err)
	} else {
		p = points["points"].(int)
	}
	return p
}

func DataGetAllScore() []Result {
	session, collection := GetSessionAndCollection("results")
	defer session.Close()
	tmScore := []Result{}

	teams := DataGetTeams()
	pipe := collection.Pipe([]bson.M{
		{"$group": bson.M{"_id": bson.M{"tname": "$teamname", "tnum": "$teamnumber"}, "points": bson.M{"$sum": "$points"}}},
		{"$project": bson.M{"_id": 0, "points": 1, "teamnumber": "$_id.tnum", "teamname": "$_id.tname"}},
		{"$sort": bson.M{"teamnumber": 1}},
	})
	err := pipe.All(&tmScore)
	if err != nil {
		Logger.Printf("Error getting all team scores: %v\n", err)
	}
	if l := len(tmScore); l < len(teams) {
		if l == 0 {
			for _, t := range teams {
				tmScore = append(tmScore, Result{Teamname: t.Name, Teamnumber: t.Number, Points: 0})
			}
		} else {
			for i := l; i < len(teams); i++ {
				tmScore = append(tmScore, Result{Teamname: teams[i].Name, Teamnumber: teams[i].Number, Points: 0})
			}
		}
	}
	return tmScore
}

func DataGetServiceStatus() []Result {
	r := []Result{}

	session, results := GetSessionAndCollection("results")
	defer session.Close()

	err := results.Find(bson.M{"type": "Service"}).All(&r)
	if err != nil {
		Logger.Printf("Error getting service status: %v\n", err)
	}
	return r
}

func DataGetLastServiceResult() time.Time {
	session, results := GetSessionAndCollection("results")
	defer session.Close()
	id := bson.M{}
	pipe := results.Pipe([]bson.M{
		{"$match": bson.M{"type": "Service"}},
		{"$sort": bson.M{"_id": 1}},
		{"$group": bson.M{"_id": nil, "last": bson.M{"$last": "$_id"}}},
		{"$project": bson.M{"_id": 0, "last": 1}},
	})
	err := pipe.One(&id)
	if err != nil {
		Logger.Printf("Error getting last Service result: %v\n", err)
	}
	time := id["last"].(bson.ObjectId).Time()
	return time
}

func DataGetLastResult() time.Time {
	session, results := GetSessionAndCollection("results")
	defer session.Close()
	id := bson.M{}
	pipe := results.Pipe([]bson.M{
		{"$sort": bson.M{"_id": 1}},
		{"$group": bson.M{"_id": nil, "last": bson.M{"$last": "$_id"}}},
		{"$project": bson.M{"_id": 0, "last": 1}},
	})
	err := pipe.One(&id)
	var time time.Time
	if err != nil {
		Logger.Printf("Error getting last document: %v\n", err)
	} else {
		time = id["last"].(bson.ObjectId).Time()
	}
	return time
}

// Insert statements
func DataAddResult(result Result) error {
	session, collection := GetSessionAndCollection("results")
	defer session.Close()
	err := collection.Insert(result)
	if err != nil {
		//Logger.Printf("Error inserting %s to team %s: %v", result.Details, result.Teamname, err)
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"
	"time"

	"github.com/ninjasphere/go-ninja/channels"
	"github.com/ninjasphere/go-zigbee/gateway"
)

type ColorChannel struct {
	Channel
	channel *channels.ColorChannel
}

// -------- Color Protocol --------

func (c *ColorChannel) init() error {
	log.Debugf("Initialising color channel of device %d", *c.device.deviceInfo.IeeeAddress)

	/*clusterID := uint32(0x06)
	  attributeID := uint32(0)
	  minReportInterval := uint32(1)
	  maxReportInterval := uint32(120)

	  request := &gateway.GwSetAttributeReportingReq{
	    DstAddress: &gateway.GwAddressStructT{
	      AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
	      IeeeAddr:    c.device.deviceInfo.IeeeAddress,
	    },
	    ClusterId: &clusterID,
	    AttributeReportList: []*gateway.GwAttributeReportT{{
	      AttributeId:       &attributeID,
	      AttributeType:     gateway.GwZclAttributeDataTypesT_ZCL_DATATYPE_BOOLEAN.Enum(),
	      MinReportInterval: &minReportInterval,
	      MaxReportInterval: &maxReportInterval,
	    }},
	  }

	  response := &gateway.GwSetAttributeReportingRspInd{}

	  err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 20*time.Second)
	  if err != nil {
	    log.Errorf("Error enabling on/off reporting: %s", err)
	  } else if response.Status.String() != "STATUS_SUCCESS" {
	    log.Errorf("Failed to enable on/off reporting. status: %s", response.Status.String())
	  }*/

	//mosquitto_pub -m '{"id":123, "params": [0.1],"jsonrpc": "2.0","method":"set","time":132123123}' -t '$device/26b4f71484/channel/11-8'

	c.channel = channels.NewColorChannel(c)
	err := c.device.driver.Conn.ExportChannel(c.device, c.channel, c.ID)
	if err != nil {
		log.Fatalf("Failed to announce color channel: %s", err)
	}

	go func() {
		for {
			log.Debugf("Polling for color")
			err := c.fetchState()
			if err != nil {
				log.Errorf("Failed to poll for color state %s", err)
			}
			time.Sleep(10 * time.Second)
		}
	}()

	return nil

}

func (c *ColorChannel) SetColor(state *channels.ColorState) error {

	if state.Mode != "hue" {
		return fmt.Errorf("TODO: Only color mode 'hue' is supported atm.")
	}

	hue := uint32(*state.Hue * float64(math.MaxUint8))
	saturation := uint32(*state.Saturation * float64(math.MaxUint8))

	request := &gateway.DevSetColorReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
		HueValue:        &hue,
		SaturationValue: &saturation,
	}

	response := &gateway.GwZigbeeGenericRspInd{}
	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 2*time.Second)
	if err != nil {
		return fmt.Errorf("Error setting color state : %s", err)
	}
	if response.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to set color state. status: %s", response.Status.String())
	}

	return c.fetchState()
}

func (c *ColorChannel) fetchState() error {
	request := &gateway.DevGetColorReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
	}

	response := &gateway.DevGetColorRspInd{}
	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 10*time.Second)
	if err != nil {
		return fmt.Errorf("Error getting color state : %s", err)
	}
	if response.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to get color state. status: %s", response.Status.String())
	}

	saturation := float64(float64(*response.SatValue) / float64(math.MaxUint8))
	hue := float64(float64(*response.HueValue) / float64(math.MaxUint8))

	state := &channels.ColorState{
		Mode:       "Hue",
		Saturation: &saturation,
		Hue:        &hue,
	}

	c.channel.SendState(state)

	return nil
}
<commit_msg>Color: hue+sat values are 0-254 inclusive, not 255<commit_after>package main

import (
	"fmt"
	"math"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/ninjasphere/go-ninja/channels"
	"github.com/ninjasphere/go-zigbee/gateway"
)

type ColorChannel struct {
	Channel
	channel *channels.ColorChannel
}

// -------- Color Protocol --------

func (c *ColorChannel) init() error {
	log.Debugf("Initialising color channel of device %d", *c.device.deviceInfo.IeeeAddress)

	/*clusterID := uint32(0x06)
	  attributeID := uint32(0)
	  minReportInterval := uint32(1)
	  maxReportInterval := uint32(120)

	  request := &gateway.GwSetAttributeReportingReq{
	    DstAddress: &gateway.GwAddressStructT{
	      AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
	      IeeeAddr:    c.device.deviceInfo.IeeeAddress,
	    },
	    ClusterId: &clusterID,
	    AttributeReportList: []*gateway.GwAttributeReportT{{
	      AttributeId:       &attributeID,
	      AttributeType:     gateway.GwZclAttributeDataTypesT_ZCL_DATATYPE_BOOLEAN.Enum(),
	      MinReportInterval: &minReportInterval,
	      MaxReportInterval: &maxReportInterval,
	    }},
	  }

	  response := &gateway.GwSetAttributeReportingRspInd{}

	  err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 20*time.Second)
	  if err != nil {
	    log.Errorf("Error enabling on/off reporting: %s", err)
	  } else if response.Status.String() != "STATUS_SUCCESS" {
	    log.Errorf("Failed to enable on/off reporting. status: %s", response.Status.String())
	  }*/

	//mosquitto_pub -m '{"id":123, "params": [0.1],"jsonrpc": "2.0","method":"set","time":132123123}' -t '$device/26b4f71484/channel/11-8'

	c.channel = channels.NewColorChannel(c)
	err := c.device.driver.Conn.ExportChannel(c.device, c.channel, c.ID)
	if err != nil {
		log.Fatalf("Failed to announce color channel: %s", err)
	}

	go func() {
		for {
			log.Debugf("Polling for color")
			err := c.fetchState()
			if err != nil {
				log.Errorf("Failed to poll for color state %s", err)
			}
			time.Sleep(10 * time.Second)
		}
	}()

	return nil

}

func (c *ColorChannel) SetColor(state *channels.ColorState) error {

	if state.Mode != "hue" {
		return fmt.Errorf("TODO: Only color mode 'hue' is supported atm.")
	}

	spew.Dump("setting color", state)

	hue := uint32(*state.Hue * float64(math.MaxUint8-1))
	saturation := uint32(*state.Saturation * float64(math.MaxUint8-1))

	request := &gateway.DevSetColorReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
		HueValue:        &hue,
		SaturationValue: &saturation,
	}

	spew.Dump(request)

	response := &gateway.GwZigbeeGenericRspInd{}
	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 2*time.Second)
	if err != nil {
		return fmt.Errorf("Error setting color state : %s", err)
	}
	if response.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to set color state. status: %s", response.Status.String())
	}

	return c.fetchState()
}

func (c *ColorChannel) fetchState() error {
	request := &gateway.DevGetColorReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
	}

	response := &gateway.DevGetColorRspInd{}
	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 10*time.Second)
	if err != nil {
		return fmt.Errorf("Error getting color state : %s", err)
	}
	if response.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to get color state. status: %s", response.Status.String())
	}

	saturation := float64(float64(*response.SatValue) / float64(math.MaxUint8-1))
	hue := float64(float64(*response.HueValue) / float64(math.MaxUint8-1))

	state := &channels.ColorState{
		Mode:       "Hue",
		Saturation: &saturation,
		Hue:        &hue,
	}

	c.channel.SendState(state)

	return nil
}
<|endoftext|>


<commit_before>// Copyright © 2016 Jip J. Dekker <jip@dekker.li>
//
// 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 settings

// Score represents the settings for a specific score file
type Score struct {
	Name       string   // The name of the score in the songbook
	Categories []string // Categories to which the scores belong
	Path       string   // The path to the scores (uncompiled) file
}
<commit_msg>Adds function to read a score structure from a JSON file<commit_after>// Copyright © 2016 Jip J. Dekker <jip@dekker.li>
//
// 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 settings

import (
	"encoding/json"
	"io/ioutil"
)

// Score represents the settings for a specific score file
type Score struct {
	Name       string   // The name of the score in the songbook
	Categories []string // Categories to which the scores belong
	Path       string   // The path to the scores (uncompiled) file
}

// FromJSON reads the settings of a score from a JSON file
func FromJSON(path string) (*Score, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	var s Score
	err = json.Unmarshal(data, &s)
	if err != nil {
		return nil, err
	}
	return &s, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. 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 integration

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"testing"
	"time"

	heapster_api "github.com/GoogleCloudPlatform/heapster/api/v1"
	sink_api "github.com/GoogleCloudPlatform/heapster/sinks/api/v1"
	"github.com/GoogleCloudPlatform/heapster/sinks/cache"
	kube_api "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/golang/glog"
	"github.com/stretchr/testify/require"
)

const (
	kTestZone        = "us-central1-b"
	targetTags       = "kubernetes-minion"
	heapsterBuildDir = "../deploy/docker"
	makefile         = "../Makefile"
)

var (
	kubeVersions           = flag.String("kube_versions", "", "Comma separated list of kube versions to test against. By default will run the test against an existing cluster")
	heapsterControllerFile = flag.String("heapster_controller", "../deploy/kube-config/standalone/heapster-controller.json", "Path to heapster replication controller file.")
	heapsterServiceFile    = flag.String("heapster_service", "../deploy/kube-config/standalone/heapster-service.json", "Path to heapster service file.")
	heapsterImage          = flag.String("heapster_image", "heapster:e2e_test", "heapster docker image that needs to be tested.")
	avoidBuild             = flag.Bool("nobuild", false, "When true, a new heapster docker image will not be created and pushed to test cluster nodes.")
	namespace              = flag.String("namespace", "default", "namespace to be used for testing")
	maxRetries             = flag.Int("retries", 100, "Number of attempts before failing this test.")
	runForever             = flag.Bool("run_forever", false, "If true, the tests are run in a loop forever.")
)

func deleteAll(fm kubeFramework, ns string, service *kube_api.Service, rc *kube_api.ReplicationController) {
	if err := fm.DeleteRC(ns, rc); err != nil {
		glog.Error(err)
	}

	if err := fm.DeleteService(ns, service); err != nil {
		glog.Error(err)
	}
}

func createAll(fm kubeFramework, ns string, service **kube_api.Service, rc **kube_api.ReplicationController) error {
	if newRc, err := fm.CreateRC(ns, *rc); err != nil {
		return err
	} else {
		*rc = newRc
	}

	if newSvc, err := fm.CreateService(ns, *service); err != nil {
		return err
	} else {
		*service = newSvc
	}
	return nil
}

func removeHeapsterImage(fm kubeFramework) error {
	if err := removeDockerImage(*heapsterImage); err != nil {
		glog.Error(err)
	}
	if nodes, err := fm.GetNodes(); err == nil {
		for _, node := range nodes {
			host := strings.Split(node, ".")[0]
			cleanupRemoteHost(host, kTestZone)
		}
	} else {
		glog.Errorf("failed to cleanup nodes - %v", err)
	}
	return nil
}

func buildAndPushHeapsterImage(hostnames []string) error {
	curwd, err := os.Getwd()
	if err != nil {
		return err
	}
	if err := os.Chdir(heapsterBuildDir); err != nil {
		return err
	}
	if err := make("build", path.Join(curwd, makefile)); err != nil {
		return err
	}
	if err := buildDockerImage(*heapsterImage); err != nil {
		return err
	}
	for _, host := range hostnames {
		if err := copyDockerImage(*heapsterImage, host, kTestZone); err != nil {
			return err
		}
	}
	glog.V(2).Info("built and pushed heapster image")
	return os.Chdir(curwd)
}

func getHeapsterRcAndSvc(fm kubeFramework) (*kube_api.Service, *kube_api.ReplicationController, error) {
	// Add test docker image
	rc, err := fm.ParseRC(*heapsterControllerFile)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to parse heapster controller - %v", err)
	}
	rc.Spec.Template.Spec.Containers[0].Image = *heapsterImage
	rc.Spec.Template.Spec.Containers[0].ImagePullPolicy = kube_api.PullNever
	// increase logging level
	rc.Spec.Template.Spec.Containers[0].Env = append(rc.Spec.Template.Spec.Containers[0].Env, kube_api.EnvVar{Name: "FLAGS", Value: "--vmodule=*=3"})
	svc, err := fm.ParseService(*heapsterServiceFile)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to parse heapster service - %v", err)
	}

	return svc, rc, nil
}

func buildAndPushDockerImages(fm kubeFramework) error {
	if *avoidBuild {
		return nil
	}
	nodes, err := fm.GetNodes()
	if err != nil {
		return err
	}
	hostnames := []string{}
	for _, node := range nodes {
		hostnames = append(hostnames, strings.Split(node, ".")[0])
	}

	return buildAndPushHeapsterImage(hostnames)
}

const (
	metricsEndpoint       = "/api/v1/metric-export"
	metricsSchemaEndpoint = "/api/v1/metric-export-schema"
)

func runHeapsterMetricsTest(fm kubeFramework, svc *kube_api.Service, expectedNodes, expectedPods []string) error {
	body, err := fm.Client().Get().
		Namespace(svc.Namespace).
		Prefix("proxy").
		Resource("services").
		Name(svc.Name).
		Suffix(metricsEndpoint).
		Do().Raw()
	if err != nil {
		return err
	}
	var timeseries []*heapster_api.Timeseries
	if err := json.Unmarshal(body, &timeseries); err != nil {
		glog.V(2).Infof("response body: %v", string(body))
		return err
	}

	if len(timeseries) == 0 {
		return fmt.Errorf("expected non zero timeseries")
	}
	body, err = fm.Client().Get().
		Namespace(svc.Namespace).
		Prefix("proxy").
		Resource("services").
		Name(svc.Name).
		Suffix(metricsSchemaEndpoint).
		Do().Raw()
	if err != nil {
		return err
	}
	var timeseriesSchema heapster_api.TimeseriesSchema
	if err := json.Unmarshal(body, &timeseriesSchema); err != nil {
		glog.V(2).Infof("response body: %v", string(body))
		return err
	}
	// Build a map of metric names to metric descriptors.
	mdMap := map[string]*heapster_api.MetricDescriptor{}
	for idx := range timeseriesSchema.Metrics {
		mdMap[timeseriesSchema.Metrics[idx].Name] = &timeseriesSchema.Metrics[idx]
	}
	actualPods := map[string]bool{}
	actualNodes := map[string]bool{}
	for _, ts := range timeseries {
		// Verify the relevant labels are present.
		// All common labels must be present.
		for _, label := range sink_api.CommonLabels() {
			_, exists := ts.Labels[label.Key]
			if !exists {
				return fmt.Errorf("timeseries: %v does not contain common label: %v", ts, label)
			}
		}
		podName, podMetric := ts.Labels[sink_api.LabelPodName.Key]
		if podMetric {
			for _, label := range sink_api.PodLabels() {
				_, exists := ts.Labels[label.Key]
				if !exists {
					return fmt.Errorf("timeseries: %v does not contain pod label: %v", ts, label)
				}
			}
		}
		if podMetric {
			actualPods[podName] = true

		} else if cName, ok := ts.Labels[sink_api.LabelContainerName.Key]; ok && cName == cache.NodeContainerName {
			hostname, ok := ts.Labels[sink_api.LabelHostname.Key]
			if !ok {
				return fmt.Errorf("hostname label missing on node container %+v", ts.Labels)
			}
			actualNodes[hostname] = true
		}

		for metricName, points := range ts.Metrics {
			md, exists := mdMap[metricName]
			if !exists {
				return fmt.Errorf("unexpected metric %q", metricName)
			}
			for _, point := range points {
				for _, label := range md.Labels {
					_, exists := point.Labels[label.Key]
					if !exists {
						return fmt.Errorf("metric %q point %v does not contain metric label: %v", metricName, point, label)
					}
				}
			}

		}
	}
	if !expectedItemsExist(expectedPods, actualPods) {
		return fmt.Errorf("expected pods don't exist.\nExpected: %v\nActual:%v", expectedPods, actualPods)
	}
	if !expectedItemsExist(expectedNodes, actualNodes) {
		return fmt.Errorf("expected nodes don't exist.\nExpected: %v\nActual:%v", expectedNodes, actualNodes)
	}

	return nil
}

func expectedItemsExist(expectedItems []string, actualItems map[string]bool) bool {
	if len(actualItems) < len(expectedItems) {
		return false
	}
	for _, item := range expectedItems {
		if _, found := actualItems[item]; !found {
			return false
		}
	}
	return true
}

func apiTest(kubeVersion string) error {
	fm, err := newKubeFramework(kubeVersion)
	if err != nil {
		return err
	}
	if err := buildAndPushDockerImages(fm); err != nil {
		return err
	}
	// Create heapster pod and service.
	svc, rc, err := getHeapsterRcAndSvc(fm)
	if err != nil {
		return err
	}
	ns := *namespace
	deleteAll(fm, ns, svc, rc)
	if err = createAll(fm, ns, &svc, &rc); err != nil {
		return err
	}
	defer func() {
		deleteAll(fm, ns, svc, rc)
		removeHeapsterImage(fm)
	}()
	if err := fm.WaitUntilPodRunning(ns, rc.Spec.Template.Labels, time.Minute); err != nil {
		return err
	}
	if err := fm.WaitUntilServiceActive(svc, time.Minute); err != nil {
		return err
	}
	expectedPods, err := fm.GetPods()
	if err != nil {
		return err
	}
	expectedNodes, err := fm.GetNodes()
	if err != nil {
		return err
	}
	attempts := *maxRetries
	for {
		err := runHeapsterMetricsTest(fm, svc, expectedNodes, expectedPods)
		if *runForever {
			continue
		}
		if err == nil {
			return nil
		}
		if attempts == 0 {
			return err
		}
		attempts--
		time.Sleep(time.Second)
	}
}

func runApiTest() error {
	tempDir, err := ioutil.TempDir("", "deploy")
	if err != nil {
		return nil
	}
	defer os.RemoveAll(tempDir)
	if *kubeVersions == "" {
		return apiTest("")
	}
	kubeVersionsList := strings.Split(*kubeVersions, ",")
	for _, kubeVersion := range kubeVersionsList {
		return apiTest(kubeVersion)
	}
	return nil
}

func TestHeapster(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping heapster kubernetes integration test.")
	}
	require.NoError(t, runApiTest())
}
<commit_msg>Updating integration tests to validate that system contaienrs are present on all the cluster hosts.<commit_after>// Copyright 2015 Google Inc. 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 integration

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"testing"
	"time"

	heapster_api "github.com/GoogleCloudPlatform/heapster/api/v1"
	sink_api "github.com/GoogleCloudPlatform/heapster/sinks/api/v1"
	"github.com/GoogleCloudPlatform/heapster/sinks/cache"
	kube_api "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/golang/glog"
	"github.com/stretchr/testify/require"
)

const (
	kTestZone        = "us-central1-b"
	targetTags       = "kubernetes-minion"
	heapsterBuildDir = "../deploy/docker"
	makefile         = "../Makefile"
)

var (
	kubeVersions           = flag.String("kube_versions", "", "Comma separated list of kube versions to test against. By default will run the test against an existing cluster")
	heapsterControllerFile = flag.String("heapster_controller", "../deploy/kube-config/standalone/heapster-controller.json", "Path to heapster replication controller file.")
	heapsterServiceFile    = flag.String("heapster_service", "../deploy/kube-config/standalone/heapster-service.json", "Path to heapster service file.")
	heapsterImage          = flag.String("heapster_image", "heapster:e2e_test", "heapster docker image that needs to be tested.")
	avoidBuild             = flag.Bool("nobuild", false, "When true, a new heapster docker image will not be created and pushed to test cluster nodes.")
	namespace              = flag.String("namespace", "default", "namespace to be used for testing")
	maxRetries             = flag.Int("retries", 100, "Number of attempts before failing this test.")
	runForever             = flag.Bool("run_forever", false, "If true, the tests are run in a loop forever.")
)

func deleteAll(fm kubeFramework, ns string, service *kube_api.Service, rc *kube_api.ReplicationController) {
	if err := fm.DeleteRC(ns, rc); err != nil {
		glog.Error(err)
	}

	if err := fm.DeleteService(ns, service); err != nil {
		glog.Error(err)
	}
}

func createAll(fm kubeFramework, ns string, service **kube_api.Service, rc **kube_api.ReplicationController) error {
	if newRc, err := fm.CreateRC(ns, *rc); err != nil {
		return err
	} else {
		*rc = newRc
	}

	if newSvc, err := fm.CreateService(ns, *service); err != nil {
		return err
	} else {
		*service = newSvc
	}
	return nil
}

func removeHeapsterImage(fm kubeFramework) error {
	if err := removeDockerImage(*heapsterImage); err != nil {
		glog.Error(err)
	}
	if nodes, err := fm.GetNodes(); err == nil {
		for _, node := range nodes {
			host := strings.Split(node, ".")[0]
			cleanupRemoteHost(host, kTestZone)
		}
	} else {
		glog.Errorf("failed to cleanup nodes - %v", err)
	}
	return nil
}

func buildAndPushHeapsterImage(hostnames []string) error {
	curwd, err := os.Getwd()
	if err != nil {
		return err
	}
	if err := os.Chdir(heapsterBuildDir); err != nil {
		return err
	}
	if err := make("build", path.Join(curwd, makefile)); err != nil {
		return err
	}
	if err := buildDockerImage(*heapsterImage); err != nil {
		return err
	}
	for _, host := range hostnames {
		if err := copyDockerImage(*heapsterImage, host, kTestZone); err != nil {
			return err
		}
	}
	glog.V(2).Info("built and pushed heapster image")
	return os.Chdir(curwd)
}

func getHeapsterRcAndSvc(fm kubeFramework) (*kube_api.Service, *kube_api.ReplicationController, error) {
	// Add test docker image
	rc, err := fm.ParseRC(*heapsterControllerFile)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to parse heapster controller - %v", err)
	}
	rc.Spec.Template.Spec.Containers[0].Image = *heapsterImage
	rc.Spec.Template.Spec.Containers[0].ImagePullPolicy = kube_api.PullNever
	// increase logging level
	rc.Spec.Template.Spec.Containers[0].Env = append(rc.Spec.Template.Spec.Containers[0].Env, kube_api.EnvVar{Name: "FLAGS", Value: "--vmodule=*=3"})
	svc, err := fm.ParseService(*heapsterServiceFile)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to parse heapster service - %v", err)
	}

	return svc, rc, nil
}

func buildAndPushDockerImages(fm kubeFramework) error {
	if *avoidBuild {
		return nil
	}
	nodes, err := fm.GetNodes()
	if err != nil {
		return err
	}
	hostnames := []string{}
	for _, node := range nodes {
		hostnames = append(hostnames, strings.Split(node, ".")[0])
	}

	return buildAndPushHeapsterImage(hostnames)
}

const (
	metricsEndpoint       = "/api/v1/metric-export"
	metricsSchemaEndpoint = "/api/v1/metric-export-schema"
)

func getTimeseries(fm kubeFramework, svc *kube_api.Service) ([]*heapster_api.Timeseries, error) {
	body, err := fm.Client().Get().
		Namespace(svc.Namespace).
		Prefix("proxy").
		Resource("services").
		Name(svc.Name).
		Suffix(metricsEndpoint).
		Do().Raw()
	if err != nil {
		return nil, err
	}
	var timeseries []*heapster_api.Timeseries
	if err := json.Unmarshal(body, &timeseries); err != nil {
		glog.V(2).Infof("response body: %v", string(body))
		return nil, err
	}
	return timeseries, nil
}

func getSchema(fm kubeFramework, svc *kube_api.Service) (*heapster_api.TimeseriesSchema, error) {
	body, err := fm.Client().Get().
		Namespace(svc.Namespace).
		Prefix("proxy").
		Resource("services").
		Name(svc.Name).
		Suffix(metricsSchemaEndpoint).
		Do().Raw()
	if err != nil {
		return nil, err
	}
	var timeseriesSchema heapster_api.TimeseriesSchema
	if err := json.Unmarshal(body, &timeseriesSchema); err != nil {
		glog.V(2).Infof("response body: %v", string(body))
		return nil, err
	}
	return &timeseriesSchema, nil
}

var expectedSystemContainers = map[string]struct{}{
	"machine":       struct{}{},
	"kubelet":       struct{}{},
	"kube-proxy":    struct{}{},
	"system":        struct{}{},
	"docker-daemon": struct{}{},
}

func runHeapsterMetricsTest(fm kubeFramework, svc *kube_api.Service, expectedNodes, expectedPods []string) error {
	timeseries, err := getTimeseries(fm, svc)
	if err != nil {
		return err
	}
	if len(timeseries) == 0 {
		return fmt.Errorf("expected non zero timeseries")
	}
	schema, err := getSchema(fm, svc)
	if err != nil {
		return err
	}
	// Build a map of metric names to metric descriptors.
	mdMap := map[string]*heapster_api.MetricDescriptor{}
	for idx := range schema.Metrics {
		mdMap[schema.Metrics[idx].Name] = &schema.Metrics[idx]
	}
	actualPods := map[string]bool{}
	actualNodes := map[string]bool{}
	actualSystemContainers := map[string]map[string]struct{}{}
	for _, ts := range timeseries {
		// Verify the relevant labels are present.
		// All common labels must be present.
		for _, label := range sink_api.CommonLabels() {
			_, exists := ts.Labels[label.Key]
			if !exists {
				return fmt.Errorf("timeseries: %v does not contain common label: %v", ts, label)
			}
		}
		podName, podMetric := ts.Labels[sink_api.LabelPodName.Key]
		if podMetric {
			for _, label := range sink_api.PodLabels() {
				_, exists := ts.Labels[label.Key]
				if !exists {
					return fmt.Errorf("timeseries: %v does not contain pod label: %v", ts, label)
				}
			}
		}
		if podMetric {
			actualPods[podName] = true
		} else {
			if cName, ok := ts.Labels[sink_api.LabelContainerName.Key]; ok {
				hostname, ok := ts.Labels[sink_api.LabelHostname.Key]
				if !ok {
					return fmt.Errorf("hostname label missing on container %+v", ts)
				}

				if cName == cache.NodeContainerName {
					actualNodes[hostname] = true
				}
				if _, exists := expectedSystemContainers[cName]; exists {
					if actualSystemContainers[cName] == nil {
						actualSystemContainers[cName] = map[string]struct{}{}
					}
					actualSystemContainers[cName][hostname] = struct{}{}
				}
			} else {
				return fmt.Errorf("container_name label missing on timeseries - %v", ts)
			}
		}

		for metricName, points := range ts.Metrics {
			md, exists := mdMap[metricName]
			if !exists {
				return fmt.Errorf("unexpected metric %q", metricName)
			}
			for _, point := range points {
				for _, label := range md.Labels {
					_, exists := point.Labels[label.Key]
					if !exists {
						return fmt.Errorf("metric %q point %v does not contain metric label: %v", metricName, point, label)
					}
				}
			}

		}
	}
	// Validate that system containers are running on all the nodes.
	// This test could fail if one of the containers was down while the metrics sample was collected.
	for cName, hosts := range actualSystemContainers {
		for _, host := range expectedNodes {
			if _, ok := hosts[host]; !ok {
				return fmt.Errorf("System container %q not found on host: %q - %v", cName, host, actualSystemContainers)
			}
		}
	}

	if !expectedItemsExist(expectedPods, actualPods) {
		return fmt.Errorf("expected pods don't exist.\nExpected: %v\nActual:%v", expectedPods, actualPods)
	}
	if !expectedItemsExist(expectedNodes, actualNodes) {
		return fmt.Errorf("expected nodes don't exist.\nExpected: %v\nActual:%v", expectedNodes, actualNodes)
	}

	return nil
}

func expectedItemsExist(expectedItems []string, actualItems map[string]bool) bool {
	if len(actualItems) < len(expectedItems) {
		return false
	}
	for _, item := range expectedItems {
		if _, found := actualItems[item]; !found {
			return false
		}
	}
	return true
}

func apiTest(kubeVersion string) error {
	fm, err := newKubeFramework(kubeVersion)
	if err != nil {
		return err
	}
	if err := buildAndPushDockerImages(fm); err != nil {
		return err
	}
	// Create heapster pod and service.
	svc, rc, err := getHeapsterRcAndSvc(fm)
	if err != nil {
		return err
	}
	ns := *namespace
	deleteAll(fm, ns, svc, rc)
	if err = createAll(fm, ns, &svc, &rc); err != nil {
		return err
	}
	defer func() {
		deleteAll(fm, ns, svc, rc)
		removeHeapsterImage(fm)
	}()
	if err := fm.WaitUntilPodRunning(ns, rc.Spec.Template.Labels, time.Minute); err != nil {
		return err
	}
	if err := fm.WaitUntilServiceActive(svc, time.Minute); err != nil {
		return err
	}
	expectedPods, err := fm.GetPods()
	if err != nil {
		return err
	}
	expectedNodes, err := fm.GetNodes()
	if err != nil {
		return err
	}
	attempts := *maxRetries
	for {
		err := runHeapsterMetricsTest(fm, svc, expectedNodes, expectedPods)
		if *runForever {
			continue
		}
		if err == nil {
			return nil
		}
		if attempts == 0 {
			return err
		}
		attempts--
		time.Sleep(time.Second)
	}
}

func runApiTest() error {
	tempDir, err := ioutil.TempDir("", "deploy")
	if err != nil {
		return nil
	}
	defer os.RemoveAll(tempDir)
	if *kubeVersions == "" {
		return apiTest("")
	}
	kubeVersionsList := strings.Split(*kubeVersions, ",")
	for _, kubeVersion := range kubeVersionsList {
		return apiTest(kubeVersion)
	}
	return nil
}

func TestHeapster(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping heapster kubernetes integration test.")
	}
	require.NoError(t, runApiTest())
}
<|endoftext|>


<commit_before>// Copyright 2012 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 fs

import (
	"fmt"
	"os"
	"syscall"

	"github.com/jacobsa/comeback/internal/sys"
)

const (
	// A permissions bit mask that matches chmod(2)'s notion of permissions.
	permissionBits os.FileMode = os.ModePerm | os.ModeSetuid | os.ModeSetgid | os.ModeSticky
)

var gUserRegistry = sys.NewUserRegistry()
var gGroupRegistry = sys.NewGroupRegistry()

// Convert the result of os.Lstat or os.Stat to a local FileInfo struct.
// symlinkTarget should be empty if this is not a symlink.
func ConvertFileInfo(
	in os.FileInfo,
	symlinkTarget string) (out *FileInfo, err error) {
	out, err = convertFileInfo(in, symlinkTarget, gUserRegistry, gGroupRegistry)
	return
}

// Like ConvertFileInfo, but allows injecting registries.
func convertFileInfo(
	in os.FileInfo,
	symlinkTarget string,
	userRegistry sys.UserRegistry,
	groupRegistry sys.GroupRegistry) (out *FileInfo, err error) {
	// Grab system-specific info.
	statT, ok := in.Sys().(*syscall.Stat_t)
	if !ok {
		return nil, fmt.Errorf("Unexpected sys value: %v", in.Sys())
	}

	if statT.Size < 0 {
		panic(fmt.Sprintf("Unexpected size: %d", statT.Size))
	}

	// Create the basic struct.
	out = &FileInfo{
		Name:             in.Name(),
		Permissions:      in.Mode() & permissionBits,
		Uid:              sys.UserId(statT.Uid),
		Gid:              sys.GroupId(statT.Gid),
		MTime:            in.ModTime(),
		Size:             uint64(statT.Size),
		ContainingDevice: statT.Dev,
		Inode:            statT.Ino,
		Target:           symlinkTarget,
	}

	// Attempt to look up user info.
	username, err := userRegistry.FindById(out.Uid)

	if _, ok := err.(sys.NotFoundError); ok {
		err = nil
	} else if err != nil {
		return nil, fmt.Errorf("Looking up user: %v", err)
	} else {
		out.Username = &username
	}

	// Attempt to look up group info.
	groupname, err := groupRegistry.FindById(out.Gid)

	if _, ok := err.(sys.NotFoundError); ok {
		err = nil
	} else if err != nil {
		return nil, fmt.Errorf("Looking up group: %v", err)
	} else {
		out.Groupname = &groupname
	}

	// Convert the type.
	typeBits := in.Mode() & (os.ModeType | os.ModeCharDevice)
	switch typeBits {
	case 0:
		out.Type = TypeFile
	case os.ModeDir:
		out.Type = TypeDirectory
	case os.ModeSymlink:
		out.Type = TypeSymlink
	case os.ModeDevice:
		out.Type = TypeBlockDevice
		out.DeviceNumber = statT.Rdev
	case os.ModeDevice | os.ModeCharDevice:
		out.Type = TypeCharDevice
		out.DeviceNumber = statT.Rdev
	case os.ModeNamedPipe:
		out.Type = TypeNamedPipe
	case os.ModeSocket:
		out.Type = TypeSocket
	default:
		return out, fmt.Errorf("Unhandled mode: %v %u", in.Mode(), in.Mode())
	}

	return out, nil
}
<commit_msg>convert_file_info.go: fix a build error.<commit_after>// Copyright 2012 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 fs

import (
	"fmt"
	"os"
	"syscall"

	"github.com/jacobsa/comeback/internal/sys"
)

const (
	// A permissions bit mask that matches chmod(2)'s notion of permissions.
	permissionBits os.FileMode = os.ModePerm | os.ModeSetuid | os.ModeSetgid | os.ModeSticky
)

var gUserRegistry = sys.NewUserRegistry()
var gGroupRegistry = sys.NewGroupRegistry()

// Convert the result of os.Lstat or os.Stat to a local FileInfo struct.
// symlinkTarget should be empty if this is not a symlink.
func ConvertFileInfo(
	in os.FileInfo,
	symlinkTarget string) (out *FileInfo, err error) {
	out, err = convertFileInfo(in, symlinkTarget, gUserRegistry, gGroupRegistry)
	return
}

// Like ConvertFileInfo, but allows injecting registries.
func convertFileInfo(
	in os.FileInfo,
	symlinkTarget string,
	userRegistry sys.UserRegistry,
	groupRegistry sys.GroupRegistry) (out *FileInfo, err error) {
	// Grab system-specific info.
	statT, ok := in.Sys().(*syscall.Stat_t)
	if !ok {
		return nil, fmt.Errorf("Unexpected sys value: %v", in.Sys())
	}

	if statT.Size < 0 {
		panic(fmt.Sprintf("Unexpected size: %d", statT.Size))
	}

	// Create the basic struct.
	out = &FileInfo{
		Name:             in.Name(),
		Permissions:      in.Mode() & permissionBits,
		Uid:              sys.UserId(statT.Uid),
		Gid:              sys.GroupId(statT.Gid),
		MTime:            in.ModTime(),
		Size:             uint64(statT.Size),
		ContainingDevice: statT.Dev,
		Inode:            statT.Ino,
		Target:           symlinkTarget,
	}

	// Attempt to look up user info.
	username, err := userRegistry.FindById(out.Uid)

	if _, ok := err.(sys.NotFoundError); ok {
		err = nil
	} else if err != nil {
		return nil, fmt.Errorf("Looking up user: %v", err)
	} else {
		out.Username = &username
	}

	// Attempt to look up group info.
	groupname, err := groupRegistry.FindById(out.Gid)

	if _, ok := err.(sys.NotFoundError); ok {
		err = nil
	} else if err != nil {
		return nil, fmt.Errorf("Looking up group: %v", err)
	} else {
		out.Groupname = &groupname
	}

	// Convert the type.
	typeBits := in.Mode() & (os.ModeType | os.ModeCharDevice)
	switch typeBits {
	case 0:
		out.Type = TypeFile
	case os.ModeDir:
		out.Type = TypeDirectory
	case os.ModeSymlink:
		out.Type = TypeSymlink
	case os.ModeDevice:
		out.Type = TypeBlockDevice
		out.DeviceNumber = statT.Rdev
	case os.ModeDevice | os.ModeCharDevice:
		out.Type = TypeCharDevice
		out.DeviceNumber = statT.Rdev
	case os.ModeNamedPipe:
		out.Type = TypeNamedPipe
	case os.ModeSocket:
		out.Type = TypeSocket
	default:
		return out, fmt.Errorf("Unhandled mode: %v", in.Mode())
	}

	return out, nil
}
<|endoftext|>


<commit_before>package shared

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"time"

	"github.com/gorilla/websocket"

	"github.com/lxc/lxd/shared/logger"
)

func RFC3493Dialer(network, address string) (net.Conn, error) {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}

	addrs, err := net.LookupHost(host)
	if err != nil {
		return nil, err
	}
	for _, a := range addrs {
		c, err := net.DialTimeout(network, net.JoinHostPort(a, port), 10*time.Second)
		if err != nil {
			continue
		}
		return c, err
	}
	return nil, fmt.Errorf("Unable to connect to: " + address)
}

// InitTLSConfig returns a tls.Config populated with default encryption
// parameters. This is used as baseline config for both client and server
// certificates used by LXD.
func InitTLSConfig() *tls.Config {
	return &tls.Config{
		MinVersion: tls.VersionTLS12,
		MaxVersion: tls.VersionTLS12,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
		},
		PreferServerCipherSuites: true,
	}
}

func finalizeTLSConfig(tlsConfig *tls.Config, tlsRemoteCert *x509.Certificate) {
	// Trusted certificates
	if tlsRemoteCert != nil {
		caCertPool := tlsConfig.RootCAs
		if caCertPool == nil {
			caCertPool = x509.NewCertPool()
		}

		// Make it a valid RootCA
		tlsRemoteCert.IsCA = true
		tlsRemoteCert.KeyUsage = x509.KeyUsageCertSign

		// Setup the pool
		caCertPool.AddCert(tlsRemoteCert)
		tlsConfig.RootCAs = caCertPool

		// Set the ServerName
		if tlsRemoteCert.DNSNames != nil {
			tlsConfig.ServerName = tlsRemoteCert.DNSNames[0]
		}
	}

	tlsConfig.BuildNameToCertificate()
}

func GetTLSConfig(tlsClientCertFile string, tlsClientKeyFile string, tlsClientCAFile string, tlsRemoteCert *x509.Certificate) (*tls.Config, error) {
	tlsConfig := InitTLSConfig()

	// Client authentication
	if tlsClientCertFile != "" && tlsClientKeyFile != "" {
		cert, err := tls.LoadX509KeyPair(tlsClientCertFile, tlsClientKeyFile)
		if err != nil {
			return nil, err
		}

		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	if tlsClientCAFile != "" {
		caCertificates, err := ioutil.ReadFile(tlsClientCAFile)
		if err != nil {
			return nil, err
		}

		caPool := x509.NewCertPool()
		caPool.AppendCertsFromPEM(caCertificates)

		tlsConfig.RootCAs = caPool
	}

	finalizeTLSConfig(tlsConfig, tlsRemoteCert)
	return tlsConfig, nil
}

func GetTLSConfigMem(tlsClientCert string, tlsClientKey string, tlsClientCA string, tlsRemoteCertPEM string, insecureSkipVerify bool) (*tls.Config, error) {
	tlsConfig := InitTLSConfig()
	tlsConfig.InsecureSkipVerify = insecureSkipVerify
	// Client authentication
	if tlsClientCert != "" && tlsClientKey != "" {
		cert, err := tls.X509KeyPair([]byte(tlsClientCert), []byte(tlsClientKey))
		if err != nil {
			return nil, err
		}

		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	var tlsRemoteCert *x509.Certificate
	if tlsRemoteCertPEM != "" {
		// Ignore any content outside of the PEM bytes we care about
		certBlock, _ := pem.Decode([]byte(tlsRemoteCertPEM))
		if certBlock == nil {
			return nil, fmt.Errorf("Invalid remote certificate")
		}

		var err error
		tlsRemoteCert, err = x509.ParseCertificate(certBlock.Bytes)
		if err != nil {
			return nil, err
		}
	}

	if tlsClientCA != "" {
		caPool := x509.NewCertPool()
		caPool.AppendCertsFromPEM([]byte(tlsClientCA))

		tlsConfig.RootCAs = caPool
	}

	finalizeTLSConfig(tlsConfig, tlsRemoteCert)

	return tlsConfig, nil
}

func IsLoopback(iface *net.Interface) bool {
	return int(iface.Flags&net.FlagLoopback) > 0
}

func WebsocketSendStream(conn *websocket.Conn, r io.Reader, bufferSize int) chan bool {
	ch := make(chan bool)

	if r == nil {
		close(ch)
		return ch
	}

	go func(conn *websocket.Conn, r io.Reader) {
		in := ReaderToChannel(r, bufferSize)
		for {
			buf, ok := <-in
			if !ok {
				break
			}

			w, err := conn.NextWriter(websocket.BinaryMessage)
			if err != nil {
				logger.Debugf("Got error getting next writer %s", err)
				break
			}

			_, err = w.Write(buf)
			w.Close()
			if err != nil {
				logger.Debugf("Got err writing %s", err)
				break
			}
		}
		conn.WriteMessage(websocket.TextMessage, []byte{})
		ch <- true
	}(conn, r)

	return ch
}

func WebsocketRecvStream(w io.Writer, conn *websocket.Conn) chan bool {
	ch := make(chan bool)

	go func(w io.Writer, conn *websocket.Conn) {
		for {
			mt, r, err := conn.NextReader()
			if mt == websocket.CloseMessage {
				logger.Debugf("Got close message for reader")
				break
			}

			if mt == websocket.TextMessage {
				logger.Debugf("got message barrier")
				break
			}

			if err != nil {
				logger.Debugf("Got error getting next reader %s, %s", err, w)
				break
			}

			buf, err := ioutil.ReadAll(r)
			if err != nil {
				logger.Debugf("Got error writing to writer %s", err)
				break
			}

			if w == nil {
				continue
			}

			i, err := w.Write(buf)
			if i != len(buf) {
				logger.Debugf("Didn't write all of buf")
				break
			}
			if err != nil {
				logger.Debugf("Error writing buf %s", err)
				break
			}
		}
		ch <- true
	}(w, conn)

	return ch
}

func WebsocketProxy(source *websocket.Conn, target *websocket.Conn) chan bool {
	forward := func(in *websocket.Conn, out *websocket.Conn, ch chan bool) {
		for {
			mt, r, err := in.NextReader()
			if err != nil {
				break
			}

			w, err := out.NextWriter(mt)
			if err != nil {
				break
			}

			_, err = io.Copy(w, r)
			w.Close()
			if err != nil {
				break
			}
		}

		ch <- true
	}

	chSend := make(chan bool)
	go forward(source, target, chSend)

	chRecv := make(chan bool)
	go forward(target, source, chRecv)

	ch := make(chan bool)
	go func() {
		select {
		case <-chSend:
		case <-chRecv:
		}

		source.Close()
		target.Close()

		ch <- true
	}()

	return ch
}

func defaultReader(conn *websocket.Conn, r io.ReadCloser, readDone chan<- bool) {
	/* For now, we don't need to adjust buffer sizes in
	* WebsocketMirror, since it's used for interactive things like
	* exec.
	 */
	in := ReaderToChannel(r, -1)
	for {
		buf, ok := <-in
		if !ok {
			r.Close()
			logger.Debugf("sending write barrier")
			conn.WriteMessage(websocket.TextMessage, []byte{})
			readDone <- true
			return
		}
		w, err := conn.NextWriter(websocket.BinaryMessage)
		if err != nil {
			logger.Debugf("Got error getting next writer %s", err)
			break
		}

		_, err = w.Write(buf)
		w.Close()
		if err != nil {
			logger.Debugf("Got err writing %s", err)
			break
		}
	}
	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
	conn.WriteMessage(websocket.CloseMessage, closeMsg)
	readDone <- true
	r.Close()
}

func defaultWriter(conn *websocket.Conn, w io.WriteCloser, writeDone chan<- bool) {
	for {
		mt, r, err := conn.NextReader()
		if err != nil {
			logger.Debugf("Got error getting next reader %s, %s", err, w)
			break
		}

		if mt == websocket.CloseMessage {
			logger.Debugf("Got close message for reader")
			break
		}

		if mt == websocket.TextMessage {
			logger.Debugf("Got message barrier, resetting stream")
			break
		}

		buf, err := ioutil.ReadAll(r)
		if err != nil {
			logger.Debugf("Got error writing to writer %s", err)
			break
		}
		i, err := w.Write(buf)
		if i != len(buf) {
			logger.Debugf("Didn't write all of buf")
			break
		}
		if err != nil {
			logger.Debugf("Error writing buf %s", err)
			break
		}
	}
	writeDone <- true
	w.Close()
}

// WebsocketMirror allows mirroring a reader to a websocket and taking the
// result and writing it to a writer. This function allows for multiple
// mirrorings and correctly negotiates stream endings. However, it means any
// websocket.Conns passed to it are live when it returns, and must be closed
// explicitly.
type WebSocketMirrorReader func(conn *websocket.Conn, r io.ReadCloser, readDone chan<- bool)
type WebSocketMirrorWriter func(conn *websocket.Conn, w io.WriteCloser, writeDone chan<- bool)

func WebsocketMirror(conn *websocket.Conn, w io.WriteCloser, r io.ReadCloser, Reader WebSocketMirrorReader, Writer WebSocketMirrorWriter) (chan bool, chan bool) {
	readDone := make(chan bool, 1)
	writeDone := make(chan bool, 1)

	ReadFunc := Reader
	if ReadFunc == nil {
		ReadFunc = defaultReader
	}

	WriteFunc := Writer
	if WriteFunc == nil {
		WriteFunc = defaultWriter
	}

	go ReadFunc(conn, r, readDone)
	go WriteFunc(conn, w, writeDone)

	return readDone, writeDone
}

var WebsocketUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// AllocatePort asks the kernel for a free open port that is ready to use
func AllocatePort() (int, error) {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		return -1, err
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return -1, err
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}
<commit_msg>shared/network: add WebsocketConsoleMirror()<commit_after>package shared

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"time"

	"github.com/gorilla/websocket"

	"github.com/lxc/lxd/shared/logger"
)

func RFC3493Dialer(network, address string) (net.Conn, error) {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}

	addrs, err := net.LookupHost(host)
	if err != nil {
		return nil, err
	}
	for _, a := range addrs {
		c, err := net.DialTimeout(network, net.JoinHostPort(a, port), 10*time.Second)
		if err != nil {
			continue
		}
		return c, err
	}
	return nil, fmt.Errorf("Unable to connect to: " + address)
}

// InitTLSConfig returns a tls.Config populated with default encryption
// parameters. This is used as baseline config for both client and server
// certificates used by LXD.
func InitTLSConfig() *tls.Config {
	return &tls.Config{
		MinVersion: tls.VersionTLS12,
		MaxVersion: tls.VersionTLS12,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
		},
		PreferServerCipherSuites: true,
	}
}

func finalizeTLSConfig(tlsConfig *tls.Config, tlsRemoteCert *x509.Certificate) {
	// Trusted certificates
	if tlsRemoteCert != nil {
		caCertPool := tlsConfig.RootCAs
		if caCertPool == nil {
			caCertPool = x509.NewCertPool()
		}

		// Make it a valid RootCA
		tlsRemoteCert.IsCA = true
		tlsRemoteCert.KeyUsage = x509.KeyUsageCertSign

		// Setup the pool
		caCertPool.AddCert(tlsRemoteCert)
		tlsConfig.RootCAs = caCertPool

		// Set the ServerName
		if tlsRemoteCert.DNSNames != nil {
			tlsConfig.ServerName = tlsRemoteCert.DNSNames[0]
		}
	}

	tlsConfig.BuildNameToCertificate()
}

func GetTLSConfig(tlsClientCertFile string, tlsClientKeyFile string, tlsClientCAFile string, tlsRemoteCert *x509.Certificate) (*tls.Config, error) {
	tlsConfig := InitTLSConfig()

	// Client authentication
	if tlsClientCertFile != "" && tlsClientKeyFile != "" {
		cert, err := tls.LoadX509KeyPair(tlsClientCertFile, tlsClientKeyFile)
		if err != nil {
			return nil, err
		}

		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	if tlsClientCAFile != "" {
		caCertificates, err := ioutil.ReadFile(tlsClientCAFile)
		if err != nil {
			return nil, err
		}

		caPool := x509.NewCertPool()
		caPool.AppendCertsFromPEM(caCertificates)

		tlsConfig.RootCAs = caPool
	}

	finalizeTLSConfig(tlsConfig, tlsRemoteCert)
	return tlsConfig, nil
}

func GetTLSConfigMem(tlsClientCert string, tlsClientKey string, tlsClientCA string, tlsRemoteCertPEM string, insecureSkipVerify bool) (*tls.Config, error) {
	tlsConfig := InitTLSConfig()
	tlsConfig.InsecureSkipVerify = insecureSkipVerify
	// Client authentication
	if tlsClientCert != "" && tlsClientKey != "" {
		cert, err := tls.X509KeyPair([]byte(tlsClientCert), []byte(tlsClientKey))
		if err != nil {
			return nil, err
		}

		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	var tlsRemoteCert *x509.Certificate
	if tlsRemoteCertPEM != "" {
		// Ignore any content outside of the PEM bytes we care about
		certBlock, _ := pem.Decode([]byte(tlsRemoteCertPEM))
		if certBlock == nil {
			return nil, fmt.Errorf("Invalid remote certificate")
		}

		var err error
		tlsRemoteCert, err = x509.ParseCertificate(certBlock.Bytes)
		if err != nil {
			return nil, err
		}
	}

	if tlsClientCA != "" {
		caPool := x509.NewCertPool()
		caPool.AppendCertsFromPEM([]byte(tlsClientCA))

		tlsConfig.RootCAs = caPool
	}

	finalizeTLSConfig(tlsConfig, tlsRemoteCert)

	return tlsConfig, nil
}

func IsLoopback(iface *net.Interface) bool {
	return int(iface.Flags&net.FlagLoopback) > 0
}

func WebsocketSendStream(conn *websocket.Conn, r io.Reader, bufferSize int) chan bool {
	ch := make(chan bool)

	if r == nil {
		close(ch)
		return ch
	}

	go func(conn *websocket.Conn, r io.Reader) {
		in := ReaderToChannel(r, bufferSize)
		for {
			buf, ok := <-in
			if !ok {
				break
			}

			w, err := conn.NextWriter(websocket.BinaryMessage)
			if err != nil {
				logger.Debugf("Got error getting next writer %s", err)
				break
			}

			_, err = w.Write(buf)
			w.Close()
			if err != nil {
				logger.Debugf("Got err writing %s", err)
				break
			}
		}
		conn.WriteMessage(websocket.TextMessage, []byte{})
		ch <- true
	}(conn, r)

	return ch
}

func WebsocketRecvStream(w io.Writer, conn *websocket.Conn) chan bool {
	ch := make(chan bool)

	go func(w io.Writer, conn *websocket.Conn) {
		for {
			mt, r, err := conn.NextReader()
			if mt == websocket.CloseMessage {
				logger.Debugf("Got close message for reader")
				break
			}

			if mt == websocket.TextMessage {
				logger.Debugf("got message barrier")
				break
			}

			if err != nil {
				logger.Debugf("Got error getting next reader %s, %s", err, w)
				break
			}

			buf, err := ioutil.ReadAll(r)
			if err != nil {
				logger.Debugf("Got error writing to writer %s", err)
				break
			}

			if w == nil {
				continue
			}

			i, err := w.Write(buf)
			if i != len(buf) {
				logger.Debugf("Didn't write all of buf")
				break
			}
			if err != nil {
				logger.Debugf("Error writing buf %s", err)
				break
			}
		}
		ch <- true
	}(w, conn)

	return ch
}

func WebsocketProxy(source *websocket.Conn, target *websocket.Conn) chan bool {
	forward := func(in *websocket.Conn, out *websocket.Conn, ch chan bool) {
		for {
			mt, r, err := in.NextReader()
			if err != nil {
				break
			}

			w, err := out.NextWriter(mt)
			if err != nil {
				break
			}

			_, err = io.Copy(w, r)
			w.Close()
			if err != nil {
				break
			}
		}

		ch <- true
	}

	chSend := make(chan bool)
	go forward(source, target, chSend)

	chRecv := make(chan bool)
	go forward(target, source, chRecv)

	ch := make(chan bool)
	go func() {
		select {
		case <-chSend:
		case <-chRecv:
		}

		source.Close()
		target.Close()

		ch <- true
	}()

	return ch
}

func defaultReader(conn *websocket.Conn, r io.ReadCloser, readDone chan<- bool) {
	/* For now, we don't need to adjust buffer sizes in
	* WebsocketMirror, since it's used for interactive things like
	* exec.
	 */
	in := ReaderToChannel(r, -1)
	for {
		buf, ok := <-in
		if !ok {
			r.Close()
			logger.Debugf("sending write barrier")
			conn.WriteMessage(websocket.TextMessage, []byte{})
			readDone <- true
			return
		}
		w, err := conn.NextWriter(websocket.BinaryMessage)
		if err != nil {
			logger.Debugf("Got error getting next writer %s", err)
			break
		}

		_, err = w.Write(buf)
		w.Close()
		if err != nil {
			logger.Debugf("Got err writing %s", err)
			break
		}
	}
	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
	conn.WriteMessage(websocket.CloseMessage, closeMsg)
	readDone <- true
	r.Close()
}

func defaultWriter(conn *websocket.Conn, w io.WriteCloser, writeDone chan<- bool) {
	for {
		mt, r, err := conn.NextReader()
		if err != nil {
			logger.Debugf("Got error getting next reader %s, %s", err, w)
			break
		}

		if mt == websocket.CloseMessage {
			logger.Debugf("Got close message for reader")
			break
		}

		if mt == websocket.TextMessage {
			logger.Debugf("Got message barrier, resetting stream")
			break
		}

		buf, err := ioutil.ReadAll(r)
		if err != nil {
			logger.Debugf("Got error writing to writer %s", err)
			break
		}
		i, err := w.Write(buf)
		if i != len(buf) {
			logger.Debugf("Didn't write all of buf")
			break
		}
		if err != nil {
			logger.Debugf("Error writing buf %s", err)
			break
		}
	}
	writeDone <- true
	w.Close()
}

// WebsocketMirror allows mirroring a reader to a websocket and taking the
// result and writing it to a writer. This function allows for multiple
// mirrorings and correctly negotiates stream endings. However, it means any
// websocket.Conns passed to it are live when it returns, and must be closed
// explicitly.
type WebSocketMirrorReader func(conn *websocket.Conn, r io.ReadCloser, readDone chan<- bool)
type WebSocketMirrorWriter func(conn *websocket.Conn, w io.WriteCloser, writeDone chan<- bool)

func WebsocketMirror(conn *websocket.Conn, w io.WriteCloser, r io.ReadCloser, Reader WebSocketMirrorReader, Writer WebSocketMirrorWriter) (chan bool, chan bool) {
	readDone := make(chan bool, 1)
	writeDone := make(chan bool, 1)

	ReadFunc := Reader
	if ReadFunc == nil {
		ReadFunc = defaultReader
	}

	WriteFunc := Writer
	if WriteFunc == nil {
		WriteFunc = defaultWriter
	}

	go ReadFunc(conn, r, readDone)
	go WriteFunc(conn, w, writeDone)

	return readDone, writeDone
}

func WebsocketConsoleMirror(conn *websocket.Conn, w io.WriteCloser, r io.ReadCloser) (chan bool, chan bool) {
	readDone := make(chan bool, 1)
	writeDone := make(chan bool, 1)

	go defaultWriter(conn, w, writeDone)

	go func(conn *websocket.Conn, r io.ReadCloser) {
		in := ReaderToChannel(r, -1)
		for {
			buf, ok := <-in
			if !ok {
				r.Close()
				logger.Debugf("sending write barrier")
				conn.WriteMessage(websocket.TextMessage, []byte{})
				readDone <- true
				return
			}
			w, err := conn.NextWriter(websocket.BinaryMessage)
			if err != nil {
				logger.Debugf("Got error getting next writer %s", err)
				break
			}

			_, err = w.Write(buf)
			w.Close()
			if err != nil {
				logger.Debugf("Got err writing %s", err)
				break
			}
		}
		closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
		conn.WriteMessage(websocket.CloseMessage, closeMsg)
		readDone <- true
		r.Close()
	}(conn, r)

	return readDone, writeDone
}

var WebsocketUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// AllocatePort asks the kernel for a free open port that is ready to use
func AllocatePort() (int, error) {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		return -1, err
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return -1, err
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}
<|endoftext|>


<commit_before>//  Copyright (c) 2015 ikawaha.
//  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 lattice

import (
	"bytes"
	"testing"
	"unicode/utf8"

	"github.com/ikawaha/kagome/internal/dic"
)

func TestLatticeBuild01(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Error("cannot new a lattice")
	}
	defer la.Free()

	inp := ""
	la.Build(inp)
	if la.Input != inp {
		t.Errorf("got %v, expected %v", la.Input, inp)
	}
	boseos := node{ID: -1}
	if len(la.list) != 2 {
		t.Errorf("lattice initialize error: got %v, expected has 2 eos/bos nodes", la.list)
	} else if len(la.list[0]) != 1 || *la.list[0][0] != boseos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[0][0], boseos)
	} else if len(la.list[1]) != 1 || *la.list[1][0] != boseos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[1][0], boseos)
	}
	if len(la.Output) != 0 {
		t.Errorf("lattice initialize error: got %v, expected empty", la.Output)
	}

	if la.dic == nil {
		t.Errorf("lattice initialize error: dic is nil")
	}
	if la.udic != nil {
		t.Errorf("lattice initialize error: got %v, expected empty", la.udic)
	}
}

func TestLatticeBuild02(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	inp := "あ"
	la.Build(inp)
	if la.Input != inp {
		t.Errorf("got %v, expected %v", la.Input, inp)
	}
	bos := node{ID: -1}
	eos := node{ID: -1, Start: 1}
	if len(la.list) != 3 {
		t.Errorf("lattice initialize error: got %v, expected has 2 eos/bos nodes", la.list)
	} else if len(la.list[0]) != 1 || *la.list[0][0] != bos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[0][0], bos)
	} else if len(la.list[2]) != 1 || *la.list[2][0] != eos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[2][0], eos)
	}

	expected := 4
	if len(la.list[1]) != expected {
		t.Errorf("lattice initialize error: got %v, expected %v", len(la.list[1]), expected)
	} else {
		l := la.list[1]
		callAndResponse := []struct {
			in  int
			out node
		}{
			{in: 0, out: node{122, 0, KNOWN, 0, 3, 3, 5549, "あ", nil}},
			{in: 1, out: node{123, 0, KNOWN, 0, 776, 776, 6690, "あ", nil}},
			{in: 2, out: node{124, 0, KNOWN, 0, 2, 2, 4262, "あ", nil}},
			{in: 3, out: node{125, 0, KNOWN, 0, 1118, 1118, 9035, "あ", nil}},
		}
		for _, cr := range callAndResponse {
			if *l[cr.in] != cr.out {
				t.Errorf("lattice initialize error: got %v, expected %v", l[cr.in], cr.out)
			}
		}
	}
	if len(la.Output) != 0 {
		t.Errorf("lattice initialize error: got %v, expected empty", la.Output)
	}
	if la.dic == nil {
		t.Errorf("lattice initialize error: dic is nil")
	}
	if la.udic != nil {
		t.Errorf("lattice initialize error: got %v, expected empty", la.udic)
	}
}

func TestLatticeBuild03(t *testing.T) {

	const udicPath = "../../_sample/userdic.txt"

	udic, e := dic.NewUserDic(udicPath)
	if e != nil {
		t.Fatalf("unexpected error: cannot load user dic, %v", e)
	}
	la := New(dic.SysDic(), udic)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	inp := "朝青龍"
	la.Build(inp)
	if la.Input != inp {
		t.Errorf("got %v, expected %v", la.Input, inp)
	}

	if la.list[3][0].Class != USER {
		t.Errorf("%+v", la)
	}

	if len(la.Output) != 0 {
		t.Errorf("lattice initialize error: got %v, expected empty", la.Output)
	}
	if la.dic == nil {
		t.Errorf("lattice initialize error: dic is nil")
	}
	if la.udic == nil {
		t.Errorf("lattice initialize error: got %v, expected not empty", la.udic)
	}
}

func TestLatticeBuild04(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	inp := "ポポピ"
	la.Build(inp)
	if la.Input != inp {
		t.Errorf("got %v, expected %v", la.Input, inp)
	}
	bos := node{ID: -1}
	eos := node{ID: -1, Start: 3}
	if len(la.list) != 5 {
		t.Errorf("lattice initialize error: got %v, expected has 2 eos/bos nodes", la.list)
	} else if len(la.list[0]) != 1 || *la.list[0][0] != bos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[0][0], bos)
	} else if len(la.list[len(la.list)-1]) != 1 || *la.list[len(la.list)-1][0] != eos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[len(la.list)-1][0], eos)
	}

	expected := 7
	if len(la.list[1]) != expected {
		t.Errorf("lattice initialize error: got %v, expected %v", len(la.list[1]), expected)
	} else {
		l := la.list[1]
		callAndResponse := []struct {
			in  int
			out node
		}{
			{in: 0, out: node{98477, 0, KNOWN, 0, 1285, 1285, 4279, "ポ", nil}},
			{in: 1, out: node{31, 0, UNKNOWN, 0, 1289, 1289, 13581, "ポ", nil}},
			{in: 2, out: node{32, 0, UNKNOWN, 0, 1285, 1285, 9461, "ポ", nil}},
			{in: 3, out: node{33, 0, UNKNOWN, 0, 1293, 1293, 13661, "ポ", nil}},
			{in: 4, out: node{34, 0, UNKNOWN, 0, 1292, 1292, 10922, "ポ", nil}},
			{in: 5, out: node{35, 0, UNKNOWN, 0, 1288, 1288, 10521, "ポ", nil}},
			{in: 6, out: node{36, 0, UNKNOWN, 0, 3, 3, 14138, "ポ", nil}},
		}
		for _, cr := range callAndResponse {
			if *l[cr.in] != cr.out {
				t.Errorf("lattice initialize error: got %v, expected %v", l[cr.in], cr.out)
			}
		}
	}
	if len(la.Output) != 0 {
		t.Errorf("lattice initialize error: got %v, expected empty", la.Output)
	}
	if la.dic == nil {
		t.Errorf("lattice initialize error: dic is nil")
	}
	if la.udic != nil {
		t.Errorf("lattice initialize error: got %v, expected empty", la.udic)
	}
}

func TestLatticeBuild05(t *testing.T) {

	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	inp := "ポポピポンポコナーノ"
	var b bytes.Buffer
	for i, step := 0, utf8.RuneCountInString(inp); i < maximumUnknownWordLength; i = i + step {
		if _, e := b.WriteString(inp); e != nil {
			t.Fatalf("unexpected error: create the test input, %v", b.String())
		}
	}
	la.Build(b.String())
	for i := range la.list {
		for j := range la.list[i] {
			l := utf8.RuneCountInString(la.list[i][j].Surface)
			if l > maximumUnknownWordLength {
				t.Errorf("too long unknown word, %v", l)
			}
		}
	}
}

func TestKanjiOnly01(t *testing.T) {
	callAndResponse := []struct {
		in  string
		out bool
	}{
		{in: "ひらがな", out: false},
		{in: "カタカナ", out: false},
		{in: "漢字", out: true},
		{in: "かな漢字交じり", out: false},
		{in: "123", out: false},
		{in: "#$%", out: false},
		{in: "", out: false},
	}
	for _, cr := range callAndResponse {
		if rsp := kanjiOnly(cr.in); rsp != cr.out {
			t.Errorf("in: %v, got %v, expected: %v", cr.in, rsp, cr.out)
		}
	}
}

func TestLatticeString(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	expected := ""
	str := la.String()
	if str != expected {
		t.Errorf("got %v, expected: %v", str, expected)
	}

	la.Build("わたしまけましたわ")
	str = la.String()
	if str == "" {
		t.Errorf("got empty string")
	}
}

func TestLatticeDot(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	expected := "graph lattice {\n\tdpi=48;\n\tgraph [style=filled, rankdir=LR]\n}\n"
	var b bytes.Buffer
	la.Dot(&b)
	if b.String() != expected {
		t.Errorf("got %v, expected: %v", b.String(), expected)
	}
	b.Reset()
	la.Build("わたしまけましたわ")
	la.Dot(&b)
	if b.String() == "" {
		t.Errorf("got empty string")
	}
}

func TestLatticeNewAndFree(t *testing.T) {
	for i := 0; i < 100; i++ {
		la := New(dic.SysDic(), nil)
		if la == nil {
			t.Fatal("unexpected error: cannot new a lattice")
		}
		if la.Input != "" {
			t.Fatalf("unexpected error: lattice input initialize error, %+v", la.Input)
		}
		if len(la.Output) != 0 {
			t.Fatalf("unexpected error: lattice output initialize error, %+v", la.Output)
		}
		if len(la.list) != 0 {
			t.Fatalf("unexpected error: lattice list initialize error, %+v", la.list)
		}
		la.Build("すべては科学する心に宿るのだ")
		la.Free()

		// renew
		la = New(dic.SysDic(), nil)
		if la == nil {
			t.Fatal("unexpected error: cannot new a lattice")
		}
		if la.Input != "" {
			t.Fatalf("unexpected error: lattice input initialize error, %+v", la.Input)
		}
		if len(la.Output) != 0 {
			t.Fatalf("unexpected error: lattice output initialize error, %+v", la.Output)
		}
		if len(la.list) != 0 {
			t.Fatalf("unexpected error: lattice list initialize error, %+v", la.list)
		}
		la.Free()
	}
}

func TestForward(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("unexpected error: cannot new a lattice")
	}

	la.Forward(Normal)
	la.Forward(Search)
	la.Forward(Extended)

	for _, m := range []TokenizeMode{Normal, Search, Extended} {
		la.Build("わたしまけましたわ")
		la.Forward(m)
	}
}

func TestBackward(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("unexpected error: cannot new a lattice")
	}

	// only run
	la.Backward(Normal)
	la.Backward(Search)
	la.Backward(Extended)

	for _, m := range []TokenizeMode{Normal, Search, Extended} {
		la.Build("わたしまけましたわ")
		la.Forward(m)
		la.Backward(m)
	}
}
<commit_msg>Add a test<commit_after>//  Copyright (c) 2015 ikawaha.
//  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 lattice

import (
	"bytes"
	"testing"
	"unicode/utf8"

	"github.com/ikawaha/kagome/internal/dic"
)

func TestLatticeBuild01(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Error("cannot new a lattice")
	}
	defer la.Free()

	inp := ""
	la.Build(inp)
	if la.Input != inp {
		t.Errorf("got %v, expected %v", la.Input, inp)
	}
	boseos := node{ID: -1}
	if len(la.list) != 2 {
		t.Errorf("lattice initialize error: got %v, expected has 2 eos/bos nodes", la.list)
	} else if len(la.list[0]) != 1 || *la.list[0][0] != boseos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[0][0], boseos)
	} else if len(la.list[1]) != 1 || *la.list[1][0] != boseos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[1][0], boseos)
	}
	if len(la.Output) != 0 {
		t.Errorf("lattice initialize error: got %v, expected empty", la.Output)
	}

	if la.dic == nil {
		t.Errorf("lattice initialize error: dic is nil")
	}
	if la.udic != nil {
		t.Errorf("lattice initialize error: got %v, expected empty", la.udic)
	}
}

func TestLatticeBuild02(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	inp := "あ"
	la.Build(inp)
	if la.Input != inp {
		t.Errorf("got %v, expected %v", la.Input, inp)
	}
	bos := node{ID: -1}
	eos := node{ID: -1, Start: 1}
	if len(la.list) != 3 {
		t.Errorf("lattice initialize error: got %v, expected has 2 eos/bos nodes", la.list)
	} else if len(la.list[0]) != 1 || *la.list[0][0] != bos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[0][0], bos)
	} else if len(la.list[2]) != 1 || *la.list[2][0] != eos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[2][0], eos)
	}

	expected := 4
	if len(la.list[1]) != expected {
		t.Errorf("lattice initialize error: got %v, expected %v", len(la.list[1]), expected)
	} else {
		l := la.list[1]
		callAndResponse := []struct {
			in  int
			out node
		}{
			{in: 0, out: node{122, 0, KNOWN, 0, 3, 3, 5549, "あ", nil}},
			{in: 1, out: node{123, 0, KNOWN, 0, 776, 776, 6690, "あ", nil}},
			{in: 2, out: node{124, 0, KNOWN, 0, 2, 2, 4262, "あ", nil}},
			{in: 3, out: node{125, 0, KNOWN, 0, 1118, 1118, 9035, "あ", nil}},
		}
		for _, cr := range callAndResponse {
			if *l[cr.in] != cr.out {
				t.Errorf("lattice initialize error: got %v, expected %v", l[cr.in], cr.out)
			}
		}
	}
	if len(la.Output) != 0 {
		t.Errorf("lattice initialize error: got %v, expected empty", la.Output)
	}
	if la.dic == nil {
		t.Errorf("lattice initialize error: dic is nil")
	}
	if la.udic != nil {
		t.Errorf("lattice initialize error: got %v, expected empty", la.udic)
	}
}

func TestLatticeBuild03(t *testing.T) {

	const udicPath = "../../_sample/userdic.txt"

	udic, e := dic.NewUserDic(udicPath)
	if e != nil {
		t.Fatalf("unexpected error: cannot load user dic, %v", e)
	}
	la := New(dic.SysDic(), udic)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	inp := "朝青龍"
	la.Build(inp)
	if la.Input != inp {
		t.Errorf("got %v, expected %v", la.Input, inp)
	}

	if la.list[3][0].Class != USER {
		t.Errorf("%+v", la)
	}

	if len(la.Output) != 0 {
		t.Errorf("lattice initialize error: got %v, expected empty", la.Output)
	}
	if la.dic == nil {
		t.Errorf("lattice initialize error: dic is nil")
	}
	if la.udic == nil {
		t.Errorf("lattice initialize error: got %v, expected not empty", la.udic)
	}
}

func TestLatticeBuild04(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	inp := "ポポピ"
	la.Build(inp)
	if la.Input != inp {
		t.Errorf("got %v, expected %v", la.Input, inp)
	}
	bos := node{ID: -1}
	eos := node{ID: -1, Start: 3}
	if len(la.list) != 5 {
		t.Errorf("lattice initialize error: got %v, expected has 2 eos/bos nodes", la.list)
	} else if len(la.list[0]) != 1 || *la.list[0][0] != bos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[0][0], bos)
	} else if len(la.list[len(la.list)-1]) != 1 || *la.list[len(la.list)-1][0] != eos {
		t.Errorf("lattice initialize error: got %v, expected %v", *la.list[len(la.list)-1][0], eos)
	}

	expected := 7
	if len(la.list[1]) != expected {
		t.Errorf("lattice initialize error: got %v, expected %v", len(la.list[1]), expected)
	} else {
		l := la.list[1]
		callAndResponse := []struct {
			in  int
			out node
		}{
			{in: 0, out: node{98477, 0, KNOWN, 0, 1285, 1285, 4279, "ポ", nil}},
			{in: 1, out: node{31, 0, UNKNOWN, 0, 1289, 1289, 13581, "ポ", nil}},
			{in: 2, out: node{32, 0, UNKNOWN, 0, 1285, 1285, 9461, "ポ", nil}},
			{in: 3, out: node{33, 0, UNKNOWN, 0, 1293, 1293, 13661, "ポ", nil}},
			{in: 4, out: node{34, 0, UNKNOWN, 0, 1292, 1292, 10922, "ポ", nil}},
			{in: 5, out: node{35, 0, UNKNOWN, 0, 1288, 1288, 10521, "ポ", nil}},
			{in: 6, out: node{36, 0, UNKNOWN, 0, 3, 3, 14138, "ポ", nil}},
		}
		for _, cr := range callAndResponse {
			if *l[cr.in] != cr.out {
				t.Errorf("lattice initialize error: got %v, expected %v", l[cr.in], cr.out)
			}
		}
	}
	if len(la.Output) != 0 {
		t.Errorf("lattice initialize error: got %v, expected empty", la.Output)
	}
	if la.dic == nil {
		t.Errorf("lattice initialize error: dic is nil")
	}
	if la.udic != nil {
		t.Errorf("lattice initialize error: got %v, expected empty", la.udic)
	}
}

func TestLatticeBuild05(t *testing.T) {

	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	inp := "ポポピポンポコナーノ"
	var b bytes.Buffer
	for i, step := 0, utf8.RuneCountInString(inp); i < maximumUnknownWordLength; i = i + step {
		if _, e := b.WriteString(inp); e != nil {
			t.Fatalf("unexpected error: create the test input, %v", b.String())
		}
	}
	la.Build(b.String())
	for i := range la.list {
		for j := range la.list[i] {
			l := utf8.RuneCountInString(la.list[i][j].Surface)
			if l > maximumUnknownWordLength {
				t.Errorf("too long unknown word, %v", l)
			}
		}
	}
}

func TestKanjiOnly01(t *testing.T) {
	callAndResponse := []struct {
		in  string
		out bool
	}{
		{in: "ひらがな", out: false},
		{in: "カタカナ", out: false},
		{in: "漢字", out: true},
		{in: "かな漢字交じり", out: false},
		{in: "123", out: false},
		{in: "#$%", out: false},
		{in: "", out: false},
	}
	for _, cr := range callAndResponse {
		if rsp := kanjiOnly(cr.in); rsp != cr.out {
			t.Errorf("in: %v, got %v, expected: %v", cr.in, rsp, cr.out)
		}
	}
}

func TestLatticeString(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	expected := ""
	str := la.String()
	if str != expected {
		t.Errorf("got %v, expected: %v", str, expected)
	}

	la.Build("わたしまけましたわ")
	str = la.String()
	if str == "" {
		t.Errorf("got empty string")
	}
}

func TestLatticeDot(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("cannot new a lattice")
	}
	defer la.Free()

	expected := "graph lattice {\n\tdpi=48;\n\tgraph [style=filled, rankdir=LR]\n}\n"
	var b bytes.Buffer
	la.Dot(&b)
	if b.String() != expected {
		t.Errorf("got %v, expected: %v", b.String(), expected)
	}
	b.Reset()
	la.Build("わたしまけましたわ")
	la.Dot(&b)
	if b.String() == "" {
		t.Errorf("got empty string")
	}
}

func TestLatticeNewAndFree(t *testing.T) {
	for i := 0; i < 100; i++ {
		la := New(dic.SysDic(), nil)
		if la == nil {
			t.Fatal("unexpected error: cannot new a lattice")
		}
		if la.Input != "" {
			t.Fatalf("unexpected error: lattice input initialize error, %+v", la.Input)
		}
		if len(la.Output) != 0 {
			t.Fatalf("unexpected error: lattice output initialize error, %+v", la.Output)
		}
		if len(la.list) != 0 {
			t.Fatalf("unexpected error: lattice list initialize error, %+v", la.list)
		}
		la.Build("すべては科学する心に宿るのだ")
		la.Free()

		// renew
		la = New(dic.SysDic(), nil)
		if la == nil {
			t.Fatal("unexpected error: cannot new a lattice")
		}
		if la.Input != "" {
			t.Fatalf("unexpected error: lattice input initialize error, %+v", la.Input)
		}
		if len(la.Output) != 0 {
			t.Fatalf("unexpected error: lattice output initialize error, %+v", la.Output)
		}
		if len(la.list) != 0 {
			t.Fatalf("unexpected error: lattice list initialize error, %+v", la.list)
		}
		la.Free()
	}
}

func TestForward(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("unexpected error: cannot new a lattice")
	}

	la.Forward(Normal)
	la.Forward(Search)
	la.Forward(Extended)

	for _, m := range []TokenizeMode{Normal, Search, Extended} {
		la.Build("わたしまけましたわ")
		la.Forward(m)
	}
}

func TestBackward01(t *testing.T) {
	la := New(dic.SysDic(), nil)
	if la == nil {
		t.Fatal("unexpected error: cannot new a lattice")
	}

	// only run
	la.Backward(Normal)
	la.Backward(Search)
	la.Backward(Extended)

	for _, m := range []TokenizeMode{Normal, Search, Extended} {
		la.Build("わたしまけましたわ．ポポピ")
		la.Forward(m)
		la.Backward(m)
	}
}
<|endoftext|>